Research map / Architectures and efficiency / Attention and transformers
dKV-Cache: The Cache for Diffusion Language Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes delayed KV-Cache (dKV-Cache) for masked diffusion language models, caching key/value states of decoded tokens with a one-step delay, in two variants (Decode and Greedy), giving 2-10x inference speedup training-free.
Problem
Diffusion language models are slow at inference because bidirectional attention and non-sequential decoding order make standard KV-Cache inapplicable, requiring full recomputation at every denoising step.
Approach
Cache key/value states only for decoded tokens with a one-step delay (since K/V change most at decoding step), reorder tokens so cached ones are contiguous (concat_reorder), and optionally refresh the cache periodically; a greedy variant also recomputes the previous decoded token and a local window to reach O(L^2) per-step cost.
Key results
- On LLaDA-8B-Instruct, dKV-Cache-Decode with confidence remasking reaches 46.34 on HumanEval at 1.83x speed vs 39.63 for the base model, and 78.85 on GSM8K vs 77.56 base (1.92x).
- On Dream-Base-7B GSM8K (L=256, T=128), dKV-Cache-Prefill reaches 65.96 at 7.11x speedup vs 68.81 for the base model.
- Speedups range from 1.75x to 3.3x in memory/speed analysis, with dKV-Cache-Greedy giving higher speedups than Decode under the same refresh interval.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Apply dKV-Cache-Decode instead of the un-cached base model | LLaDA-8B-Instruct · GSM8K | accuracy: 72.25 -> 78.85 | Helped |
| Apply dKV-Cache-Decode instead of the un-cached base model | LLaDA-8B-Instruct · HumanEval | accuracy: 19.88 -> 46.34 | Helped |
| Apply dKV-Cache-Decode instead of Half-Steps | LLaDA-8B-Instruct · GSM8K | accuracy: 77.91 -> 78.85 | Helped |
| Apply dKV-Cache-Decode instead of the un-cached model | LLaDA-8B-Instruct · MMLU | token/s: 28.27 -> 66.52 | Helped |
| Use dKV-Cache-Greedy with a local window rather than without a window (ablation) | LLaDA-8B-Instruct · GSM8K | accuracy: 67.93 -> 68.23 | Helped |
| Use dKV-Cache-Greedy with a local window rather than without a window (ablation) | LLaDA-8B-Instruct · GPQA | accuracy: 26.79 -> 28.35 | Helped |
| Use dKV-Cache-Greedy with a local window rather than without a window (ablation) | LLaDA-8B-Instruct · HumanEval | accuracy: 15.13 -> 15.37 | Mixed |
| Use a one-step delay for caching (ablation) | dKV-Cache · DLM generation | accuracy | Helped |
| Cache prefill tokens without refreshing (ablation) | Dream-Base-7B · GSM8K (8-shot), L=256, T=256 | token/s: 15.1 -> 53.6 (3.55×) | Helped |
| Use dKV-Cache-PD instead of Half-Steps | Dream-Base-7B · GSM8K (8-shot), L=256, T=128 | accuracy: 46.63 -> 63.31 | Helped |
| Shift cached representations one position right instead of using unshifted representations (ablation) | Dream-Base-7B · MMLU | accuracy: 71.78 -> 64.60 | Hurt |
| Condition caching on both input stability and decoding completion instead of unshifted caching (ablation) | Dream-Base-7B · GSM8K | accuracy: 76.34 -> 77.71 | Helped |
Surprising findings
- Caching can improve accuracy on long sequences (e.g., HumanEval 39.63 to 46.34 on LLaDA), suggesting DLMs may underutilize contextual information.
- Without a one-step delay in caching, accuracy collapses to near zero at high cache ratios.
Limitations
- Algorithm-only work; system-level optimization (gather/scatter overhead, memory management) not addressed; batch size 1 can be slower than baseline.
- dKV-Cache-Greedy relies on random remasking and degrades accuracy relative to confidence remasking.
- Dream requires a position shift handling; the best caching variant (Un&Right-Shift) is incompatible with concat_reorder.
Experimental setup
Datasets: MMLU GPQA GSM8K Math500 HumanEval MBPP
Benchmarks: MMLU GPQA GSM8K Math500 HumanEval MBPP
Metrics: accuracy pass@1 tokens per second speedup cache ratio
Baselines: Few-Steps (random remasking) Half-Steps (confidence remasking) Dream-7B base decoding
Builds on: KV-Cache RoPE masked diffusion language models
Closest papers by meaning
- EntropyCache: Decoded Token Entropy Guided KV Caching for Diffusion Language Models (NeurIPS 2026 · similarity 0.94)
- Polestar: Drift-Aware Cache Calibration and Token Commitment for Efficient Inference of Diffusion LLMs (NeurIPS 2026 · similarity 0.92)
- Fast and Fluent Diffusion Language Models via Convolutional Decoding and Rejective Fine-tuning (NeurIPS 2025 · similarity 0.92)
- Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models (ICLR 2025 · similarity 0.91)
- From Next-Token to Next-Block: A Principled Adaptation Path for Diffusion LLMs (NeurIPS 2026 · similarity 0.91)
- Simple and Effective Masked Diffusion Language Models (NeurIPS 2024 · similarity 0.91)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.