La Atención Es Todo Lo Que Necesitas para la Caché KV en LLMs de Difusión

Attention Is All You Need for KV Cache in Diffusion LLMs

October 16, 2025

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

cs.AI

Resumen

Este trabajo estudia cómo recalcular de manera adaptativa las cachés de clave-valor (KV) para modelos de lenguaje grandes de difusión (DLMs) con el fin de maximizar la precisión en las predicciones mientras se minimiza la latencia en la decodificación. Los métodos previos recalculan QKV para todos los tokens en cada paso de eliminación de ruido y en cada capa, a pesar de que los estados KV cambian poco en la mayoría de los pasos, especialmente en las capas superficiales, lo que genera una redundancia considerable. Hacemos tres observaciones: (1) los tokens {bf MASK} distantes actúan principalmente como un sesgo de longitud y pueden almacenarse en caché por bloques más allá de la ventana de predicción activa; (2) la dinámica de KV aumenta con la profundidad, lo que sugiere que una actualización selectiva a partir de capas más profundas es suficiente; y (3) el token más atendido exhibe la menor deriva en KV, proporcionando un límite inferior conservador en el cambio de caché para otros tokens. Basándonos en esto, proponemos {bf Elastic-Cache}, una estrategia libre de entrenamiento y agnóstica a la arquitectura que decide conjuntamente {cuándo} actualizar (mediante una prueba de deriva consciente de la atención en el token más atendido) y {dónde} actualizar (mediante un programa consciente de la profundidad que recalcula a partir de una capa elegida en adelante, reutilizando las cachés de capas superficiales y las cachés de MASK fuera de la ventana). A diferencia de los esquemas de período fijo, Elastic-Cache realiza actualizaciones de caché adaptativas y conscientes de la capa para DLMs de difusión, reduciendo el cómputo redundante y acelerando la decodificación con una pérdida insignificante en la calidad de la generación. Los experimentos en LLaDA-Instruct, LLaDA-1.5 y LLaDA-V en tareas de razonamiento matemático y generación de código demuestran aceleraciones consistentes: 8.7 veces en GSM8K (256 tokens), 45.1 veces en secuencias más largas y 4.8 veces en HumanEval, manteniendo consistentemente una mayor precisión que la línea base. Nuestro método logra un rendimiento significativamente mayor (6.8 veces en GSM8K) que los enfoques basados en confianza existentes, preservando la calidad de la generación y permitiendo la implementación práctica de DLMs de difusión.

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.