dKV-Cache: The Cache for Diffusion Language Models

We are deeply grateful for the reviewer’s insightful comments and thoughtful suggestions, which have helped us improve the quality of our work. We would revise our manuscript accordingly.

W1: Speedup v.s. prefill-to-generation ratio

W1. The "2-10x speedup" claim while technically supported by the data, could be nuanced. The 10x (or even 18x) speedup appears in specific long-prefill scenarios, which are important but not universal. In more general cases, the speedup is closer to 2-4x. It would strengthen the paper to more clearly delineate the conditions under which different levels of speedup are achieved, perhaps by discussing the speedup as a function of the prefill-to-generation ratio.

Thanks for your question and your great suggestion. The acceleration ratio is relevant to both the generation length and also the prefill-to-generation ratio (ratio in the below table). We here give more results about the acceleration ratio regarding this two factors:

• decode length = 32 | ratio | 0.25 | 0.5 | 1 | 2 | 4 | 8 | 16 |32 | 64 | | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | | dKV-Cache-Prefill | 1.15$\times$ | 1.31$\times$ | 1.68$\times$ | 2.34$\times$ | 3.51$\times$ | 5.48$\times$ | 8.21$\times$ | 11.40$\times$ | 13.99$\times$ | | dKV-Cache-Decode | 1.41$\times$ | 1.48$\times$ | 1.66$\times$ | 1.83$\times$ | 2.04$\times$ | 2.18$\times$ | 2.30$\times$ | 2.37$\times$ | 2.43$\times$ |

• decode length = 256 | ratio | 0.03125 | 0.0625 | 0.125 |0.25 | 0.5 | 1 | 2 | 4 | 8 | | -- | -- | -- | -- | -- | -- | -- | -- | -- | -- | | dKV-Cache-Prefill | 1.00$\times$ | 1.03$\times$ | 1.07$\times$ | 1.18$\times$ | 1.39$\times$ | 1.78$\times$ | 2.52$\times$ | 3.92$\times$ | 6.58$\times$ | | dKV-Cache-Decode | 1.39$\times$ | 1.40$\times$ | 1.44$\times$ | 1.52$\times$ | 1.62$\times$ | 1.78$\times$ | 1.97$\times$ | 2.18$\times$ | 2.39$\times$ |

According to the results, dKV-Cache-Prefill can achieve a significantly higher speedup when the prefill-to-generation ratio is large, due to its ability to reuse cached key/value representations for the long prompt. dKV-Cache-Decode is less influenced by the prefill-to-generation ratio.

In our experiments, we extensively evaluated the method under diverse configurations, including a wide range of zero-shot in-context learning settings, to systematically vary the prefill-to-generation ratio and generation length.

W2: Relationship with semi-autoregressive approaches

W2. The chosen baseline (reducing the number of sampling steps) is reasonable as it's a common way to trade quality for speed. However, the paper could be strengthened by briefly discussing how dKV-Cache compares to or complements other DLM acceleration strategies mentioned in the related work, such as semi-autoregressive approaches. This would better position the work within the broader landscape of DLM inference optimization.

Thanks for your valuable suggestion. We identify two key connections between our method and semi-autoregressive approaches:

1. Complementarity. Given a pre-trained semi-autoregressive model (e.g., Block Diffusion[1]), our approach can be seamlessly integrated to provide intra-block KV-Cache, while the semi-autoregressive design itself naturally induces inter-block KV-Cache.

2. Transformation from a non-autoregressive DLMs toward semi-autoregressive DLMs. This is achieved by applying the proposed dKV-Cache-Greedy strategy and replaces random sampling with a deterministic, left-to-right generation order with a fixed-size window. The left-to-right sampling means that we decoded the n-th token at the n-th step, which has the same decoding order as an auto-regressive model, but with a sliding window sliding from left to right in dKV-Cache-greedy. This makes the model act as a semi-autoregressive one, constraining the decoding and the non-cached tokens inside the window.

