Atenção é Tudo o que Você Precisa para o Cache KV em LLMs de Difusão

Attention Is All You Need for KV Cache in Diffusion LLMs

October 16, 2025

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

cs.AI

Resumo

Este trabalho estuda como recalcular de forma adaptativa os caches de chave-valor (KV) para modelos de linguagem de difusão em larga escala (DLMs) a fim de maximizar a precisão das previsões enquanto minimiza a latência de decodificação. Métodos anteriores recalcularam QKV para todos os tokens em cada etapa de remoção de ruído e camada, apesar dos estados KV mudarem pouco na maioria das etapas, especialmente em camadas mais superficiais, resultando em redundância substancial. Fizemos três observações: (1) tokens {bf MASK} distantes atuam principalmente como um viés de comprimento e podem ser armazenados em cache em blocos além da janela de previsão ativa; (2) a dinâmica KV aumenta com a profundidade, sugerindo que uma atualização seletiva a partir de camadas mais profundas é suficiente; e (3) o token mais atendido exibe o menor desvio KV, fornecendo um limite inferior conservador para a mudança de cache em outros tokens. Com base nisso, propomos {bf Elastic-Cache}, uma estratégia livre de treinamento e independente de arquitetura que decide conjuntamente {quando} atualizar (por meio de um teste de desvio baseado em atenção no token mais atendido) e {onde} atualizar (por meio de um cronograma consciente da profundidade que recalcula a partir de uma camada escolhida enquanto reutiliza caches de camadas superficiais e caches MASK fora da janela). Diferente de esquemas de período fixo, o Elastic-Cache realiza atualizações de cache adaptativas e conscientes da camada para DLMs, reduzindo a computação redundante e acelerando a decodificação com perda insignificante na qualidade da geração. Experimentos em LLaDA-Instruct, LLaDA-1.5 e LLaDA-V em tarefas de raciocínio matemático e geração de código demonstram acelerações consistentes: 8,7 vezes no GSM8K (256 tokens), 45,1 vezes em sequências mais longas e 4,8 vezes no HumanEval, mantendo consistentemente maior precisão que a linha de base. Nosso método alcança uma taxa de transferência significativamente maior (6,8 vezes no GSM8K) que abordagens baseadas em confiança existentes, preservando a qualidade da geração e permitindo a implantação prática de DLMs.

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.