Flash-dLLM: IO-Aware KV Caching and Parallel Decoding for Fast, Memory-Efficient Diffusion LLMs

Flash-dLLM. IO-Aware KV Caching and Parallel Decoding for Fast, Memory-Efficient Diffusion LLMs. By Quan Nguyen-Tri and two others, from MBZUAI. Published on the twenty-second of September. Its code repository is public and real: twelve source files.

Diffusion Large Language Models, or dLLMs, have emerged as a promising alternative to autoregressive models, because they can generate text non-autoregressively. In practice, their deployment is held back by inefficient inference. Much of that comes from the absence of effective Key-Value, or KV, caching, and of parallel decoding that scales. Existing acceleration methods study caching and parallel decoding separately. That overlooks the input and output bottlenecks that appear when cache reuse and parallel token verification are used together. Two techniques that each speed things up in isolation can run into a new limit when combined, and that limit is what this work targets.

Flash-dLLM is a training-free inference acceleration framework that treats the two together. It first identifies GPU memory input and output as a dominant bottleneck in dLLM inference with a KV cache. It then addresses that bottleneck directly, and builds its decoding strategy on top of the fix.

The first component is an I/O-aware fused KV-cache kernel that reduces redundant memory movement. The second, built on that optimized cache, is a draft-and-verify decoding strategy driven by the KV cache. The dLLM itself serves as both drafter and verifier, so no auxiliary model is required. Together, these let decoding run faster while preserving generation quality, and scale better to longer sequences and larger batch sizes. The order matters. The kernel makes the cache cheap to use, and the drafting strategy then leans on that cache heavily. Because the whole framework is training-free, it applies to an existing dLLM as it stands, with no retraining.

Extensive experiments on mathematical reasoning and code-generation benchmarks show Flash-dLLM consistently outperforming existing state-of-the-art dLLM acceleration methods, in both inference speed and memory efficiency. Against the strongest prior baseline, Elastic-Cache, it achieves speedups of five point one times on GSM8K and eleven point zero times on HumanEval. Those gains are measured against the best existing method, not against unaccelerated inference, and they come with better memory efficiency as well.

The argument is about interaction between optimizations. Caching and parallel decoding each help on their own, but combining them exposes a memory bottleneck neither addresses. Fixing that bottleneck first is what lets the two work together, and using the model as its own drafter keeps the whole scheme to a single model.