We conduct exploratory experiments for this (LLaDA-instruct, GSM8K, official performance = 78.6%):

The experimental results show that with this special sampling strategy under dKV-Cache-Greedy, the model performance remains largely unaffected while enabling the use of a significantly higher cache ratio. This, in turn, leads to a semi-autoregressive model with improved inference speed without compromising accuracy.

We would add this paragraph into the related work section and also add another analysis section for dKV-Cache-Greedy in the revised manuscript.

W3: Explanations for concat_reorder

W3. The concat_reorder operator is critical for implementing a non-sequential cache efficiently, especially with rotational position embeddings (RoPE). While it's deferred to the appendix, a sentence or two in the main text explaining its function and importance would be beneficial for clarity.

Thanks for your suggestion. We would modify and add this short description about the concat_reorder into Line 184-186 to introduce more about this operator:

The concat_reorder operator first reorders the input sequence by placing cached tokens on the left and newly generated (uncached) tokens on the right before passing it into the model. This reordering affects RoPE, so we also reorder the corresponding rotary position matrix accordingly prior to the model forward pass. This design minimizes token-level reordering overhead by applying it only once per decoding step. During attention computation, each attention module then concatenates the cached and new key/value states, applies the appropriate index reordering, and retrieves the updated cached representations from the combined key/value matrices.

Q1: Why dKV-Cache-Decode can improve performance on long sequences

Q1. A very interesting result is that dKV-Cache-Decode can improve performance on long sequences (e.g., HumanEval). The paper suggests this might be due to inefficiencies in how bidirectional attention aggregates context. This finding is quite interesting, what is the authors' thoughts on this? Why does freezing older representations help? Does this imply that the model is "overthinking" or that later-stage token refinements are destructively interfering with earlier, stable parts of the sequence?

Thanks for your insightful question. We also think that this phenomenon is very interesting.

This observation, as you noted, suggests that updating a token and thereby updating the key and value states can in fact degrade the prediction quality of future tokens. We attribute this to error accumulation during intermediate generation stages, where incorrectly generated tokens propagate and negatively impact subsequent decoding. Unlike some autoregressive models with potential for correction (e.g., reflection), current DLMs ( LLaDA and Dream ) lack any refinement mechanism, meaning that early mistakes cannot be revised and persist throughout the sequence. The cache mechanism helps mitigate this issue by using the previous key and value states, effectively bypassing the influence of erroneous tokens and leading to more stable decoding.

[1] Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models

We are thankful for the reviewer’s insightful comments and suggestions, and we will incorporate the following analysis and provide further clarification in the revised manuscript.

W1: Lack of theoretical justification

W1: The method is based purely on empirical observations rather than formal theoretical guarantees. In particular, the reuse of key/value states lacks a rigorous foundation, unlike AR models where caching is mathematically supported by causal masking. The use of the term “KV-Cache” may therefore be misleading.

Thank you for the insightful question. While diffusion language models do not naturally admit autoregressive causal masking, we provide a theoretical justification for the validity of reusing stale key/value states from earlier steps

Our intuition is to treat the cached key/value states as propagation states across decoding steps, similar to the hidden states in RNNs or in linearized attention variants. We aim to derive the bound if we reuse the key/value states from previous steps, e.g., using $k_i^{t-k}$ and $v_i^{t-k}$ at step $t$ instead of freshly computed ones $k_i^{t}$ and $v_i^{t}$. $L$ below means the sequence length.

Consider the standard attention output at decoding step $t$. Following prior work on kernelized attention, we approximate the softmax with a positive kernel $\phi(\cdot)$:

And we have:

Here for dKV-Cache:

• Some $k_i^t$ and $v_i^t$ are fresh, while others are replaced with cached $\tilde{k}_i^t = {k}_i^{t-k}$ and $\tilde{v}_i^t = {v}_i^{t-k}$.
• We denote $\tilde{P}_t$ as $P_t$ cached counterpart; similarly $\tilde{Z}_t$ and $Z_t$.

