Supercharging LLM inference on Google TPUs: Achieving 3X speedups with diffusion-style speculative decoding

This article presents DFlash, a diffusion-style speculative decoding method that achieves 3x speedups for LLM inference on Google TPUs by generating entire token blocks in a single forward pass instead of sequentially.

• •DFlash replaces sequential O(K) autoregressive drafting with O(1) block diffusion, enabling a single forward pass to generate K candidate tokens
• •UCSD researchers integrated DFlash into vLLM TPU framework, achieving 3.13x average speedup (up to 6x for math tasks) on TPU v5p
• •Head-to-head comparison shows DFlash delivered 2.29x speedup versus 1.30x for EAGLE-3, with 10 draft tokens per block compared to 2 for EAGLE-3
• •Implementation required dual-cache architecture for paged attention, power-of-2 padding for context management, and metadata synchronization to prevent sequence length inflation