Aufmerksamkeit ist alles, was Sie für den KV-Cache in diffusiven LLMs benötigen.

Attention Is All You Need for KV Cache in Diffusion LLMs

October 16, 2025

papers.authors: Quan Nguyen-Tri, Mukul Ranjan, Zhiqiang Shen

cs.AI

papers.abstract

Diese Arbeit untersucht, wie Schlüssel-Wert (KV)-Caches für Diffusionsgroßsprachmodelle (DLMs) adaptiv neu berechnet werden können, um die Vorhersagegenauigkeit zu maximieren und gleichzeitig die Dekodierungslatenz zu minimieren. Bisherige Methoden berechnen QKV für alle Tokens in jedem Denoising-Schritt und jeder Schicht neu, obwohl sich die KV-Zustände in den meisten Schritten, insbesondere in flachen Schichten, nur wenig ändern, was zu erheblicher Redundanz führt. Wir machen drei Beobachtungen: (1) entfernte {bf MASK}-Tokens wirken hauptsächlich als Längenverzerrung und können blockweise über das aktive Vorhersagefenster hinaus zwischengespeichert werden; (2) die KV-Dynamik nimmt mit der Tiefe zu, was darauf hindeutet, dass ein selektives Aktualisieren ab tieferen Schichten ausreichend ist; und (3) der am meisten beachtete Token zeigt die geringste KV-Drift, was eine konservative untere Grenze für die Cache-Änderung bei anderen Tokens liefert. Basierend darauf schlagen wir {bf Elastic-Cache} vor, eine trainingsfreie, architekturunabhängige Strategie, die gemeinsam entscheidet, {wann} aktualisiert werden soll (über einen aufmerksamkeitsbasierten Drift-Test am am meisten beachteten Token) und {wo} aktualisiert werden soll (über einen tiefenabhängigen Zeitplan, der ab einer gewählten Schicht neu berechnet, während flache Schicht-Caches und off-window MASK-Caches wiederverwendet werden). Im Gegensatz zu festen Periodenschemata führt Elastic-Cache adaptive, schichtbewusste Cache-Aktualisierungen für Diffusions-LLMs durch, reduziert redundante Berechnungen und beschleunigt die Dekodierung mit vernachlässigbarem Verlust an Generierungsqualität. Experimente mit LLaDA-Instruct, LLaDA-1.5 und LLaDA-V in mathematischen Denkaufgaben und Codegenerierungsaufgaben zeigen konsistente Beschleunigungen: 8,7-fach auf GSM8K (256 Tokens), 45,1-fach auf längeren Sequenzen und 4,8-fach auf HumanEval, wobei stets eine höhere Genauigkeit als beim Baseline-Modell erreicht wird. Unsere Methode erzielt einen signifikant höheren Durchsatz (6,8-fach auf GSM8K) als bestehende konfidenzbasierte Ansätze bei gleichbleibender Generierungsqualität, was die praktische Anwendung von Diffusions-LLMs ermöglicht.

English

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.7times on GSM8K (256 tokens), 45.1times on longer sequences, and 4.8times on HumanEval, while consistently maintaining higher accuracy than the baseline. Our method achieves significantly higher throughput (6.8times on GSM8K) than existing confidence-based approaches while preserving generation quality, enabling practical deployment of diffusion LLMs.