Let’s define the output error for attention as: $\delta_t=\left\|| o_t- \tilde{o}_t\right\||$, we would get:

First, for

Let’s denote $\mathcal{M}_t \subseteq \\{1,\dots,t\\}$ as the set of indices for which cached states are used, i.e., $\tilde{k}_i^t = k_i^{t-k}$ and $\tilde{v}_i^t = v_i^{t-k}$. For fresh tokens ($i \notin \mathcal{M}_t$), $\tilde{k}_i = k_i^t$ and $\tilde{v}_i = v_i^t$, so the difference is zero. The sum becomes:

with $||\phi(k_i^t) - \phi(\tilde{k}_i)|| \leq \epsilon_k$, $||v_i^t - \tilde{v}_i|| \leq \epsilon_v$, $||\phi(k_i^t)|| \leq C_k$, $||\tilde{v}_i|| \leq C_v$.

For the denominator drift part, we aim to bound:

Assume $\phi(q_t) Z_t \geq \gamma$ and $\phi(q_t) \tilde{Z}_t \geq \gamma'$, similar to the previous one, we can get

Thus, summarizing the above two, we would get $\delta_t$ is bounded by

Interpretation of the bound:

• The more stable the cached $k, v$, the smaller $\epsilon_k, \epsilon_v$ → smaller error.
• The fewer tokens we cache ($|\mathcal{M}_t|$), the smaller the error.
• If dKV-Cache avoids caching tokens (how to select $\mathcal{M}_t$) with high dynamics (i.e., large $|\phi(k_i^t) - \phi(k_i^{t-k})|$), the smaller the error.

W2: Method is somewhat incremental

W2: Similar ideas (e.g., exploiting representation stability across steps) have already appeared in the diffusion literature, particularly in text-to-image generation. The contribution here mainly lies in applying them to DLMs.

Thank you for your question. While our method does use the representation stability across steps, the core contribution lies in the delayed cache mechanism, which is specifically designed for discrete diffusion language models. This delayed caching strategy is unique to the discrete nature of text-based diffusion and is not applicable to standard text-to-image diffusion models. This comes from the special evolving dynamics for the DLMs, and importantly, this design has a substantial impact on performance. As illustrated in Figure 3, when the cache ratio is fixed at 48%, the model’s accuracy improves dramatically, from 3.33% to 74.3%, when using the delayed cache compared to a naïve caching strategy.

Building on this, we propose two variants for dKV-Cache: the greedy one, which uses a tolerant window to allow approximate caching, and the decode one, which aligns caching with decoding dynamics for better performance. Both of these two designs are designed for discrete diffusion language models and cannot be used on image diffusion models.

W3: Missing comparison to AR models

W3: The paper does not show how the proposed caching mechanism compares against autoregressive models in terms of both throughput and generation quality, which is crucial for justifying practical relevance

Thanks for your great question. Despite incorporating dKV-Cache, current Diffusion Language Models still fail to outperform AR models in terms of decoding speed. Below, we present a speed and complexity analysis comparing LLaDA and LLaMA3 under two scenarios: long decoding length (L=1024) and short decoding length (L=32).

• Long Decoding (L=1024, zero-shot, LLaDA-Instruct)

• Short Decoding (L=32, zero-shot, LLaDA-Instruct)

• FLOPs analysis

We provide a FLOPs analysis for attention computation per token during generation:

1. For DLMs, the FLOPs is $4L^2d + 4Ld^2$ per token
2. For ARs, the FLOPs is $4Ld + 4d^2$ per token
3. With dKV-Cache-decode, the expected FLOPs per token is: $2L^2d + 2Ld^2$
4. With dKV-Cache-greedy, the FLOPs is $4wLd + 4wd^2$, where $w$ denotes the attention window size (typically $w \leq 4$).

