Flash-dLLM: IO-Aware KV Caching and Parallel Decoding for Fast, Memory-Efficient Diffusion LLMs
Abstract
Diffusion Large Language Models (dLLMs) have recently emerged as a promising alternative to autoregressive LLMs by enabling non-autoregressive text generation. However, their practical deployment remains limited by inefficient inference, largely due to the absence of effective Key-Value (KV) caching and scalable parallel decoding mechanisms. Existing acceleration methods typically study KV caching and parallel decoding in isolation, overlooking the I/O bottlenecks that arise when cache reuse and parallel token verification are jointly applied. In this work, we introduce Flash-dLLM, a training-free inference acceleration framework for fast and memory-efficient dLLMs. Flash-dLLM first identifies GPU memory I/O as a dominant bottleneck in KV-cache-enabled dLLM inference and addresses it with an I/O-aware fused KV-cache kernel that reduces redundant memory movement. Building on this optimized cache mechanism, Flash-dLLM further proposes an efficient KV-cache-driven draft-and-verify decoding strategy, where the dLLM itself serves as both drafter and verifier without requiring an auxiliary model. This unified design enables faster decoding while preserving generation quality and improving scalability to longer sequences and larger batch size. Extensive experiments on mathematical reasoning and code-generation benchmarks demonstrate that Flash-dLLM consistently outperforms existing state-of-the-art dLLM acceleration methods in both inference speed and memory efficiency. In particular, it achieves 5.1times and 11.0times speedups over prior strongest baseline Elastic-Cache on GSM8K and HumanEval, respectively.
Community
- Problem: KV caching and parallel decoding in diffusion LLMs are usually studied separately, and repeated KV-cache reads and writes make inference memory-bound, which limits real speedups.
- Flash-Cache: A fused Triton kernel combines QKV projection, RoPE, and cache writes. Block-scheduled attention handles variable query lengths across a batch. Only newly decoded tokens and a fixed set of the most-attended decoded tokens are refreshed.
- Flash-Verify: A training-free draft-and-verify scheme in which the dLLM acts as its own drafter and verifier. It needs no auxiliary model and uses a two-view causal attention mask, which roughly doubles the tokens accepted per step.
- Results (LLaDA-1.5): 148–211 tokens/s, 22.3x –148.2x faster than greedy decoding without caching, and 5.1x/11.0x faster than Elastic-Cache on GSM8K/HumanEval. It also uses about 48% less GPU memory than Fast-dLLM and scales to batch size 32.
This is an automated message from the Librarian Bot. I found the following papers similar to this paper.
The following papers were recommended by the Semantic Scholar API
- Fixed State, Long Reach: What a Constant-Size Cache Buys Block Diffusion at Scale (2026)
- Strong Drafts Need Compact Memories: Long-Context Speculative Decoding with Compressed KV Cache (2026)
- TierKV: Long-Context On-Device LLMs via Predictive Multi-Tier KV Caching (2026)
- S2-MoE: Enabling Efficient Self-Speculative Decoding for Mixture-of-Experts on Edge Devices (2026)
- Faster Than Flash: Exploiting Attention Sparsity for Efficient Long-Context Decoding (2026)
- ASPIRE: Asynchronous Batched Self-Speculative Decoding for Long-Context LLM Inference (2026)
- GroupKV: Hierarchical KV Cache Management for Long-Context Diffusion LLM Inference (2026)
Please give a thumbs up to this comment if you found it helpful!
If you want recommendations for any Paper on Hugging Face checkout this Space
You can directly ask Librarian Bot for paper recommendations by tagging it in a comment: @librarian-bot recommend
Get this paper in your agent:
hf papers read 2609.26796 Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash Models citing this paper 0
No model linking this paper
Datasets citing this paper 0
No dataset linking this paper
Spaces citing this paper 0
No Space linking this paper