Éviction Stochastique du Cache KV Consciente de la Valeur pour les Modèles de Raisonnement

Résumé

Les modèles de raisonnement améliorent leur précision grâce à des chaînes de pensée étendues, mais leurs longues sorties créent un goulot d'étranglement en termes de mémoire et de calcul. Les méthodes d'éviction du cache KV réduisent ce coût en évacuant du cache les paires clé-valeur peu importantes, mais elles donnent souvent une précision inférieure à celle des alternatives d'attention sparse basées sur la sélection, qui conservent l'intégralité du cache KV. Nous identifions des facteurs clés cruciaux pour la précision de l'éviction du cache KV. Premièrement, une petite fraction des états de valeur présente des magnitudes anormalement élevées, et leur éviction provoque une défaillance catastrophique où les modèles entrent dans des boucles de raisonnement répétitives. Deuxièmement, l'introduction de stochasticité lors de l'éviction améliore la précision en augmentant la diversité du cache. Sur la base de ces constatations, nous proposons l'Éviction Stochastique du Cache KV Sensible à la Valeur (VaSE), une recette sans entraînement qui protège les états de valeur de grande magnitude et favorise des décisions d'éviction diversifiées. Sur six tâches de raisonnement, les modèles Qwen3 utilisant VaSE avec une compression 4x du cache KV atteignent des précisions moyennes plus élevées que la méthode de sélection de pointe à la même rareté, tout en surpassant la méthode d'éviction la plus performante de plus de 4 %. Globalement, VaSE comble le fossé entre efficacité et précision, prend en charge FlashAttention2 et permet une empreinte mémoire statique pour les modèles de raisonnement.

English

Reasoning models improve accuracy through extended chains of thought, but their long outputs create a memory and compute bottleneck. KV cache eviction methods reduce this cost by evicting unimportant key-value pairs from the cache, yet they often yield worse accuracy than selection-based sparse attention alternatives, which keep the full KV cache. We identify key factors crucial to KV cache eviction accuracy. First, a small fraction of value states have abnormally large magnitudes, and evicting them causes catastrophic failure where models enter repetitive reasoning loops. Second, introducing stochasticity during eviction improves accuracy by increasing cache diversity. Based on these findings, we propose Value-aware Stochastic KV Cache Eviction (VaSE), a training-free recipe that protects large-magnitude value states and promotes diverse eviction decisions. Across six reasoning tasks, Qwen3 models using VaSE with 4x KV cache compression yield higher average accuracies than SOTA selection method at the same sparsity, while outperforming the strongest eviction method by more than 4%. Overall, VaSE bridges the gap between efficiency and accuracy, supporting FlashAttention2 and enabling a static memory footprint for reasoning models.