This indicates that unaccelerated DLMs require approximately $L\times$ more FLOPs than ARs for attention computation, leading to significantly slower inference on long sequences. The proposed dKV-Cache-decode method mitigates this overhead by reducing the FLOPs by half, while dKV-Cache-greedy further improves efficiency by reducing the cost by a factor of L/w.

Summarization:

1. Despite utilizing dKV-Cache to narrow the gap, LLaDA still lags significantly behind ARs in terms of decoding speed, although it achieves competitive performance in terms of generation quality.

2. One scenario where DLMs outperform AR models is under short decoding lengths. Specifically, when the decoding length is limited to L=32, AR models like LLaMA struggle to produce coherent or complete answers, achieving only a score of 1.74, whereas LLaDA maintains strong performance with a score of 38.7. This demonstrates that DLMs are more effective under constrained generation budgets.

Solutions:

To make DLMs as fast or faster than ARs, several key components are needed:

1. dKV-Cache-Greedy to match ARs' complexity with respect to sequence length $L$.
2. Parallel decoding to reduce the number of sampling steps.
3. Semi-autoregressive decoding to bound $L$ and limit computation.
4. System-level optimizations such as vLLM and FlashAttention for improved runtime efficiency.

W4: No architectural improvements for lossless caching

The work focuses solely on inference-time heuristics. A more ambitious direction—modifying DLM architecture to enable theoretically lossless caching—would have been more compelling and impactful

Thanks for your suggestion. We fully agree that this problem maybe can be resolved from the model architecture, such as by redesigning the attention mechanism or the causal mask in a way that ensures strict equivalence for attention calculation and then pretrain this model. It would be highly interesting, impactful and also an ambitious idea and of significant theoretical and practical value. We hope to explore it in the future, making the architecture of DLMs to natively support the KV-Cache mechanisms.

Dear Reviewer 2PJU,

Thank you again for your thoughtful and constructive comments, and also for the time and effort in reviewing our paper.

We sincerely look forward to hearing any further insightful comments you may have. We would greatly appreciate it if you could let us know whether any concerns may need further clarification. We are more than happy to provide additional explanations.

Best Regards,
Authors of Submission 3030

We deeply appreciate the reviewer’s constructive feedback, and we will revise the manuscript accordingly

W1: Lack of analysis of batch size

W1: The paper lacks analysis of some key factors, such as the effect of batch size on acceleration. I’m particularly curious how the speedup of dKV-Cache changes under varying batch sizes (i.e., different GPU utilization levels).

Thanks for your great question. We also think this is a crucial problem and thus we have conducted an analysis on the effects of batch size and decoding length in Appendix D. Please refer to Figure 7 in the Appendix for full details. Below, we summarize the key observations and conclusions:

• At low GPU utilization (e.g., small batch size or short decoding length), the GPU is underutilized, and therefore the speedup is limited.

• When batch size = 1, the improvement in latency is marginal due to insufficient utilization of the GPU.

• When batch size >= 2, the speedup becomes pronounced, reaching approximately 1.4×.

• When batch size >= 64, the acceleration ratio stabilizes at around 2.5×, indicating that our method has better speedup with full GPU utilization.

W2: Comparisons with autoregressive baselines

W2. The paper does not include comparisons with autoregressive baselines. After applying dKV-Cache, is the sampling speed of DLMs actually faster than that of autoregressive models?

Thanks for your great question. Despite incorporating dKV-Cache, current Diffusion Language Models still fail to outperform AR models in terms of decoding speed. Below, we present a speed and complexity analysis comparing LLaDA and LLaMA3 under two scenarios: long decoding length (L=1024) and short decoding length (L=32).

• Long Decoding (L=1024, zero-shot, LLaDA-Instruct)

• Short Decoding (L=32, zero-shot, LLaDA-Instruct)

• FLOPs analysis

We provide a FLOPs analysis for attention computation per token during generation:

