CONF-KV : Éviction du cache KV sensible à la confiance avec stockage à précision mixte pour les LLM à horizon long

Résumé

L'inférence des LLM à long horizon transforme le cache clé-valeur (KV) en principal consommateur de mémoire GPU et rend l'attention par jeton de plus en plus coûteuse. De nombreuses politiques d'éviction courantes utilisent des fenêtres de récence statiques ou l'attention historique, laissant inexploité un signal calculé à chaque étape de décodage : l'incertitude actuelle du modèle. Nous présentons CONF-KV, un gestionnaire de cache KV qui convertit la distribution du jeton suivant en un score de confiance scalaire et l'utilise pour choisir le budget de cache par étape, en retenant plus de contexte lorsque le modèle est incertain et en élaguant de manière agressive lorsqu'il est confiant. Dans chaque budget, les jetons sont classés selon une combinaison de la masse d'attention accumulée et de la récence, tandis qu'une fenêtre récente protégée préserve la cohérence locale. Nous associons cette politique à une attention softmax en ligne par blocs, un stockage mixte FP16/INT8 et une variante de budget pyramidal par couche. Sur quatre familles de modèles et des longueurs de génération allant jusqu'à 4K, CONF-KV reste proche de l'empreinte d'une fenêtre glissante fixe de 512 jetons, tout en se maintenant à 1,5–2,1 points de perplexité du cache KV complet. Sur le test Needle-in-a-Haystack jusqu'à 32K jetons, CONF-KV atteint 91,4 % de précision de récupération contre 53,8 % pour les fenêtres glissantes et 80,6 % pour H2O ; sur 75 tâches VisualWebArena, il conserve 95,3 % du succès du KV complet avec une mémoire maximale 2,8 fois inférieure.

English

Long-horizon LLM inference turns the key--value (KV) cache into the dominant GPU memory consumer and makes per-token attention increasingly expensive. Many common eviction policies use static recency windows or historical attention, leaving unused a signal computed on every decoding step: the model's current uncertainty. We introduce CONF-KV, a KV-cache manager that converts the next-token distribution into a scalar confidence score and uses it to choose the per-step cache budget, retaining more context when the model is uncertain and pruning aggressively when it is confident. Within each budget, tokens are ranked by a composite of accumulated attention mass and recency, while a protected recent window preserves local coherence. We combine the policy with blockwise online-softmax attention, mixed FP16/INT8 storage, and a pyramidal per-layer budget variant. Across four model families and generated lengths up to 4K, CONF-KV stays near the footprint of a fixed 512-token sliding window while remaining within 1.5--2.1 perplexity points of full KV. On Needle-in-a-Haystack up to 32K tokens, CONF-KV reaches 91.4% retrieval accuracy versus 53.8% for sliding windows and 80.6% for H2O; on 75 VisualWebArena tasks it retains 95.3% of full-KV success at 2.8 times lower peak memory.