1. Bibliographic Information

• Title: Attention Is All You Need for KV Cache in Diffusion LLMs
• Authors:
  • Quan Nguyen-Tri (FPT AI Residency, Hanoi, Vietnam)
  • Mukul Ranjan (VILA Lab, MBZUAI, Abu Dhabi, UAE)
  • Zhiqiang Shen (VILA Lab, MBZUAI, Abu Dhabi, UAE)
  • The authors are affiliated with research institutions known for contributions to AI and machine learning.
• Journal/Conference: This paper is available as a preprint on arXiv. It has not yet been published in a peer-reviewed journal or conference. Preprints on arXiv are common in the fast-moving field of AI, allowing for rapid dissemination of new research. The arXiv identifier 2510.14973v1 suggests a submission intended for a future date (e.g., October 2025), a convention sometimes used to reserve an identifier.
• Publication Year: The paper was submitted to arXiv in 2024 (inferred from the context and similar paper IDs, despite the futuristic identifier).
• Abstract: The abstract introduces the problem of high decoding latency in diffusion large language models (DLMs) due to redundant computation. Standard DLM decoders recompute the entire Query-Key-Value (QKV) matrix for all tokens at every denoising step and in every layer, even though the internal states (Key-Value or KV cache) change minimally. The authors propose Elastic-Cache, a training-free, adaptive strategy to manage the KV cache. This method decides when to refresh the cache (using an attention-based "drift test") and where to refresh it (using a depth-aware schedule that recomputes only deeper layers). The approach is based on three observations about DLM behavior. Experiments show significant speedups (e.g., up to 45.1x on long sequences) and higher throughput compared to existing methods, all while maintaining or even improving generation accuracy on tasks like mathematical reasoning and code generation.
• Original Source Link:
  • Official Source: https://arxiv.org/abs/2510.14973v1
  • PDF Link: https://arxiv.org/pdf/2510.14973v1.pdf
  • Status: Preprint.

2. Executive Summary

• Background & Motivation (Why):

  • Core Problem: Diffusion Large Language Models (DLMs) are a promising alternative to standard autoregressive models (like GPT) because they can generate text in parallel, potentially offering much faster inference. However, their iterative "denoising" process is computationally expensive. At each step, the model refines its predictions for all masked tokens, requiring a full-pass computation through all layers. This involves re-calculating the Key-Value (KV) cache for every token, creating massive computational redundancy, as most token representations change very little from one step to the next, especially in the early layers of the network.
  • Importance & Gaps: Prior acceleration methods for DLMs, such as Fast-dLLM, use a fixed schedule (e.g., refreshing the cache every $k$ steps or after a block of tokens is decoded). These "one-size-fits-all" approaches are inefficient because they recompute the cache even when it's stable and fail to update it when rapid changes occur. They also treat all layers equally, wasting computation on shallow layers that have already converged while under-serving deeper layers where semantic adjustments are still happening.
  • Innovation: This paper introduces an adaptive, fine-grained caching strategy called Elastic-Cache. Instead of a fixed schedule, it uses the model's own attention patterns as a signal to determine if an update is needed and which parts of the model need updating. This aligns computation with the model's actual information flow, drastically reducing waste.

• Main Contributions / Findings (What):

  1. Diagnosis of Redundancy: The paper empirically demonstrates and quantifies the redundancy in DLM decoding, showing that Key-Value states (KV drift) change minimally in shallow layers and for most tokens across denoising steps.
  2. Elastic-Cache Algorithm: It proposes a novel, training-free, and architecture-agnostic caching policy with two core components:
     • Attention-Aware Refresh Trigger (When): A lightweight test that monitors the change in attention patterns for the single "most-attended" token. If this token's attention distribution shifts significantly, it triggers a cache refresh, acting as a conservative proxy for the entire cache's stability.
     • Depth-Aware Refresh Schedule (Where): When a refresh is triggered at a certain layer, Elastic-Cache only recomputes the KV cache for that layer and all subsequent (deeper) layers, reusing the stable caches of the shallow layers.
  3. Block-wise MASK Caching: It introduces a strategy to cache the representations of distant, un-decoded (MASK) tokens, which are observed to have negligible influence on the current prediction window.
  4. Significant Performance Gains: Elastic-Cache achieves substantial speedups (e.g., 8.7x on GSM8K, 45.1x on longer sequences) and higher throughput (6.8x on GSM8K) compared to baselines, often with a slight improvement in accuracy. This makes DLMs more practical for real-world deployment.

