Aandacht Is Alles Wat Je Nodig Hebt voor KV Cache in Diffusie LLM's

Attention Is All You Need for KV Cache in Diffusion LLMs

October 16, 2025

Auteurs: Quan Nguyen-Tri, Mukul Ranjan, Zhiqiang Shen

cs.AI

Samenvatting

Dit werk onderzoekt hoe sleutel-waarde (KV) caches adaptief opnieuw kunnen worden berekend voor diffusie grote taalmodellen (DLMs) om de voorspellingsnauwkeurigheid te maximaliseren terwijl de decodeerlatentie wordt geminimaliseerd. Eerdere methoden herberekenen QKV voor alle tokens bij elke denoisestap en laag, ondanks dat KV-statussen weinig veranderen over de meeste stappen, vooral in ondiepe lagen, wat leidt tot aanzienlijke redundantie. We doen drie observaties: (1) verre {bf MASK} tokens fungeren voornamelijk als een lengtebias en kunnen bloksgewijs worden gecached buiten het actieve voorspellingsvenster; (2) KV-dynamiek neemt toe met de diepte, wat suggereert dat selectief verversen vanaf diepere lagen voldoende is; en (3) het meest-aandachtige token vertoont de kleinste KV-drift, wat een conservatieve ondergrens biedt voor cacheverandering bij andere tokens. Op basis hiervan stellen we {bf Elastic-Cache} voor, een trainingsvrije, architectuuronafhankelijke strategie die gezamenlijk beslist {wanneer} te verversen (via een aandachtbewuste drifttest op het meest-aandachtige token) en {waar} te verversen (via een dieptebewust schema dat opnieuw berekent vanaf een gekozen laag terwijl ondiepe laagcaches en off-window MASK caches worden hergebruikt). In tegenstelling tot vaste-periodeschema's voert Elastic-Cache adaptieve, laagbewuste cache-updates uit voor diffusie LLMs, wat redundante berekeningen vermindert en het decoderen versnelt met verwaarloosbaar verlies in generatiekwaliteit. Experimenten op LLaDA-Instruct, LLaDA-1.5 en LLaDA-V over wiskundige redeneer- en codegeneratietaken tonen consistente snelheidswinsten: 8.7x op GSM8K (256 tokens), 45.1x op langere sequenties en 4.8x op HumanEval, terwijl consistent een hogere nauwkeurigheid wordt behouden dan de baseline. Onze methode bereikt een aanzienlijk hogere doorvoer (6.8x op GSM8K) dan bestaande op vertrouwen gebaseerde benaderingen terwijl de generatiekwaliteit behouden blijft, wat praktische implementatie van diffusie LLMs mogelijk maakt.

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.