DiffusionGemma: The First Diffusion LLM Natively Supported in vLLM

vLLM has integrated DiffusionGemma, a 26B-parameter discrete diffusion language model built on the Gemma4 backbone. This integration marks the first native support for a diffusion LLM (dLLM) in vLLM, enabling a shift from sequential token generation to iterative block refinement, which significantly reduces latency at low batch sizes.

DiffusionGemma Architecture and Decoding

DiffusionGemma replaces standard autoregressive decoding—where tokens are generated one by one from left to right—with a process that iteratively denoises a fixed-length canvas of 256 tokens. This approach trades memory bandwidth pressure for additional compute, allowing the model to refine multiple tokens in parallel across several denoising steps.

Dual-Mode Operation

DiffusionGemma uses a single set of weights operating in two distinct modes:

• Encoder Mode: Employs causal attention to prefill the prompt and "commit" finished blocks to the KV cache.
• Decoder Mode: Employs bidirectional attention, allowing every position in the canvas to attend to every other position for simultaneous refinement.

Because the encoder mode uses standard causal attention, vLLM's automatic prefix caching remains compatible and works without modification.

The Sampling Loop and Entropy-Bound Denoising

Generation occurs in 256-token blocks. After the prompt is prefilled, the canvas is initialized with random tokens. The model then enters a denoising loop where it samples candidate tokens at every position.

To determine which tokens to keep, DiffusionGemma uses an entropy-bound rule: it accepts tokens from most confident to least confident until the accumulated entropy exceeds a fixed budget. Once a canvas is converged—defined as the argmax prediction remaining stable for several steps and the mean per-token entropy falling below a threshold (or hitting a step limit)—the tokens are committed via an encoder pass and the process repeats for the next block.

Self-Conditioning for Stability

To accelerate convergence and increase stability, the model uses self-conditioning. Between denoising steps, the model is conditioned on the full softmax distribution of its previous prediction. This distribution is converted into a probability-weighted average of token embeddings and added to the canvas embeddings via a gated MLP, providing the model with a memory of its previous beliefs.

vLLM Implementation Details

Integrating DiffusionGemma required supporting a non-autoregressive decoding pattern. vLLM achieved this by leveraging several architectural abstractions:

Speculative Decoding Data Path

vLLM reused its existing speculative decoding path, treating the current canvas as a set of draft tokens that are either fully rejected or fully accepted. This allowed the team to maintain the core scheduler and model runner with minimal changes.

The ModelState Interface

To avoid forking the model runner, vLLM utilized the ModelState abstraction. This provides hooks that allow DiffusionGemma to define custom behavior without altering the generic runner:

Hook	Purpose in DiffusionGemma
prepare_inputs()	Handles canvas embedding and self-conditioning application.
prepare_attn()	Switches between causal (encoder) and bidirectional (denoise) attention per request.
custom_sampler()	Installs the DiffusionSampler to replace the default sampler.
add_request() / remove_request()	Manages the lifecycle of per-request diffusion state (canvas, probabilities).

Dynamic Per-Sequence Causal Attention

Because a single batch can contain requests at different stages (prefill, denoise, commit), vLLM implemented dynamic per-sequence causal attention. This allows the attention mask to adapt to each request's specific causality requirements. This functionality is supported in both the Triton Attention (TRITON_ATTN) and FlashAttention 4 (FLASH_ATTN) backends.

Symmetric Sliding Window Attention

For layers using sliding window attention, DiffusionGemma requires symmetric windows during the denoising phase. While causal requests use a one-sided window (attending to the $W$ tokens before), bidirectional requests attend to $W$ tokens on both sides for a total window size of $2W + 1$.

Performance and Quantization

Throughput Results

DiffusionGemma's architecture enables extremely low-latency inference. Benchmarks conducted at batch size 1 on a single GPU show significant throughput gains over autoregressive baselines:

• H200 (FP8): 1,288 generation tokens per second (~6$ imes$ standard autoregressive baseline, ~3$ imes$ multi-token prediction).
• H100 (FP8): 1,008 generation tokens per second (~5$ imes$ standard autoregressive baseline, ~2.6$ imes$ multi-token prediction).

Quantized Checkpoints

Quantized versions of the model were created using LLM Compressor and are available in the compressed-tensors format. Supported formats include:

• FP8: Quantized weights with fully dynamic activations.
• NVFP4: Both weights and activations quantized to NVFP4 format.