3. Prerequisite Knowledge & Related Work

• Foundational Concepts:

  • Transformer & Self-Attention: The core building block of modern LLMs. In a Transformer, for each token, the model computes three vectors: a Query (Q), a Key (K), and a Value (V). The Q vector of a token is compared with the K vectors of all other tokens to compute "attention scores," which determine how much focus to place on each of those other tokens. These scores are then used to create a weighted sum of the V vectors, producing the token's final representation.
  • KV Cache in Autoregressive LLMs: In standard models like GPT, text is generated one token at a time. To generate the 100th token, the model needs to attend to the previous 99. Instead of recomputing the K and V vectors for all 99 tokens every time, they are calculated once and stored in a KV Cache. For the 101st token, the model only needs to compute the K and V for the 100th token and append them to the cache. This is efficient because past token representations are fixed.
  • Diffusion Models (DMs): A class of generative models that learn to reverse a "noising" process. For images, this involves adding Gaussian noise. For text, a discrete version is used.
  • Masked Diffusion Models (DLMs): A specific type of diffusion model for text. The "noising" process involves progressively replacing original tokens with a special [MASK] token. Generation starts with a sequence of all [MASK] tokens. In each "denoising" step, the model predicts the original tokens for some of the [MASK] positions. A key feature is bidirectional attention: every token can attend to every other token, including future ones (which are [MASK]s). This breaks the assumption of standard KV caching, as the representations of all tokens can change at every step, rendering the simple append-only cache invalid.

• Previous Works:

  • LLaDA and Dream7B: These are foundational DLMs that demonstrated performance competitive with autoregressive models, establishing the viability of the diffusion paradigm for language. The baseline in this paper is a LLaDA model with no caching.
  • Fast-dLLM: A key prior work on accelerating DLMs. It introduced a block-wise caching mechanism. Instead of recomputing the cache at every step, it recomputes it only after a "block" of tokens has been decoded. This is a fixed, non-adaptive schedule. Elastic-Cache is directly compared against Fast-dLLM and shown to be superior.
  • dKV-Cache and DeepCache: Other caching methods. DeepCache is for standard diffusion models (e.g., for images) and reuses features from deeper layers. dKV-Cache also proposes caching for DLMs but relies on a fixed interval. Elastic-Cache differentiates itself by being fully adaptive and fine-grained, using attention drift as a dynamic signal rather than a fixed interval or block size.

• Technological Evolution: The field has moved from computationally intensive, full-recomputation DLMs to more efficient versions. The first step was fixed-schedule caching (Fast-dLLM). This paper represents the next logical step: intelligent, adaptive caching that is aware of the model's internal dynamics.

• Differentiation:

  • Elastic-Cache vs. Fast-dLLM: Fast-dLLM uses a rigid, block-wise schedule. Elastic-Cache uses a dynamic, event-driven schedule based on attention drift.
  • Elastic-Cache vs. Interval-based Caching: Elastic-Cache is more granular. It decides not only when to update but also where (which layers), whereas interval-based methods typically perform a full refresh.
  • Elastic-Cache vs. Autoregressive KV Cache: The mechanisms are fundamentally different. Autoregressive caching is simple and append-only. Elastic-Cache is a complex policy for managing a fully dynamic, bidirectional cache.

4. Methodology (Core Technology & Implementation)

The core of Elastic-Cache is built on three empirical observations about DLM behavior during decoding:

1. Distant MASK tokens act as a length bias: MASK tokens far from the current decoding position receive very little attention and have minimal impact on the prediction. They can be cached without frequent updates.

