DiffusionGemma Hits 1,500 Tokens/sec by Refining 256 Tokens at Once

The assumption baked into every autoregressive serving stack is that token generation is inherently sequential: you produce token N before you can produce token N+1. That constraint has driven years of inference optimization, from speculative decoding to KV-cache compression. DiffusionGemma removes it entirely, not by building a new architecture from scratch, but by converting an existing pretrained checkpoint.

The mechanism is discrete diffusion applied to blocks rather than individual tokens. Instead of predicting the next token, DiffusionGemma starts with a block of 256 masked tokens and iteratively denoises the entire block in parallel, each forward pass refining the whole block closer to coherent text. Think of it less like writing a sentence word by word and more like developing a photograph: the full image exists at every stage, progressively sharpening with each step. The model starts from Gemma 4, a mixture-of-experts architecture with 3.8B activated and 25.2B total parameters, and converts it through a two-stage training pipeline. Stage one uses supervised fine-tuning to teach bidirectional denoising, replacing the causal attention mask with full bidirectional attention so each token in the block can attend to all others. Stage two combines reinforcement learning with sampler distillation, jointly optimizing generation quality and the number of steps needed per block. The entire pipeline uses fewer than 10% of Gemma 4's original training token budget.

The result: roughly 20 tokens per forward pass on average, and approximately 1,500 output tokens per second on a single NVIDIA H100. That throughput beats autoregressive Gemma 4 even when the AR baseline uses state-of-the-art speculative decoding. The model retains thinking mode, multimodal inputs, and long-context support from the base checkpoint, and remains capable of AR generation with only minor performance degradation. For ML infrastructure teams evaluating next-generation serving architectures, the takeaway is direct: the throughput ceiling on autoregressive decoding is not a hardware problem, and a conversion path from existing checkpoints now exists.

We're thinking: The most consequential detail here is not the speed number. It is the conversion path. We've treated pretrained AR checkpoints as the starting point for capability work, not for architectural change. DiffusionGemma shows that a Gemma 4 checkpoint, trained entirely autoregressively, can be redirected into a fundamentally different decoding regime in under 10% of its original compute budget, and still pass information bidirectionally across a 256-token block. That reframes the question for every team sitting on a large pretrained model: the AR checkpoint is not a constraint you work around with speculative decoding patches. It may be the most efficient on-ramp to a parallel generation regime. The hybrid AR-diffusion capability the paper mentions as a byproduct is worth watching closely, because it suggests these two decoding modes are not mutually exclusive.

Key takeaways:

• DiffusionGemma converts an AR MoE checkpoint to discrete diffusion via two-stage fine-tuning: bidirectional denoising SFT followed by RL plus sampler distillation, enabling full 256-token parallel block refinement without training from scratch.
• The model generates approximately 1,500 output tokens per second on a single H100, surpassing speculative-decoding-augmented AR baselines, using under 10% of Gemma 4's original training token budget; the main caveat is that per-token quality at low sampling steps has not been benchmarked exhaustively across all task types.
• Teams running high-throughput inference on large MoE models should treat this conversion pipeline as a concrete alternative to speculative decoding, particularly for batch workloads where sequential decoding latency compounds.

Source: DiffusionGemma Technical Report