Diffusion LLMs can verify their own drafts, removing the need for a separate drafter

Diffusion LLMs generate text by iteratively unmasking tokens rather than producing them left to right. That parallelism is theoretically appealing, but in practice the inference machinery has lagged badly behind autoregressive systems. Flash-dLLM attacks two specific failure modes that prior work treated separately: the memory-bandwidth cost of KV caching across denoising steps, and the throughput ceiling imposed by conservative confidence-aware decoding.

The memory problem is concrete. A conventional dLLM cache pipeline launches four separate CUDA kernels per transformer layer — QKV projection, rotary positional embedding, cache write, attention — each writing intermediate tensors to GPU high-bandwidth memory before the next kernel reads them. The non-attention operations are arithmetic-light and memory-heavy, so the cache-update path becomes memory-bound rather than compute-bound. Flash-dLLM's fused kernel collapses projection, RoPE, and cache write into a single Triton kernel that keeps intermediate values in SRAM, eliminating the redundant HBM round-trips. The paper's Figure 1 caption reports a 1.37× speedup on an RTX 3090 from this kernel alone.

The decoding problem is subtler. Confidence-aware decoding only unmasks tokens above a fixed threshold per step, which is safe but wasteful: many tokens are nearly correct early in denoising but fail the threshold, forcing additional iterations. Flash-Verify addresses this by having the dLLM verify its own draft predictions without any auxiliary model. At each step, search-set tokens — those below the confidence threshold — are duplicated in the query: once filled with the draft prediction, once filled with a mask token. A causal attention mask prevents the two views from attending to each other, so the model produces independent predictions from shared context. A token is accepted only when both views agree and the mask-view confidence clears a second threshold. This is structurally similar to speculative decoding but requires no separate drafter; the dLLM checks its own consistency in one additional forward pass.

Selective cache tracking completes the picture. The paper observes that the top 32 most-attended decoded tokens account for roughly 50% of total attention weight in middle layers. Flash-dLLM therefore maintains a fixed tracking budget rather than recomputing or caching all decoded positions uniformly, bounding per-step compute while preserving most of the useful context.

The combined system scales to batch size 32 on a single A100 without running out of memory, whereas Fast-dLLM hits an out-of-memory error at batch size 24. The memory reduction relative to Fast-dLLM is attributed to a flat preallocated cache layout that avoids dynamic allocation overhead. On GSM8K at 512-token generation length, Flash-Verify with Flash-Cache achieves both the highest accuracy and the highest throughput among evaluated configurations. Accuracy trade-offs are task-dependent: code generation at shorter lengths shows larger gaps than mathematical reasoning.

The limitations are stated plainly. Evaluation covers only masked diffusion models on structured-output tasks; continuous-space diffusion LLMs and open-ended generation remain untested. The confidence and verify thresholds are fixed throughout generation rather than adapted to running statistics.