2. KV dynamics increase with depth: The representations in shallow layers of the network stabilize quickly, while deeper layers continue to refine semantic relationships. This suggests that recomputation is most beneficial in deeper layers.

3. The most-attended token has the smallest KV drift: The token that receives the most attention from the currently predicted MASKs is the most influential. Its representation tends to be the most stable. Therefore, if this token's representation starts to change significantly, it's a strong, conservative signal that other, less stable tokens are also changing, and a cache refresh is needed.

   These observations lead to the Elastic-Cache algorithm, which jointly decides when and where to refresh the cache.

   Figure 1 Analysis: This figure visually confirms the paper's core motivations. (a) Attention is concentrated on prompt tokens and nearby tokens, while "Faraway MASK tokens" receive negligible attention. (b) The heatmap shows cosine similarity of KV states between steps. The similarity is very high (blue, >0.98) for shallow layers but drops for "Deep layers," indicating larger changes. (d) The most-attended tokens show very high similarity across all layers, confirming they are the most stable.

Steps & Procedures

The Elastic-Cache pipeline is illustrated below and contrasted with the Fast-dLLM baseline.

Figure 2 Analysis: This diagram clearly contrasts the two methods. (a) Fast-dLLM uses a fixed block-based approach. It decodes a block of tokens (e.g., 4) and then does a full KV cache recomputation. (b) Elastic-Cache is more dynamic. At each step, it checks for significant attention change using the cosine similarity condition. If the condition is met at a layer $l$, it triggers a partial recomputation from layer $l+1$ to $L$, reusing the cache for shallower layers.

The algorithm proceeds as follows:

1. Initialization: At the first step ($t=0$), a full forward pass is performed, and the initial K and V vectors for all tokens (prompt + MASKs) are computed and stored in the cache.

2. Sliding Window Decoding: For subsequent steps, computation is focused on a "sliding window" of MASK tokens, denoted $\mathcal{M}_\beta^t$, which are the next $\beta$ tokens targeted for decoding. Other, "off-window" MASK tokens are assumed to be stable and their cached representations are reused without recomputation.

3. Attention-Aware Trigger (When to Refresh):

   • At each layer $l$ and step $t$, the model identifies the most-attended token $\mathcal{T}^{t,l}$ from the set of already decoded tokens. This is the token whose Key vector receives the highest cumulative attention score from the Query vectors of the MASKs in the sliding window. $$\mathcal { T } ^ { t , l } = \arg \max _ { k \in \mathcal { D } ^ { < t } } \sum _ { q \in \mathcal { M } _ { \beta } ^ { t } } \mathbf { S } _ { [ q , k ] } ^ { t , l }$$
     • $\mathcal{D}^{<t}$: Set of tokens decoded before step $t$.
     • $\mathcal{M}_\beta^t$: Set of MASK tokens in the current prediction window.
     • $\mathbf{S}^{t,l}$: Attention score matrix at step $t$, layer $l$.
   • The algorithm then checks if the attention distribution directed at this most-attended token has changed significantly from the previous step. It computes the cosine similarity between the attention score vectors for this token at step $t$ and t-1.
   • If the similarity drops below a predefined threshold $\gamma$, a refresh is triggered for that layer. $$\text{trigger if: } \mathrm{CosineSimilarity}(\mathbf{S}_{[\mathcal{T}^{t-1}]}^{t-1, l}, \mathbf{S}_{[\mathcal{T}^{t-1}]}^{t, l}) < \gamma$$

4. Layer-Aware Schedule (Where to Refresh):

   • Let $l^*$ be the first (shallowest) layer where the trigger condition is met.

   • The KV caches for all shallow layers $l < l^*+1$ are reused.

   • The KV caches for all deeper layers $l \ge l^*+1$ are recomputed. This is done by taking the hidden states from the output of layer $l^*$ and running a full forward pass from layer $l^*+1$ to the final layer $L$.

   • This "recompute from here onward" strategy saves significant computation by not re-running the stable shallow layers.

     The complete process is detailed in Algorithm 1 of the paper.

