Attention Is All You Need for KV Cache in Diffusion LLMs (2510.14973v1)

Abstract: This work studies how to adaptively recompute key-value (KV) caches for diffusion LLMs (DLMs) to maximize prediction accuracy while minimizing decoding latency. Prior methods' decoders recompute QKV for all tokens at every denoising step and layer, despite KV states changing little across most steps, especially in shallow layers, leading to substantial redundancy. We make three observations: (1) distant ${\bf MASK}$ tokens primarily act as a length-bias and can be cached block-wise beyond the active prediction window; (2) KV dynamics increase with depth, suggesting that selective refresh starting from deeper layers is sufficient; and (3) the most-attended token exhibits the smallest KV drift, providing a conservative lower bound on cache change for other tokens. Building on these, we propose ${\bf Elastic-Cache}$, a training-free, architecture-agnostic strategy that jointly decides ${when}$ to refresh (via an attention-aware drift test on the most-attended token) and ${where}$ to refresh (via a depth-aware schedule that recomputes from a chosen layer onward while reusing shallow-layer caches and off-window MASK caches). Unlike fixed-period schemes, Elastic-Cache performs adaptive, layer-aware cache updates for diffusion LLMs, reducing redundant computation and accelerating decoding with negligible loss in generation quality. Experiments on LLaDA-Instruct, LLaDA-1.5, and LLaDA-V across mathematical reasoning and code generation tasks demonstrate consistent speedups: $8.7\times$ on GSM8K (256 tokens), $45.1\times$ on longer sequences, and $4.8\times$ on HumanEval, while consistently maintaining higher accuracy than the baseline. Our method achieves significantly higher throughput ($6.8\times$ on GSM8K) than existing confidence-based approaches while preserving generation quality, enabling practical deployment of diffusion LLMs.

• The paper introduces Elastic-Cache, a training-free, architecture-agnostic method that reduces redundant QKV computations in diffusion LLMs.
• It employs a sliding window mechanism and adaptive refresh based on attention drift to achieve up to 45.1× speedup with minimal quality degradation.
• Empirical results on tasks like GSM8K and HumanEval validate its efficiency, showing increased throughput while maintaining high accuracy.

Adaptive Key-Value Caching for Diffusion LLMs

Introduction

Diffusion LLMs (DLMs) have emerged as a competitive alternative to autoregressive transformers, offering parallel decoding and flexible infilling capabilities. However, their iterative denoising process introduces significant computational overhead during inference, primarily due to the need to recompute query, key, and value (QKV) representations for all tokens and layers at every step. This paper introduces Elastic-Cache, a training-free, architecture-agnostic strategy for adaptive key-value (KV) cache management in diffusion LLMs, aiming to minimize redundant computation and accelerate decoding with negligible loss in generation quality.

Motivation and Empirical Observations

The inefficiency in standard DLM inference arises from indiscriminate QKV recomputation, despite the fact that KV states change minimally across most steps, especially in shallow layers. The authors present three key empirical observations:

1. MASK Token Influence: Distant MASK tokens primarily act as a length-bias prior and can be block-cached outside the active prediction window.
2. Layerwise KV Drift: KV drift increases with layer depth, suggesting that selective refreshes starting from deeper layers are sufficient.
3. Attention-Guided Drift: The most-attended token exhibits the smallest KV drift, providing a conservative lower bound for cache change across the context.

These observations are visualized in the following figure, which highlights the attention patterns and KV dynamics that motivate the Elastic-Cache design.

Figure 1: MASK tokens in proximity receive high attention, deeper layers show greater KV drift, and most-attended tokens have minimal KV changes.

Elastic-Cache: Methodology

Sliding Window Decoding and KV Caching

Elastic-Cache employs a sliding window mechanism over masked positions, denoted as $\mathcal{M}^t_\beta$, where $\beta$ is the window size. Only tokens within this window are actively predicted and their KV states updated, while tokens outside the window are block-cached. This approach ensures that only contextually relevant MASK tokens are recomputed, reducing unnecessary updates and maintaining prediction fidelity.

