Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Attention Is All You Need for KV Cache in Diffusion LLMs

Quan Nguyen-Tri, Mukul Ranjan, Zhiqiang Shen

TL;DR

Diffusion LLMs incur high decoding latency due to recomputing KV caches for all tokens at every step. The authors introduce Elastic-Cache, a training-free, adaptive KV caching policy that uses attention-driven drift and a learned layer boundary $\\ell^\\star$ to decide when and where to refresh, complemented by sliding-window decoding and block-caching of distant MASK tokens. They establish three KV-drift insights—drift increases with depth, distant MASKs act as length priors, and the most-attended token drifts the least—and prove a bound on drift for the most-attended token, guiding the update strategy $ (\\Delta^{t,\\ell}_{\\mathcal{T}^{t,\\ell}} \le \\bar{\\Delta}^{t,\\ell} + O\\left(\\frac{\\sqrt{d_k}}{R_\\ell\\sqrt{N}}\\right)) $. Empirically, Elastic-Cache yields up to $45.1\\times$ speedups on long sequences and $8.7\\times$ on GSM8K (256 tokens) with negligible accuracy loss, and scales to multimodal settings, enabling practical deployment of diffusion LLMs. The approach is architecture-agnostic, training-free, and complementary to existing acceleration techniques, with substantial throughput gains and reduced memory traffic. These results demonstrate that aligning KV recomputation with where and when the model's beliefs change can dramatically accelerate diffusion-based generation without sacrificing quality.

Abstract

This work studies how to adaptively recompute key-value (KV) caches for diffusion large language models (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), and $45.1\times$ on longer sequences, 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.

Attention Is All You Need for KV Cache in Diffusion LLMs

TL;DR

Diffusion LLMs incur high decoding latency due to recomputing KV caches for all tokens at every step. The authors introduce Elastic-Cache, a training-free, adaptive KV caching policy that uses attention-driven drift and a learned layer boundary to decide when and where to refresh, complemented by sliding-window decoding and block-caching of distant MASK tokens. They establish three KV-drift insights—drift increases with depth, distant MASKs act as length priors, and the most-attended token drifts the least—and prove a bound on drift for the most-attended token, guiding the update strategy . Empirically, Elastic-Cache yields up to speedups on long sequences and on GSM8K (256 tokens) with negligible accuracy loss, and scales to multimodal settings, enabling practical deployment of diffusion LLMs. The approach is architecture-agnostic, training-free, and complementary to existing acceleration techniques, with substantial throughput gains and reduced memory traffic. These results demonstrate that aligning KV recomputation with where and when the model's beliefs change can dramatically accelerate diffusion-based generation without sacrificing quality.

Abstract

This work studies how to adaptively recompute key-value (KV) caches for diffusion large language models (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 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 , a training-free, architecture-agnostic strategy that jointly decides to refresh (via an attention-aware drift test on the most-attended token) and 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: on GSM8K (256 tokens), and on longer sequences, while consistently maintaining higher accuracy than the baseline. Our method achieves significantly higher throughput ( on GSM8K) than existing confidence-based approaches while preserving generation quality, enabling practical deployment of diffusion LLMs.
Paper Structure (33 sections, 3 theorems, 69 equations, 3 figures, 23 tables)

This paper contains 33 sections, 3 theorems, 69 equations, 3 figures, 23 tables.

Table of Contents

  1. Introduction
  2. Preliminary
  3. Masked Diffusion Models
  4. Key-Value Cache in Transformers
  5. Methodology
  6. Our Framework Overview and Motivation
  7. Sliding window decoding and KV caching
  8. Attention-aware KV Cache Update
  9. Experiments
  10. Experimental Setup
  11. Performance and Efficiency Evaluation
  12. Ablations
  13. Analysis
  14. Related Work
  15. Conclusion
  16. ...and 18 more sections

Key Result

Lemma A.1

Based on the Proposition 2 in gao2017properties, the softmax function $\sigma: \mathbb{R}^n \to \Delta^{n-1}$ defined by is $1$-Lipschitz continuous with respect to the $\ell_2$ norm:

Figures (3)

  • Figure 1: Visualization of our motivation. (a) MASK tokens located near each other receive high attention, while those situated far apart have minimal influence. (b) Over time, the representations in the KV states of cached tokens evolve, with deeper layers experiencing more substantial changes. (c) The changes in attention weights of most-attended tokens exhibit similar patterns to the changes in KV states of all cached tokens. (d) KV states of the most-attended tokens have the least changes.
  • Figure 2: Illustration of the Key-Value cache method for diffusion LLMs. (a) The fast-dLLM wu2025fast block-wise decoding method caches the Key-Value of all tokens outside the current block at each step. The KV cache is updated after completing a block of decoding. (b) Our proposed method, Elastic-Cache, caches the key-value of tokens outside a sliding window that flexibly moves through the sentence from left to right at each iteration. When the attention weights corresponding to the most-attended tokens (one for each layer) change significantly at a layer $l$, we start recomputing the KV cache from layer $l+1$ to the last layer.
  • Figure 3: Ablation study and analysis of our proposed method. (a) Ablation study of our sliding window mechanism compared to block-wise decoding. (b) Analysis of cache update frequency under varying $\gamma$. The blue and orange lines represent accuracy and throughput, respectively. The numbers along the lines indicate the frequency of cache updates, assuming no baseline. (c) Analysis of cache update frequency under confident-aware decoding with varying $\epsilon$.

Theorems & Definitions (6)

  • Lemma A.1: Lipschitz Continuity of Softmax
  • Definition A.7: KV Drift
  • Theorem A.8: Layer-Wise KV Drift Monotonicity
  • proof
  • Theorem A.9: Attention Concentration and Drift
  • proof