5. Experimental Setup

• Datasets:

  • GSM8K: A dataset of grade school math word problems.
  • MATH: A more challenging dataset of competition-level mathematics problems.
  • HumanEval & MBPP: Datasets for evaluating code generation from natural language descriptions.
  • MathVista & MathVerse: Multimodal datasets that require mathematical reasoning over visual contexts (e.g., charts, diagrams). These datasets were chosen to test the method's effectiveness on complex reasoning tasks where generation quality is critical.

• Evaluation Metrics:

  • Accuracy:
    • Conceptual Definition: Measures the correctness of the generated output. The specific calculation depends on the task.
    • For GSM8K, flexible_extract is used, which checks if the final numerical answer in the generation matches the ground truth.
    • For MATH, math_verify checks the correctness of mathematical derivations.
    • For HumanEval and MBPP, pass@1 is used. It measures the percentage of problems for which the generated code passes a set of unit tests.
    • For MathVista/MathVerse, gpt-eval-score uses a powerful model like GPT-4 to judge the correctness of the response.
  • Throughput:
    • Conceptual Definition: Measures the speed of generation. It is the total number of tokens generated divided by the total time taken.
    • Mathematical Formula: $$\mathrm{Throughput} = \frac{\text{Number of Generated Tokens}}{\text{Wall-clock Time (seconds)}}$$
    • Unit: tokens/second. Higher is better.
  • Speedup:
    • Conceptual Definition: A relative measure of how much faster a method is compared to a baseline.
    • Mathematical Formula: $$\mathrm{Speedup} = \frac{\mathrm{Throughput}_{\text{Method}}}{\mathrm{Throughput}_{\text{Baseline}}}$$
    • Unit: A dimensionless factor (e.g., 2.0x means twice as fast).

• Baselines:

  • LLaDA (No Cache): The standard DLM implementation where the QKV is recomputed for all tokens at every step and layer. This is the slowest but most accurate baseline.
  • Fast-dLLM: The primary competitor, which uses a block-wise caching strategy. It represents the state-of-the-art in fixed-schedule caching for DLMs.

6. Results & Analysis

The experiments convincingly demonstrate that Elastic-Cache provides a superior trade-off between speed and accuracy.

Core Results

The following transcribed tables summarize the main results on the LLaDA-Instruct and LLaDA-1.5 models.

Table 1: Performance on LLaDA-Instruct (Manual Transcription) Each cell shows Accuracy% (top) and Throughput (tokens/s) with Speedup (bottom).

Table 2: Performance on LLaDA-1.5 (Manual Transcription) Each cell shows Accuracy% (top) and Throughput (tokens/s) with Speedup (bottom).

Analysis:

• Superior Speed-Accuracy Trade-off: In almost all cases, Elastic-Cache achieves higher throughput than Fast-dLLM while also achieving better accuracy. For example, on GSM8K with LLaDA-1.5 (length 512), Elastic-Cache achieves a staggering 45.1x speedup while maintaining the baseline's 81.35% accuracy. In contrast, Fast-dLLM's speedup is only 14.1x, and its accuracy drops to 80.82%.
• Advantage Increases with Length: The speedup advantage of Elastic-Cache over Fast-dLLM becomes more pronounced for longer generation lengths. In Table 1 (GSM8K), the speedup factor goes from 8.2x (at 256 tokens) to 25.2x (at 512 tokens). This is because the sliding window and selective refresh mechanisms become more effective as the proportion of cacheable "distant" tokens grows.
• Robustness Across Models and Tasks: The gains are consistent across different models (LLaDA-Instruct, LLaDA-1.5, LLaDA-V) and tasks (math, code, multimodal reasoning), demonstrating the general applicability of the method.

Ablations / Parameter Sensitivity

