Google TPU Gets 3x Faster Inference with DFlash

The Google Developer Blog introduced DFlash from a UCSD research team. Instead of an autoregressive draft, DFlash uses block diffusion to generate blocks of candidate tokens in a single forward pass, reducing the serial bottleneck of speculative decoding.

After integrating it into vLLM TPU inference, it recorded an average 3.13x improvement in tokens/sec on TPU v5p, with speedups of up to about 6x on complex math tasks. In a head-to-head based on Llama-3.1-8B, DFlash achieved a 2.29x end-to-end speedup, outperforming EAGLE-3's 1.30x.

The implementation was redesigned to fit the TPU/JAX environment.

• Applied a dual-cache structure where the target model uses paged KV cache and the draft model uses static on-device JAX arrays
• Reduced host-TPU transfer overhead for the context buffer via power-of-2 padding
• Synchronized the accepted token count and proposer state to eliminate mismatches between the diffusion draft and the target

Citing the fact that on TPU-class hardware, verification cost for 16 tokens and 1024 tokens is nearly the same, the blog concluded that the bottleneck lies in draft quality rather than verification. It also suggested that improving acceptance probability is a bigger performance lever than increasing block size.