Flash-dLLM: IO-Aware KV Caching and Parallel Decoding for Fast, Memory-Efficient Diffusion LLMs
Quan Nguyen-Tri, Mukul Ranjan, Zhiqiang Shen
cs.CL
2026-09-23
Flash-dLLM fuses KV-cache I/O and lets the dLLM draft-and-verify itself. Up to 148.2x vs uncached greedy; 5.1x / 11.0x vs Elastic-Cache on GSM8K / HumanEval.
Diffusion LLMs unmask tokens in parallel, so they should not wait left-to-right. Open dLLMs still crawl, because bidirectional attention and changing hidden states block a naive port of autoregressive KV caches. Prior accelerators split the job: Fast-dLLM refreshes the full cache at block edges, Elastic-Cache refreshes on attention drift, and draft-and-verify often adds an extra autoregressive model. When cache reuse and parallel checks run together, GPU memory traffic dominates FLOPs.
Flash-dLLM is training-free and has two pieces.
Flash-Cache: a stock cache path launches separate kernels for QKV projection, RoPE, cache write, and attention, bouncing intermediates through HBM. The fused kernel does projection and RoPE in SRAM and writes keys and values straight into the cache. Scheduled Flash Attention uses a block table so samples that only compute a window and samples that recompute the full sequence can share a batch without padding. Selective update follows a sparsity fact: the top-32 attended decoded tokens in middle layers carry about 50% of attention mass. Each step’s query is a sliding mask window βm plus a tracking budget βt; everyone else is served from cache.
Flash-Verify: tokens above ε commit immediately; the rest form a search set. A second forward duplicates those positions, once with the draft token and once still masked, isolated by a causal mask on the shared KV. A token is accepted if both views agree and the mask-view confidence exceeds γ, left to right until the first mismatch. The dLLM is its own drafter and verifier. Extra cost scales with 2βm, not the full sequence.
Defaults: ε=0.9, γ=0.8, βt=80, βm=64, 512-token generation. Main numbers use LLaDA-1.5 on one A100 80GB.
Against uncached greedy, Flash-Cache+Flash-Verify wins throughput in all eight settings at 148.0–210.6 tokens/s and 22.3×–148.2×. GSM8K-512: 83.02% accuracy, 210.6 tokens/s, 81.0×. Against Elastic-Cache, that is 5.1× on GSM8K (210.6 / 41.7) and 11.0× on HumanEval (185.6 / 16.8). In the longer baseline table FreeDave sits at 42.8 tokens/s; the Flash-Verify row is 210.6.
The fused kernel alone is 1.37× on an RTX 3090. Throughput stays near-linear to batch 32; Fast-dLLM OOMs at 24. At batch 16, memory is about 26GB versus 50GB for Fast-dLLM, a 48% cut. On math, the combined setting stays within 1.78 points of the best accuracy. On 256-token code, the gap is about 3.6 points. Throughput-first users pay that.
Parallelism that never shows up in wall-clock time will not displace autoregressive serving. The paper treats I/O and parallel decoding as one system: the fused cache recovers bandwidth, self-verify roughly doubles tokens committed per step, and no extra model is trained. For teams already running masked dLLMs such as LLaDA, this is an inference layer, not a retrain.
The main table is LLaDA-1.5. Continuous-space diffusion LMs are untested. Benchmarks are math and code; open-ended prose has flatter confidence, so fixed thresholds may stall. γ and βm never adapt. On code, the fastest row is not the most accurate. Llama3-8B curves are a reference overlay, not a matched architecture race.