L'attention est tout ce dont vous avez besoin pour le cache KV dans les LLMs de diffusion.

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

Ce travail étudie comment recalculer de manière adaptative les caches clé-valeur (KV) pour les grands modèles de langage à diffusion (DLM) afin de maximiser la précision des prédictions tout en minimisant la latence de décodage. Les méthodes précédentes recalculent QKV pour tous les tokens à chaque étape de débruitage et à chaque couche, bien que les états KV changent peu à travers la plupart des étapes, en particulier dans les couches superficielles, ce qui entraîne une redondance substantielle. Nous faisons trois observations : (1) les tokens {bf MASK} distants agissent principalement comme un biais de longueur et peuvent être mis en cache par blocs au-delà de la fenêtre de prédiction active ; (2) la dynamique des KV augmente avec la profondeur, suggérant qu'un rafraîchissement sélectif à partir des couches plus profondes est suffisant ; et (3) le token le plus fréquemment sollicité présente la plus faible dérive des KV, fournissant une limite inférieure conservatrice sur le changement de cache pour les autres tokens. Sur cette base, nous proposons {bf Elastic-Cache}, une stratégie sans apprentissage et indépendante de l'architecture qui décide conjointement {quand} rafraîchir (via un test de dérive basé sur l'attention pour le token le plus sollicité) et {où} rafraîchir (via un calendrier basé sur la profondeur qui recalcule à partir d'une couche choisie tout en réutilisant les caches des couches superficielles et les caches MASK hors fenêtre). Contrairement aux schémas à période fixe, Elastic-Cache effectue des mises à jour de cache adaptatives et conscientes de la couche pour les DLM, réduisant les calculs redondants et accélérant le décodage avec une perte négligeable de qualité de génération. Les expériences sur LLaDA-Instruct, LLaDA-1.5 et LLaDA-V à travers des tâches de raisonnement mathématique et de génération de code démontrent des accélérations constantes : 8,7 fois sur GSM8K (256 tokens), 45,1 fois sur des séquences plus longues, et 4,8 fois sur HumanEval, tout en maintenant systématiquement une précision plus élevée que la référence. Notre méthode atteint un débit significativement plus élevé (6,8 fois sur GSM8K) que les approches existantes basées sur la confiance tout en préservant la qualité de génération, permettant un déploiement pratique des 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.