The pipeline is illustrated below, contrasting block-wise caching in Fast-dLLM with the sliding window approach in Elastic-Cache.

Figure 2: Fast-dLLM caches KVs for all tokens outside the current block; Elastic-Cache uses a sliding window and attention-aware refreshes for efficient KV management.

Attention-Aware and Layer-Aware KV Cache Update

Elastic-Cache introduces a two-stage adaptive refresh policy:

• Attention-Aware Update: At each step and layer, the most-attended token is identified based on attention weights. The drift in its attention pattern, measured via cosine similarity between consecutive steps, serves as a trigger for cache refresh. If the similarity falls below a threshold $\gamma$, a refresh is initiated.
• Layer-Aware Update: Upon triggering, only layers deeper than a learned boundary $\ell^*$ are refreshed, while shallow layers and off-window MASK tokens retain their cached KVs.

This mechanism aligns cache updates with actual changes in model beliefs, focusing computation where it is most impactful.

Algorithmic Implementation

The Elastic-Cache algorithm operates as follows:

1. Initialize KV cache for all positions.
2. For each decoding step:
   • Identify the sliding window of masked positions.
   • For each layer, compute attention and update KVs for tokens in the window.
   • Monitor attention drift for the most-attended token; if drift exceeds $\gamma$, refresh KVs from layer $l^*+1$ onward.
   • Decode new tokens and update state.

This process is formalized in the provided pseudocode and is compatible with batch inference via concatenation and custom attention computation.

Experimental Results

Elastic-Cache is evaluated on LLaDA-Instruct, LLaDA-1.5, and LLaDA-V across mathematical reasoning and code generation tasks. Key findings include:

• Throughput: Achieves up to $45.1\times$ speedup on long sequences (GSM8K-512), $8.7\times$ on GSM8K-256, and $4.8\times$ on HumanEval, consistently outperforming Fast-dLLM and baseline DLMs.
• Accuracy: Maintains or improves accuracy, with less than $2\%$ degradation on MATH and HumanEval, and higher accuracy on GSM8K and MBPP.
• Scalability: Throughput increases with generation length, in contrast to Fast-dLLM, which suffers reduced throughput for longer sequences.

Ablation studies demonstrate the tunable trade-off between accuracy and speed via the attention threshold $\gamma$ and window size $\beta$. The sliding window mechanism is shown to outperform block-wise caching, especially for small blocks, by grouping contextually relevant MASK tokens.

Figure 3: Sliding window size impacts accuracy and throughput; larger windows enable parallel prediction but reduce cacheable MASK tokens.

Theoretical Analysis

The paper provides formal justification for the design choices:

• Layerwise Drift Monotonicity: KV drift increases with layer depth, supporting selective refresh of deeper layers.
• Attention Concentration: Most-attended tokens have minimal drift, validating their use as cache staleness indicators.
• Lipschitz Properties: The softmax attention mechanism is 1-Lipschitz, ensuring stability in attention-based drift measurement.

These results establish that Elastic-Cache aligns cache updates with the actual dynamics of the model, minimizing redundant computation without sacrificing output quality.

Practical Implications and Future Directions

Elastic-Cache enables practical deployment of diffusion LLMs by significantly reducing inference latency and resource consumption. Its training-free, architecture-agnostic nature facilitates integration into existing systems. The method is robust across model scales and tasks, including multimodal reasoning.

Future work may include:

• Learning adaptive drift thresholds for cache refresh.
• Formalizing guarantees linking attention patterns to KV drift.
• Extending the approach to autoregressive LLMs and multimodal diffusion frameworks.
• Exploring hardware-aware scheduling and speculative decoding.

Conclusion

Elastic-Cache presents a principled, adaptive strategy for KV cache management in diffusion LLMs, leveraging attention dynamics and layerwise drift to minimize redundant computation. The method achieves substantial speedups with negligible accuracy loss, supporting scalable and efficient inference for large-scale language and multimodal models. Theoretical and empirical analyses validate its effectiveness, and its modular design invites further extension and integration with broader acceleration toolkits.