Flash-dLLM speeds up memory use in language model text generation
Flash-dLLM: IO-Aware KV Caching and Parallel Decoding for Fast, Memory-Efficient Diffusion LLMs
Computation and Language
Summary
Large language models that generate text using a process called diffusion can be slow and memory-heavy when they work. The authors found that moving data between memory and the processor is a big slowdown. They created Flash-dLLM, which reduces unnecessary data movement and uses a clever method where the model writes text drafts and checks them itself without extra tools. This makes text generation faster and less memory hungry while keeping the quality good.
What this means in practice
- •For ai system developers: Build language generation systems that handle longer texts with less memory by using Flash-dLLM’s optimized caching and decoding approach.
- •For cloud service operators: Run diffusion-language-model-powered API services faster and with lower hardware costs by integrating Flash-dLLM’s inference optimizations.
Authors
Quan Nguyen-Tri, Mukul Ranjan, Zhiqiang Shen
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 $\textbf{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.1\times$ and $11.0\times$ speedups over prior strongest baseline Elastic-Cache on GSM8K and HumanEval, respectively.