Desalojo Estocástico de Caché KV Consciente del Valor para Modelos de Razonamiento

Resumen

Los modelos de razonamiento mejoran su precisión mediante cadenas de pensamiento extendidas, pero sus salidas largas crean un cuello de botella en memoria y cómputo. Los métodos de eliminación de la caché KV reducen este costo al descartar pares clave-valor poco importantes de la caché, aunque a menudo producen menor precisión que las alternativas de atención dispersa basadas en selección, que mantienen la caché KV completa. Identificamos factores clave cruciales para la precisión de la eliminación de la caché KV. En primer lugar, una pequeña fracción de los estados de valor tiene magnitudes anormalmente grandes, y eliminarlos provoca un fallo catastrófico en el que los modelos entran en bucles repetitivos de razonamiento. En segundo lugar, introducir estocasticidad durante la eliminación mejora la precisión al aumentar la diversidad de la caché. A partir de estos hallazgos, proponemos Value-aware Stochastic KV Cache Eviction (VaSE), una receta sin entrenamiento que protege los estados de valor de gran magnitud y promueve decisiones de eliminación diversas. En seis tareas de razonamiento, los modelos Qwen3 que utilizan VaSE con compresión 4× de la caché KV logran precisiones promedio más altas que el método de selección de última generación con la misma escasez, superando al método de eliminación más fuerte en más de un 4%. En general, VaSE reduce la brecha entre eficiencia y precisión, es compatible con FlashAttention2 y permite un uso de memoria estática para modelos de razonamiento.

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.