Source-linked AI summary
dKV-Cache: The Cache for Diffusion Language Models
Xinyin Ma, Runpeng Yu, Gongfan Fang, Xinchao Wang
TL;DR
DLM inference is slowed by the incompatibility of conventional KV-Cache with bidirectional, non-sequential denoising. The paper introduces delayed KV-Cache with two variants, achieving broad inference acceleration while largely maintaining performance; its scope is primarily masked diffusion language models and algorithmic optimization.
Problem
DLMs remain slower than AR models because bidirectional attention and non-sequential decoding prevent conventional KV-Cache reuse.
Method
dKV-Cache delays and conditions key-value caching during denoising, with Decode and Greedy variants designed for diffusion-based generation.
Results
2–10× inference speedup is reported across general language understanding, code-generation, and mathematical benchmarks, with minimal or negligible performance loss.
Takeaways & Limitations
The experiments support applying KV-Cache to diffusion language models without training, while dKV-Cache-Decode offers near-lossless performance and dKV-Cache-Greedy offers more aggressive acceleration.
Takeaways & Limitations
The work focuses on algorithmic caching and leaves system-level optimization, including memory management, parallelism, and hardware-aware execution, for future research.
Abstract
from arXiv · showhide
Diffusion Language Models (DLMs) have been seen as a promising competitor for autoregressive language models. However, diffusion language models have long been constrained by slow inference. A core challenge is that their non-autoregressive architecture and bidirectional attention preclude the key-value cache that accelerates decoding. We address this bottleneck by proposing a KV-cache-like mechanism, delayed KV-Cache, for the denoising process of DLMs. Our approach is motivated by the observation that different tokens have distinct representation dynamics throughout the diffusion process. Accordingly, we propose a delayed and conditioned caching strategy for key and value states. We design two complementary variants to cache key and value step-by-step: (1) dKV-Cache-Decode, which provides almost lossless acceleration, and even improves performance on long sequences, suggesting that existing DLMs may under-utilise contextual information during inference. (2) dKV-Cache-Greedy, which has aggressive caching with reduced lifespan, achieving higher speed-ups with quadratic time complexity at the cost of some performance degradation. dKV-Cache, in final, achieves from 2-10x speedup in inference, largely narrowing the gap between ARs and DLMs. We evaluate our dKV-Cache on several benchmarks, delivering acceleration across general language understanding, mathematical, and code-generation benchmarks. Experiments demonstrate that cache can also be used in DLMs, even in a training-free manner from current DLMs.
1 Introduction
DLMs can decode tokens in parallel, but bidirectional attention, flexible generation order, and repeated denoising make them slower than AR models in practice. dKV-Cache adapts key-value caching to this setting and accelerates inference across diverse benchmarks.
- DLMs refine noisy or masked sequences and can decode arbitrary numbers of tokens in parallel, unlike left-to-right AR generation.
- O(L3) denoising arises because length-L generation typically performs L full bidirectional-attention passes, whereas AR KV-Cache reduces per-step complexity to O(L2).
- Bidirectional attention makes prior key and value states change with other tokens, while flexible decoding prevents preselecting the next position for cached computation.
- dKV-Cache delays and conditions key-value reuse by caching states according to tokens’ representation dynamics and decoding status.
- 2–10× inference speedup is achieved on 7B-scale LLaDA and Dream across language understanding, code-generation, and mathematical benchmarks, with minor or negligible degradation.
- The method provides two variants: dKV-Cache-Decode supports long-term reuse, while dKV-Cache-Greedy reduces per-step complexity through more aggressive caching.
2 Related Work
Prior work developed diffusion language models and caching methods for transformers and other generative models, but KV-Cache for fully diffusion-based language models remains comparatively unexplored.
- Diffusion language modeling extends diffusion generation to discrete or continuous language representations, including categorical corruption, embedding-space diffusion, and masked discrete diffusion.
- Transformer KV-Cache stores previous tokens’ key and value tensors, while related diffusion-model caches reuse temporally similar features or attention maps.
- KV-Cache for diffusion language models is less explored; prior semi-autoregressive work requires cache-aware training and remains constrained by an autoregressive formulation.
3 Methods
This section introduces masked diffusion language models, explains why conventional KV-Cache fails under bidirectional and non-sequential denoising, and develops delayed, conditioned caching from token representation dynamics. It then formulates dKV-Cache variants that trade cache lifespan and refresh computation against accuracy and quadratic-time inference.
- 3.1 Preliminary: Masked diffusion language models corrupt tokens through timestep-controlled masking, then repeatedly predict clean tokens and remask uncertain positions during denoising.The forward process uses a transition matrix, while sampling applies the learned reverse model and leaves unmasked tokens unchanged.
- 3.2 Why KV-Cache Cannot be Used in DLMs?: Conventional KV-Cache is incompatible with DLMs because bidirectional attention makes key and value states timestep-dependent, while decoding positions are revealed non-sequentially.These properties prevent globally reusing prior states or pre-determining which token states require computation.
- 3.2 Why KV-Cache Cannot be Used in DLMs?: Key and value states remain highly similar across timesteps, and decoded tokens become relatively stable while masked-token representations continue changing.These dynamics motivate reusing states selectively rather than assuming that all previously computed states are fixed.
- 3.3 Delayed KV-Cache for Masked Diffusion Language Models: dKV-Cache generalizes caching to arbitrary token positions by gathering cached states from an index set and computing attention for dynamically selected decoded or masked-token positions.The formulation replaces contiguous autoregressive slices with reordered key and value states and extends single-token queries to multiple arbitrary positions.
- 3.3 Delayed KV-Cache for Masked Diffusion Language Models: dKV-Cache-Decode delays caching by one denoising step, using the previous masking state to select cacheable tokens and preserving accuracy and stability.The delay avoids reusing key and value states at the step where a token changes from masked to decoded; it slightly reduces efficiency but is critical for quality.
- 3.4 dKV-Cache-Greedy: Greedy Formulation of dKV-Cache: dKV-Cache-Greedy refreshes the current token, the previous decoded token, and a fixed local window, adding computation while retaining overall time complexity O(L^2).Experiments found that centering the local window on the previous decoded token performs better; the window size was at most 6.
4 Experiments
The experiments evaluate dKV-Cache variants across language understanding, mathematical reasoning, and code-generation benchmarks, measuring accuracy, token/s, and cache reuse. Results show broad acceleration with small performance costs, while dKV-Cache-Decode offers near-lossless quality and dKV-Cache-Greedy offers greater speed potential.
- Experimental Setup: Evaluations cover general language understanding, mathematical reasoning, and code-generation benchmarks using accuracy, token/s, and cache-ratio metrics.Speed tests use A6000 for LLaDA and H20 for Dream.
- Performance and Speed with dKV-Cache: dKV-Cache-Decode provides near-lossless performance with a high cache ratio and few refresh steps, achieving the best trade-off among tested strategies.Its accuracy is nearly indistinguishable from the full model.
- Performance and Speed with dKV-Cache: dKV-Cache-Greedy outperforms few-step baselines across most benchmarks, except HumanEval, while a lightweight cache window adds substantial gains with negligible computational overhead.The predefined decoding order used by dKV-Cache-Greedy can reduce accuracy relative to low-confidence remasking.
- Performance and Speed with dKV-Cache: On Dream, dKV-Cache-Prefill achieves up to a 10× acceleration on MMLU and GPQA, and dKV-Cache improves over baselines across prefilling and decoding lengths.Long few-shot contexts make repeated encoding a significant source of overhead.
- Performance and Speed with dKV-Cache: On GSM8K with decoding length 256, dKV-Cache reaches 6.6× speedup and 63.31 Pass@1, versus 4× speedup and 46.63 Pass@1 for the 64-step baseline.The reported performance difference is +16.68.
- Analysis: Across configurations, dKV-Cache is robust to decoding and sampling steps, improves long-form results, and achieves 1.75×–3.3× acceleration with modest memory growth.For L = 512, GSM8K improves from 80.97% to 83.13%, while HumanEval improves from 39.63% to 46.34%.
5 Conclusions and Limitations
The paper establishes delayed KV-Cache as a practical caching mechanism for diffusion language models, with variants that accelerate inference while maintaining performance. Its main limitation is that the work focuses on algorithmic design rather than system-level optimization.
- 5 Conclusions and Limitations: Delayed KV-Cache enables caching in diffusion language models by exploiting delayed reuse of token key and value states.The approach is motivated by token-representation dynamics during diffusion, and delay is crucial for effective caching.
- 5 Conclusions and Limitations: dKV-Cache-Decode and dKV-Cache-Greedy are complementary variants designed for improved compatibility with diffusion-based generation.Both variants leverage delayed caching, but the supplied passages do not quantify their separate trade-offs.
- 5 Conclusions and Limitations: The primary limitation is that the work studies algorithmic design in isolation from system-level optimization.The authors identify memory management, parallelism, and hardware-aware execution as areas for future efficiency improvements.
- 5 Conclusions and Limitations: The implementation reduces indexing overhead by reordering cached and uncached token positions so cached tokens become contiguous during Transformer computation.This shifts some indexing work to token-level operations, while concat and reorder still introduce overhead.
- 5 Conclusions and Limitations: The one-step cache shift aligns cached tokens at step t with positions decoded at step t −1, avoiding storage of entire key/value matrices.The alignment identifies which entries to cache before the current decoding results are available.
B Design for Dream
Dream requires cache designs that account for its autoregressive origin and shifted output positions. The study compares unshifted, right-shifted, and conditioned caching strategies, finding that the conditioned variant performs best but is incompatible with the main concat_reorder implementation.
- B Design for Dream: Dream shifts output positions because it is adapted from pretrained autoregressive models, changing how caching strategies must align token representations.Its output position corresponds to the probability of the next token rather than the current token in traditional masked diffusion.
- B Design for Dream: Un-Shift caches the t-th token’s key and value states at position t.This is one of the three caching strategies compared for Dream.
- B Design for Dream: Right-Shift caches Kt+1 and Vt+1 for the t-th token because hidden states are sensitive to input changes.The one-step shift feeds the token at position t+1 in the next step.
- B Design for Dream: Un&Right-Shift caches features only after the input is fixed and the token has been decoded.This conditions caching on both input stability and decoding completion.
- B Design for Dream: Un&Right-Shift has the best performance, whereas Right-Shift largely harms performance; Un-Shift is used in the main experiment because the best variant is incompatible with concat_reorder.The one-step shift remains part of the conditioned design, but implementation compatibility determines the main experimental choice.
C.1 For LLaDA
The LLaDA evaluation reimplements reported datasets with answer generation and extraction rather than multiple-choice log-probability comparison. Experiment configurations are documented separately.
- C.1 For LLaDA: LLaDA evaluation generates and extracts final answers instead of comparing log probabilities for multiple-choice questions.This evaluation choice affects how MMLU and GPQA results are obtained.
- C.1 For LLaDA: MMLU and GPQA scores are lower than reported because some outputs do not contain answers in the required format.The passage attributes the lower results to generation-format failures or missing answers.
- C.1 For LLaDA: The experiment configurations for LLaDA-Instruct are presented in Table 5.The supplied passage identifies the table but does not enumerate its settings.
C.2 For Dream
The Dream evaluation follows its original pipeline while adding MMLU and GPQA answer-generation evaluations, retaining the original experimental hyperparameters.
- C.2 For Dream: Dream evaluation follows the original evaluation pipeline and adds MMLU and GPQA answer-generation evaluations.The added datasets evaluate generated answers rather than probability comparison.
- C.2 For Dream: The evaluation retains the script’s temperature, remasking strategy, top_p, and few-shot in-context-learning settings.These hyperparameters are followed as configured in the original evaluation script.
D Impact of batch size on speed
Batch size materially affects dKV-Cache decoding speed: batch-one indexing overhead can make caching memory-bound, whereas larger batches provide more stable speedups.
- Batch size is evaluated for decoding speed on LLaDA with a single NVIDIA H20 and a fixed 100-token prefill length.
- At batch size one, uncoalesced gather and scatter reads can leave GPU compute units idle, making inference memory-bound.
- Larger batch sizes can alleviate indexing bottlenecks and deliver greater, more stable speedup ratios over the unaccelerated baseline.
E Case Study
The case-study section presents three tables covering LLaDA and Dream-Instruct examples in mathematical reasoning and story generation.
- The paper presents case studies in Tables 6, 7, and 8.
- Table 6 uses LLaDA-8B-Instruct to examine a multistage running-distance question.
- Tables 7 and 8 use Dream-Instruct-7B for a story-completion request and a running-distance question, respectively.