CONF-KV: Confidence-Aware KV Cache Eviction with Mixed-Precision Storage for Long-Horizon LLM

Resumen Técnico: CONF-KV

Enunciado del Problema
La inferencia de Modelos de Lenguaje Grandes (LLM) de horizonte largo enfrenta un cuello de botella crítico donde la memoria caché de Clave-Valor (KV) se convierte en el consumidor dominante de memoria de la GPU, provocando que la latencia de atención por token escale linealmente con la longitud de la secuencia. Las políticas de expulsión existentes, como las ventanas deslizantes o los métodos basados en atención histórica (por ejemplo, H2O, SCISSORHANDS), dependen de la recencia estática o de señales pasadas para determinar qué tokens retener. Estos enfoques a menudo fallan al no utilizar una señal disponible en cada paso de decodificación: la incertidumbre actual del modelo. En consecuencia, las políticas estáticas pueden expulsar prematuramente el contexto durante span de generación difíciles o retener tokens innecesarios cuando el modelo está seguro, lo que conduce a un compromiso subóptimo entre el uso de memoria y la calidad de la generación.

Metodología: CONF-KV
Los autores proponen CONF-KV, un gestor de caché KV consciente de la confianza que ajusta dinámicamente el presupuesto de la caché en función de la distribución del siguiente token actual del modelo. El sistema opera mediante los siguientes mecanismos:

1. Estimación de Confianza: En cada paso de decodificación, el método convierte los logits en una puntuación de confianza escalar ($c$). Esta puntuación es una combinación ponderada de la entropía normalizada, el margen de log-probabilidad entre los dos tokens superiores y la probabilidad del token superior.
   $$c = w_H(1 - \hat{H}) + w_m\sigma(m) + w_p p_{(1)}$$
   donde $\hat{H}$ es la entropía normalizada, $m$ es el margen de logaritmos y los pesos se establecen en $(0.4, 0.3, 0.3)$.

2. Selección de Presupuesto Adaptativo: Un umbral $\tau$ determina el presupuesto de la caché ($N$) para el paso actual.

   • Si $c \ge \tau$ (alta confianza), se selecciona un presupuesto ajustado ($N_{high}$), permitiendo una expulsión agresiva.
   • Si $c < \tau$ (baja confianza/inseguridad), se selecciona un presupuesto más holgado ($N_{low}$) para retener más contexto.
     Esto crea un perfil de memoria en "diente de sierra" donde la caché se expande durante span de incertidumbre y se contrae durante los de confianza.

3. Clasificación y Expulsión de Tokens: Dentro del presupuesto seleccionado, los tokens se clasifican según una puntuación compuesta que combina la masa de atención de Media Móvil Exponencial (EMA) y la recencia. Una ventana protegida rígida ($P$) asegura que los tokens más recientes nunca sean expulsados para preservar la coherencia local.

4. Integración del Sistema:

   • Almacenamiento de Precisión Mixta: Los $W$ tokens más recientes se almacenan en FP16, mientras que los tokens retenidos más antiguos se cuantizan a INT8 por cabeza y canal. Esto reduce la huella de memoria mientras minimiza la degradación de la perplejidad en comparación con la cuantización de bits inferiores.
   • Presupuestos Piramidales (CONF-KV-L): Una variante opcional asigna presupuestos de manera no uniforme a través de las capas, asignando presupuestos más pequeños a las capas más profundas basándose en la observación de que las capas más profundas a menudo concentran la información en menos tokens.
   • Implementación: El gestor utiliza compactación contigua para mantener un diseño denso para los kernels de atención, evitando tablas de indirectas, y fusiona la dequantización en las lecturas de atención por bloques.

Contribuciones Clave

• Política Consciente de la Confianza: Una estrategia de expulsión novedosa que utiliza la distribución de salida actual para impulsar las decisiones de presupuesto de caché por paso, en lugar de depender únicamente de señales históricas.
• Diseño de Sistemas: Una implementación completa que combina expulsión adaptativa con atención de softmax en línea por bloques, almacenamiento mixto FP16/INT8 y presupuestos de capas piramidales opcionales.
• Validación Mecanística: Una prueba que demuestra que la puntuación de confianza se anticorrelaciona fuertemente (Pearson $r = -0.77$) con el desplazamiento de KL inducido por la ablación del contexto reciente, validando la suposición de que una baja confianza indica la necesidad de más contexto retenido.
• Evaluación Exhaustiva: Benchmarks extensos en cuatro familias de modelos (GPT-2, Qwen-14B/32B, gpt-oss-20b) y cargas de trabajo diversas, incluyendo perplejidad, recuperación de contexto largo y tareas de agentes web.

Resultados Experimentales

• Compromiso Memoria-Calidad: En GPT-2 con una longitud de generación de 2048 tokens, CONF-KV+INT8 iguala la huella de memoria de una ventana deslizante fija de 512 tokens (~38.7 MB) pero reduce la perplejidad en 3.11 puntos en comparación con la ventana deslizante. CONF-KV-L (variante piramidal) reduce aún más la memoria a 34.2 MB mientras mejora la perplejidad en 3.89 puntos, cerrando el 74% de la brecha de calidad entre una ventana deslizante y KV completo.
• Recuperación de Contexto Largo: En pruebas de "Aguja en un Pajonal" (NIAH) hasta 32K tokens, CONF-KV logra una precisión de recuperación del 91.4%, superando significativamente a las ventanas deslizantes (53.8%) y a H2O (80.6%).
• Tareas de Agente: En 75 tareas de VisualWebArena, CONF-KV retiene el 95.3% de la tasa de éxito de KV completo mientras reduce la memoria pico en 2.8 veces.
• Rendimiento y Latencia: CONF-KV reduce la latencia por paso en 1.8 veces en comparación con KV completo en GPT-2 y Qwen-32B. El rendimiento escala con el tamaño del lote; con un tamaño de lote de 8, CONF-KV alcanza 2.06 veces el rendimiento de KV completo.
• Ablación: Aislar la señal de confianza muestra que la expulsión aleatoria a la misma tasa degrada el rendimiento a 36.54 PPL, mientras que CONF-KV completo alcanza 30.92 PPL, confirmando que tanto el control de puerta de confianza como el mecanismo de clasificación son esenciales.

Significado y Afirmaciones
El artículo afirma que CONF-KV demuestra que la incertidumbre actual es un metadato de sistemas útil para mejorar la frontera de Pareto memoria-calidad. El método es distinto porque es prospectivo: no puede recuperar tokens expulsados, pero previene la expulsión prematura durante span difíciles. El enfoque está diseñado para ser libre de entrenamiento, no altera los pesos del modelo y es ortogonal a técnicas existentes como la asignación por cabeza, la asignación por capa y la selección de precisión, lo que permite combinarlo con ellas.

Los autores señalan limitaciones, incluyendo que el método es una "operación nula" en contextos cortos que nunca superan el umbral de expulsión y que la compactación contigua incurre en un costo de recolección (aunque esto se ve superado por los ahorros de atención). También reconocen que la señal de confianza se vuelve menos informativa bajo muestreo de alta temperatura. Finalmente, el artículo destaca una consideración de doble uso: si bien menores costos de inferencia mejoran la eficiencia, también pueden reducir las barreras para la automatización indeseable de horizonte largo, lo que sugiere que los controles de seguridad para los modelos subyacentes siguen siendo necesarios.