1. For DLMs, the FLOPs is $4L^2d + 4Ld^2$ per token
2. For ARs, the FLOPs is $4Ld + 4d^2$ per token
3. With dKV-Cache-decode, the expected FLOPs per token is: $2L^2d + 2Ld^2$
4. With dKV-Cache-greedy, the FLOPs is $4wLd + 4wd^2$, where $w$ denotes the attention window size (typically $w \leq 4$).

This indicates that unaccelerated DLMs require approximately $L\times$ more FLOPs than ARs for attention computation, leading to significantly slower inference on long sequences. The proposed dKV-Cache-decode method mitigates this overhead by reducing the FLOPs by half, while dKV-Cache-greedy further improves efficiency by reducing the cost by a factor of L/w.

Summarization:

1. Despite utilizing dKV-Cache to narrow the gap, LLaDA still lags significantly behind ARs in terms of decoding speed, although it achieves competitive performance in terms of generation quality.

2. One scenario where DLMs outperform AR models is under short decoding lengths. Specifically, when the decoding length is limited to L=32, AR models like LLaMA struggle to produce coherent or complete answers, achieving only a score of 1.74, whereas LLaDA maintains strong performance with a score of 38.7. This demonstrates that DLMs are more effective under constrained generation budgets.

Solutions:

To make DLMs as fast or faster than ARs, several key components are needed:

1. dKV-Cache-Greedy to match ARs' complexity with respect to sequence length $L$.
2. Parallel decoding to reduce the number of sampling steps.
3. Semi-autoregressive decoding to bound $L$ and limit computation.
4. System-level optimizations such as vLLM and FlashAttention for improved runtime efficiency.

Q1: Explanation for Figure 2(b)

Q1：Could you provide more explanation for Figure 2(b)? It appears that the feature distances both before and after decoding are relatively high and fairly similar. However, the paper claims that representations become stable after decoding, while they fluctuate more before decoding. I find this somewhat confusing and would appreciate a more detailed clarification.

Sorry for the unclear explanation here. For Figure 2(b), the calculation pipeline is as follows:

1. We compute the pairwise K/V distance between two adjacent decoding steps, resulting in an array dis[:] of length T-1 for a sequence of T steps.

2. To analyze the impact of decoding for the i-th token, we compute:

   • Before Distance: avg(dis[:i]),
   • After Distance: avg(dis[i:])

From this, we obtain the curve shown in Figure 2(b). The key observations are:

• The before-decoding distance is generally larger than the after-decoding distance, suggesting that the hidden states become more stable as decoding progresses.

• We observe a fluctuation pattern, and in particular, the last step shows after-distance > before-distance. This is because the maximum change in hidden states occurs exactly at the step when a token is decoded. Hence, the ranking of distance magnitudes is: at decoded step > before decode > after decode

For the last token, there’s no step beyond it, so avg(dis[-1:]) = dis[-1], which is exactly at the decode step and has a large value. We think that there is room to improve the presentation of Figure 2(b). Following your suggestion, we plan to enhance it by adding separate curves for before, at decoded, and after distances. Thank you for the valuable feedback.

Q2: LLaDA and Dream with different evaluation settings

Q2. Why do LLaDA and Dream use different evaluation settings? LLaDA is evaluated on an A6000 GPU, while Dream is evaluated on an H20. Also, since Dream uses a few-shot in-context learning setup, did LLaDA use few-shot as well? Moreover, why does this paper use the Instruct model for LLaDA but the base model for Dream?

Thanks for your question. Since there is no big difference between using a base/instruct model regarding the proposed dKV-cache algorithm, we choose the following two settings for Dream and LLaDA in our submission to cover as many cases as possible (whether ICL, Instruct/base, different prefill length):

• LLaDA: LLaDA-Instruct, no in-context learning, A6000.
• Dream: Dream-base, in-context learning, H20.

