Papers
arxiv:2609.26796

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

Published on Sep 22
· Submitted by
Mukul Ranjan
on Sep 23
Authors:
,
,

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

Paper submitter
  • 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.

Sign up or log in to comment

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

Cite arxiv.org/abs/2609.26796 in a model README.md to link it from this page.

Datasets citing this paper 0

No dataset linking this paper

Cite arxiv.org/abs/2609.26796 in a dataset README.md to link it from this page.

Spaces citing this paper 0

No Space linking this paper

Cite arxiv.org/abs/2609.26796 in a Space README.md to link it from this page.

Collections including this paper 0

No Collection including this paper

Add this paper to a collection to link it from this page.