PaperlayerView exact revision on arXiv

Source-linked AI summary

Attention Is All You Need for KV Cache in Diffusion LLMs

Quan Nguyen-Tri, Mukul Ranjan, Zhiqiang Shen

arXiv:2510.14973v2cs.CLcs.AIcs.LG

TL;DR

Diffusion LLM decoding repeatedly recomputes mostly stable KV states, creating a latency bottleneck. Elastic-Cache adaptively refreshes caches according to attention drift and layer depth, achieving consistent quality-preserving acceleration across evaluated tasks.

  • Problem

    Diffusion LLM decoders recompute QKV for all tokens and layers despite limited KV changes, motivating adaptive cache updates that preserve accuracy while reducing latency.

  • Method

    Elastic-Cache uses attention-aware drift tests to decide when to refresh, depth-selective recomputation to decide where, and block caching for distant MASK tokens.

  • Results

    Elastic-Cache preserves generation quality while substantially reducing decoding latency across tasks and model scales.

  • Takeaways & Limitations

    The policy provides a practical adaptive approach for reducing redundant QKV work in diffusion LLM decoding.

  • Takeaways & Limitations

    The attention-based measurement can incur amplified error when cached values diverge from actual values.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Diffusion LLM decoding wastes computation by recomputing largely stable KV states across all tokens and layers. Elastic-Cache addresses this with adaptive, layer-aware refreshes that reduce latency while preserving generation quality.

  • 1 INTRODUCTION: KV drift is generally small but increases with layer depth, motivating reuse of shallow-layer caches and selective deeper refreshes.
  • 1 INTRODUCTION: Attention-aware updates monitor the most-attended token as a low-overhead signal for deciding whether a global refresh is warranted.
  • 1 INTRODUCTION: Elastic-Cache jointly decides when to refresh using attention-aware drift and where to refresh using depth-selective updates.It is training-free and architecture-agnostic.
  • 1 INTRODUCTION: Distant MASK tokens can be block-cached outside the active prediction window to eliminate needless updates.
  • 1 INTRODUCTION: Elastic-Cache reduces redundant computation and latency while maintaining generation quality across tasks and model scales.

2 PRELIMINARY

Masked diffusion models iteratively unmask tokens from an initially masked sequence, enabling parallel and flexible text generation. Unlike autoregressive KV caching, bidirectional attention makes diffusion-model KV states evolve across denoising steps.

  • 2 PRELIMINARY: Masked diffusion generation begins from a fully masked sequence and iteratively unmasks tokens using a learned reverse-denoising model.
  • 2 PRELIMINARY: DLMs extend masked diffusion to language modeling while supporting parallel decoding and flexible infilling.
  • 2 PRELIMINARY: Autoregressive KV caching reuses previous key-value pairs because causal-attention states remain invariant during decoding.
  • 2 PRELIMINARY: Bidirectional attention invalidates that invariance in diffusion models, where token representations and cached keys and values evolve across denoising steps.

3 METHODOLOGY

Elastic-Cache combines sliding-window MASK caching with attention-aware drift detection and depth-selective recomputation. It reuses cached states during stable decoding and refreshes deeper layers when attention changes indicate meaningful KV drift.

  • 3.1 OUR FRAMEWORK OVERVIEW AND MOTIVATION: Elastic-Cache refreshes the KV cache from a learned boundary layer onward when most-attended-token attention changes exceed a threshold.If no layer crosses the threshold, the existing cache is reused.
  • 3.2 SLIDING WINDOW DECODING AND KV CACHING: Sliding-window decoding predicts nearby masked tokens together and safely reuses cached KV states for MASK tokens outside the active window.Nearby MASK tokens attend closely, whereas distant MASK tokens have little influence on current predictions.
  • 3.3 ATTENTION-AWARE KV CACHE UPDATE: The method uses attention-weight cosine similarity for most-attended tokens as a lightweight trigger for cache updates.Most-attended tokens are treated as conservative indicators because they typically exhibit the least KV change.
  • 3.3 ATTENTION-AWARE KV CACHE UPDATE: When a layer detects significant attention change, Elastic-Cache recomputes subsequent layers while retaining cached states in preceding layers.
  • 3.3 ATTENTION-AWARE KV CACHE UPDATE: The attention-based trigger is theoretically supported as a computationally efficient proxy for overall KV drift.The paper states that the most-attended token’s KV drift is bounded with a negligibly scaling error term.
  • 3.3 ATTENTION-AWARE KV CACHE UPDATE: The threshold γ controls the accuracy-latency trade-off by determining how frequently and extensively cache updates occur.Higher γ produces more frequent and broader updates; lower γ reduces update frequency.