Here we also show another group of results on Dream-Instruct-7B, with similar experimental results as in Dream-base-7B (no ICL, H20, same hyper-parameters as in Dream-Base-7B).

We provide here another comparison of Dream and LLaDA on the same device (H20) to directly compare their throughput (token/s), where LLaDA-8B is slightly slower than Dream-7B.

Q3: What T refers to in Table 2

Q3. Could you clarify what T refers to in Table 2? It doesn't seem to represent the number of sampling steps. Also, why is T not reported in Table 1?

Thank you for your question and for pointing out the unclear expressions. In Table 2, T refers to the total number of decoding steps (as defined in Line 123). We will clarify this in the caption of the table for better readability.

Regarding Table 1, the experimental setup follows the hyperparameter configuration used in LLaDA, which involves configuring T(sampling steps), L(decoding length), block size, cache interval and window size. For clarity, we listed the detailed hyperparameters, including T, in Table 5 in the appendix.

We sincerely thank the reviewer for the constructive and valuable feedback. We would revise and include the below experiments and analysis in our manuscript.

W1: Why cannot dKV-Cache-Greedy be used on Dream?

W1: The experiments remain somewhat confusing, as they utilize different variants and baselines across various figures for different base models. I would expect the authors to explain why we can only combine dKV-Cache-Greedy with the random version of LLada and why it cannot be used for Dream.

Thanks for your question. dKV-Cache-Greedy can also be applied to Dream, and we provide the corresponding results below. We selected hyperparameters for dKV-Cache-Greedy to achieve approximately 2× acceleration.

Dream's random sampling yields significantly worse performance compared to entropy-based sampling. In some datasets, we observed over a 60% drop in accuracy, with overall accuracy falling below 10%. We hypothesize that this is because DREAM is adapted from an auto-regressive model and thus inherently favors left-to-right generation. As a result, it is less compatible with random sampling strategies.

We think evaluating our method on top of DREAM with random sampling makes it difficult to draw some reliable and generalized conclusions to show the effectiveness of our approach. Thus, we didn't put the paper in the main body of the paper. We would add these results in the Appendix.

W2: dKV-Cache-Greedy still works in limited cases

W2: From the experiment results, it appears that dKV-Cache-Greedy still works in limited cases; therefore, it remains a question of how to reduce complexity while maintaining performance in real-world applications.

Thank you for your insightful question. We agree that dKV-Cache-Greedy, as presented in this paper, has limitations in real-world applications. We first identify that the core reason for these limited cases comes from the unknown decoding order and then propose some potential ways to solve this problem.

• Problem: Unknown decoding order limits the application of dKV-Cache-Greedy

Here, we try another sampling method, AR sampling, to decode LLaDA-instruct. AR sampling refers to enforcing a deterministic left-to-right decoding order, just as the AR models. Here are the experimental results on GSM8K:

Under this sampling method, we observe that model performance remains largely unaffected while enabling the use of a significantly higher cache ratio (less refresh of the cache). This, in turn, leads to improved inference speed without compromising accuracy.

• Solution: How to know the decoding order in advance?

Thus, the primary challenge lies in its reliance on knowing the decoding order in advance, prior to generating each token. To address this limitation and enhance the applicability of dKV-Cache-Greedy, we propose several potential solutions:

1. Decoding Position Estimation Alternatively, we can design a lightweight prediction module to estimate the most likely position to decode in the next step, based on information from the previous iteration. By learning or heuristically approximating the next decoded position, we can use it instead of the random sampling method to achieve better performance and higher cacheability, with the $\mathcal{O}(N^2)$ complexity during inference.

2. Relaxed Decoding Order via Window Tolerance Rather than requiring the exact decoding position, we can instead predict an approximate decoding range, leveraging the inherent window mechanism in dKV-Cache-Greedy to provide tolerance for the imprecise prediction of the decoding position. This relaxes the constraint from precise ordering to coarse localization.