![该图像是包含三部分的图表，分别展示了滑动窗口与块大小(a)、缓存更新时间频率(b)、以及置信感知解码阈值(c)对GSM8K 5-shot准确率和吞吐量的影响。图中用箭头标注了不同配置下的加速比，说明了Elastic-Cache在保持准确率的同时显著提升吞吐量。](/files/papers/68f43749d6dbb63266273665/images/3.jpg)
*该图像是包含三部分的图表，分别展示了滑动窗口与块大小(a)、缓存更新时间频率(b)、以及置信感知解码阈值(c)对GSM8K 5-shot准确率和吞吐量的影响。图中用箭头标注了不同配置下的加速比，说明了Elastic-Cache在保持准确率的同时显著提升吞吐量。*
*Figure 3 Analysis: This figure explores the method's sensitivity to its design choices.

(a) Sliding Window vs. Block-wise: The proposed sliding window approach (blue lines) consistently achieves higher accuracy than a block-wise approach (orange lines). While block-wise can have high throughput at a specific size, the sliding window provides a better and more stable accuracy-throughput curve. (b) Cache Update Frequency vs. Gamma ($\gamma$): As the gamma threshold increases (making the trigger more sensitive), throughput decreases because the cache is updated more often. However, accuracy increases. Elastic-Cache (blue line) maintains high accuracy even with low update frequencies, significantly outperforming the no-cache baseline's speed (orange line). (c) Confidence-Aware Decoding: Elastic-Cache (blue line) maintains a high throughput advantage over the baseline (orange line) across different confidence thresholds for parallel decoding.*

Table 4: Impact of $γ$ on GSM8K (512 tokens) (Manual Transcription) Each cell shows Accuracy% (top) and Throughput (tokens/s) with Speedup (bottom).

Analysis:

• Tunable Trade-off: Table 4 clearly shows that the hyperparameter $\gamma$ provides a direct lever to trade speed for accuracy. For LLaDA-1.5, setting $\gamma=0.5$ yields a massive 54.9x speedup but drops accuracy to 76.04%. Setting $\gamma=0.95$ yields a lower but still huge 37.8x speedup while improving accuracy to 83.02% over the baseline. This tunability is a powerful feature for practical deployment.
• Optimal $γ$: The best $γ$ depends on the model. For the more accurate LLaDA-1.5, a higher $γ$ is beneficial, suggesting that as models get better, more sensitive cache updates are needed to preserve their performance.

7. Conclusion & Reflections

• Conclusion Summary: The paper successfully identifies and addresses a major bottleneck in diffusion LLM inference: redundant KV cache recomputation. The proposed Elastic-Cache method offers a principled, training-free, and highly effective solution. By using attention dynamics to create an adaptive policy that decides when and where to refresh the cache, it achieves significant latency reductions and throughput improvements with negligible, and sometimes positive, impact on generation quality. This work makes an important contribution toward making large-scale deployment of diffusion LLMs practical.

• Limitations & Future Work: The authors propose several avenues for future work:

  • Integrating Elastic-Cache with other acceleration techniques like speculative decoding.
  • Developing hardware-aware scheduling to further optimize the cache management policy.
  • Extending the core principles to autoregressive LLMs and other multimodal diffusion frameworks.

• Personal Insights & Critique:

  • Novelty and Elegance: The central idea of using the most-attended token's stability as a proxy for the entire cache's drift is both clever and elegant. It's a simple heuristic grounded in a strong empirical observation, avoiding the complexity of a learned policy while achieving excellent results. The depth-aware refresh is another simple yet powerful optimization.
  • Practical Impact: Being training-free and architecture-agnostic makes Elastic-Cache highly practical. It can be implemented as a wrapper around existing DLMs without requiring model retraining, lowering the barrier to adoption. The dramatic speedups shown could be a game-changer for the real-world viability of DLMs.
  • Open Questions:
    • While the "most-attended token" is a good heuristic, is it always the optimal one? An alternative could be to monitor a small set of highly-attended tokens to get a more robust signal.
    • The method still relies on a manually tuned hyperparameter, $\gamma$. While it offers a useful trade-off lever, an interesting research direction would be to learn this threshold dynamically based on the input prompt or decoding state.
    • The analysis focuses on latency reduction. A deeper analysis of the memory footprint implications would be valuable, especially as sequence lengths grow.