4 EXPERIMENTS

Elastic-Cache is evaluated across diffusion LLMs, tasks, decoding strategies, and cache configurations, showing substantial throughput gains while largely preserving or improving accuracy. Ablations indicate adaptive refresh and sliding-window choices govern the speed–accuracy trade-off and scaling behavior.

  • 4.2 PERFORMANCE AND EFFICIENCY EVALUATION: 45.1× speedup on GSM8K with LLaDA-1.5 at 512-token generation maintained 81.35% accuracy, matching the baseline.Across benchmarks, Elastic-Cache achieved up to 45.1× speedup with minimal accuracy loss and higher accuracy on GSM8K and MBPP.
  • 4.3 ABLATIONS: Throughput generally improved with longer generation lengths for Elastic-Cache, while Fast-dLLM often slowed as generation length increased.The authors attribute this scaling behavior to the fixed-size sliding window and automatic cache updates.
  • 4.2 PERFORMANCE AND EFFICIENCY EVALUATION: Elastic-Cache increased LLaDA-V MathVerse-256 throughput to 32.3 t/s from Fast-dLLM’s 30.3 t/s while maintaining 29.19% accuracy.This result extends the evaluation beyond text-only tasks.
  • 4.2 PERFORMANCE AND EFFICIENCY EVALUATION: 139.4 t/s at 83.7% accuracy on GSM8K with LLaDA-1.5 outperformed dLLM-Cache and DeepCache fixed-interval variants.The compared DeepCache variants reached 58.4–60.9 t/s at 81.4–83.1% accuracy, while dLLM-Cache reached 16.84 t/s at 80.97%.
  • 4.3 ABLATIONS: Lowering γ increased throughput but tended to reduce accuracy, while larger sliding windows enabled parallel prediction until excessive size raised per-step cost.Accuracy remained stable across β values and close to No-Cache until β ≈64; aggressive denoising also raised cache-update frequency from 5.6% to 17.2%.

5 RELATED WORK

The paper situates diffusion language models within broader efforts to adapt diffusion modeling to discrete text, following multiple probabilistic formulations and applications.

  • Diffusion language modeling has developed through Markov, multinomial, and continuous-time approaches to discrete text generation.

6 CONCLUSION

Elastic-Cache adaptively controls when and where diffusion LLM KV caches are refreshed, reducing redundant computation while preserving generation quality. The authors identify learned drift thresholds, theoretical guarantees, and broader integrations as future directions.

  • Elastic-Cache adaptively refreshes KV caches by testing attention-aware drift and updating layers from a learned boundary onward.
  • Block-caching distant MASK tokens, reusing shallow-layer caches, and refreshing only on meaningful state changes reduce redundant QKV work with negligible quality impact.
  • Future work includes learned drift predictors, formal links between attention and KV drift, hardware-aware scheduling, and extensions to autoregressive and multimodal models.

ETHICS STATEMENT

The work focuses on inference-time efficiency without collecting new data or training models, while acknowledging potential misuse and retaining existing safety constraints.

  • The study introduces no new data collection or model training and evaluates publicly available datasets and third-party checkpoints under their original licenses.
  • Faster decoding may lower generation costs but could also amplify misuse, so the work does not alter safety filters or bypass model alignment constraints.
  • The authors plan to document evaluation materials, follow provider policies, and encourage human oversight for downstream high-stakes deployments.

REPRODUCIBILITY STATEMENT

