L'Attenzione è Tutto Ciò di Cui Hai Bisogno per la Cache KV nei Modelli Linguistici di Diffusione

Attention Is All You Need for KV Cache in Diffusion LLMs

October 16, 2025

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

cs.AI

Abstract

Questo lavoro studia come ricomputare in modo adattivo le cache chiave-valore (KV) per i modelli linguistici di diffusione (DLM) al fine di massimizzare l'accuratezza delle previsioni riducendo al minimo la latenza di decodifica. I decodificatori dei metodi precedenti ricomputano QKV per tutti i token a ogni passo di denoising e a ogni livello, nonostante gli stati KV cambino poco nella maggior parte dei passi, specialmente nei livelli superficiali, portando a una sostanziale ridondanza. Facciamo tre osservazioni: (1) i token {bf MASK} distanti agiscono principalmente come un bias di lunghezza e possono essere memorizzati in cache a blocchi oltre la finestra di previsione attiva; (2) le dinamiche KV aumentano con la profondità, suggerendo che un aggiornamento selettivo a partire dai livelli più profondi è sufficiente; e (3) il token con il maggior numero di attenzioni mostra la minima deriva KV, fornendo un limite inferiore conservativo sul cambiamento della cache per gli altri token. Basandoci su queste osservazioni, proponiamo {bf Elastic-Cache}, una strategia senza addestramento e indipendente dall'architettura che decide congiuntamente {quando} aggiornare (tramite un test di deriva basato sull'attenzione sul token con il maggior numero di attenzioni) e {dove} aggiornare (tramite una pianificazione basata sulla profondità che ricomputa a partire da un livello scelto in poi, riutilizzando le cache dei livelli superficiali e le cache MASK fuori dalla finestra). A differenza degli schemi a periodo fisso, Elastic-Cache esegue aggiornamenti adattivi e consapevoli del livello delle cache per i DLM, riducendo il calcolo ridondante e accelerando la decodifica con una perdita trascurabile nella qualità della generazione. Esperimenti su LLaDA-Instruct, LLaDA-1.5 e LLaDA-V su compiti di ragionamento matematico e generazione di codice dimostrano accelerazioni consistenti: 8.7 volte su GSM8K (256 token), 45.1 volte su sequenze più lunghe e 4.8 volte su HumanEval, mantenendo costantemente un'accuratezza superiore rispetto alla baseline. Il nostro metodo raggiunge un throughput significativamente più alto (6.8 volte su GSM8K) rispetto agli approcci esistenti basati sulla confidenza, preservando la qualità della generazione e consentendo la distribuzione pratica dei DLM.

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.