The appendix specifies implementation symbols, assumptions, and proof components supporting Elastic-Cache’s layer-aware KV refresh strategy. Its theoretical development links hidden-state and attention changes to layer-wise KV drift behavior.

  • Elastic-Cache is controlled by the attention threshold γ, block size, and generation length, with code, configurations, prompts, metrics, environments, and hardware details slated for release.
  • The notation defines transformer layers, denoising steps, sequence length, hidden dimensions, attention weights, decoded positions, masked positions, and sliding windows.
  • The assumptions model progressive unmasking, bounded representations, Lipschitz network components, layer-wise dynamics, and attention concentration.
  • Theorem A.8 states that KV drift becomes layer-dependent with a transition layer, supporting selective recomputation of deeper layers and reuse of shallow caches.
  • The proof proceeds by relating KV drift to hidden-state drift, bounding attention-output changes, recursively propagating layer effects, and applying layer-wise specialization.

A.5 IMPLICATIONS FOR ELASTIC-CACHE

The theoretical results support Elastic-Cache’s depth-aware refresh and most-attended-token drift test.

  • Theorem A.8 establishes that deeper layers have larger KV drift, supporting refreshes beginning at layer ℓ∗.
  • Theorem A.9 shows that most-attended tokens have minimal drift, validating them as conservative cache-staleness indicators.

B DETAILED EXPERIMENT SETUP

The experiments evaluate Elastic-Cache across multiple diffusion LLMs and reasoning tasks using standardized efficiency and quality metrics, alongside sensitivity and overhead analyses.

  • Implementation Details: Experiments cover LLaDA-Instruct, LLaDA-1.5, and LLaDA-V on language and multimodal reasoning tasks, including MBPP, HumanEval, MATH, MathVista, and MathVerse.
  • Evaluation Framework and Metrics: Inference speed is measured as throughput in tokens per second, with task-specific evaluations conducted using lm-eval-harness and multimodal evaluation using lmms-eval.
  • Hyper-parameters: Elastic-Cache uses γ = 0.9 for LLaDA and LLaDA-1.5, γ = 0.7 for LLaDA-V, and generation lengths of 256 and 512 tokens.
  • C.1 COMPREHENSIVE HYPERPARAMETER SENSITIVITY: β = 16 with γ = 0.7 delivers 138.6 t/s while maintaining 77.6% accuracy for throughput-critical deployments.
  • Top-k Sensitivity: Tracking more top-attended tokens slightly improves accuracy but reduces throughput as drift-computation overhead increases.
  • Computational Overhead: The cache-update trigger adds less than 0.5% overhead relative to attention, while drift detection requires only O(KH) additional operations.

C.3 SCALABILITY ANALYSIS

Across scalability, caching, threshold, multimodal, and baseline comparisons, Elastic-Cache improves throughput while preserving accuracy within the reported settings.

  • Multi-GPU Scalability: 225.5 t/s versus 68.0 t/s for Fast-dLLM gives Elastic-Cache a 3.3x throughput advantage on two A100 GPUs with batch size 8.
  • Block-Caching Ablation: Block-caching at β = 16 increases throughput from 82.7 t/s to 119.8 t/s while maintaining similar accuracy, 80.6% versus 81.4%.
  • Adaptive Layer Boundary: Higher γ increases cache-update frequency to 20.02% at γ = 0.95, while lower γ reduces it to 0.47% at γ = 0.5.
  • Token Stability: Most-attended tokens retain cosine similarity of 0.948-0.985 across benchmarks, supporting their use as lower-bound drift indicators.
  • Multimodal Evaluation: Parallel prediction raises LLaDA-V throughput to 44.2 t/s on MathVista and 42.2 t/s on MathVerse.
  • Comparison with Consistency Acceleration: Elastic-Cache maintains 79.2% accuracy and 109.6 t/s, compared with 56.5% accuracy and 35.5 t/s for consistency models.

E SAMPLE RESPONSE

This section presents qualitative generation examples under different γ settings and includes mathematical and coding response examples.

  • Qualitative Examples: The examples compare generated responses under different γ settings for LLaDA-1.5.
  • Mathematical Example: The mathematical examples derive constraints from midpoint symmetry and parabola equations to compute the square of segment AB’s length.
  • Coding Examples: The coding examples include tasks for detecting close elements and returning the decimal part of a number.
Loading 2510.14973v2…