Resumen Técnico: KVzap - Poda del KV Cache Rápida, Adaptativa y Fiel

1. Planteamiento del Problema

A medida que las longitudes de contexto en los modelos de lenguaje basados en transformadores se expanden a decenas o cientos de miles de tokens, el caché de Clave-Valor (KV) se ha convertido en un cuello de botella dominante para la inferencia eficiente de Modelos de Lenguaje Grande (LLM). El caché KV, que almacena estados de atención intermedios, consume una cantidad significativa de memoria GPU (por ejemplo, 335 GB para un modelo de 65B con una longitud de secuencia de 128k) y aumenta el tiempo hasta el primer token (time-to-first-token), reduciendo simultáneamente el rendimiento de decodificación (throughput).

Aunque las modificaciones arquitectónicas (por ejemplo, Atención de Consulta Agrupada, Atención Latente Multicabeza) y los mecanismos de atención dispersa han abordado la compresión a lo largo de los ejes de cabezal ($H$), dimensión ($D$) y capa ($L$), el eje de longitud de secuencia ($T$) permanece en gran medida sin abordar por los métodos ampliamente adoptados. Los métodos de poda ad-hoc actuales suelen fallar en cumplir con cuatro criterios críticos para el despliegue en producción:

1. Rápido y ligero: Sobrecarga insignificante.
2. Agnóstico a la fase: Aplicable tanto al prellenado (contexto largo) como a la decodificación (razonamiento).
3. Amigable con la optimización: Compatible con kernels como FlashAttention2 y PagedAttention.
4. Fiel: Degradación mínima de la precisión.

Los métodos actuales de vanguardia como KVzip logran una alta compresión pero requieren un paso de forward doble computacionalmente costoso en prompts extendidos, lo que los hace inadecuados para la inferencia en línea y tareas de decodificación.

2. Metodología

Los autores presentan KVzap, una aproximación rápida y adaptativa al input de KVzip que opera tanto durante el prellenado como durante la decodificación. La metodología consta de tres componentes principales:

2.1. Política de Puntuación Mejorada (KVzip+)

Los autores primero refinan el mecanismo de puntuación de KVzip. KVzip puntúa los tokens basándose en la "reconstrucción de contexto", midiendo el peso de atención máximo que un token recibe cuando el prompt se repite. Los autores mejoran esto incorporando un término de normalización inspirado en Devoto et al. (2025), que tiene en cuenta la contribución del vector de valor de un token al flujo residual.
La puntuación KVzip+ ($s^+_i$) se define como:
$$s^+_i = \max_{j \in \text{<prompt>}} \frac{a_{ji}}{\|W_O v_i\| / \|h_j\|}$$
donde $a_{ji}$ es el peso de atención, $W_O$ es la matriz de proyección de salida, $v_i$ es el vector de valor y $h_j$ es el estado oculto.

2.2. Modelado Sustituto (Surrogate Modeling)

Para evitar el costo prohibitivo del paso de forward doble de KVzip, KVzap entrena modelos sustitutos ligeros (ya sea una capa lineal o un MLP de dos capas) para predecir los log-scores ($\log(s^+)$) directamente desde los estados ocultos de entrada del modelo ($h$).

• Entrada: Estados ocultos $h_t \in \mathbb{R}^{D_h}$.
• Salida: Puntuaciones para cada cabezal KV en $\mathbb{R}^H$.
• Entrenamiento: Los modelos son entrenados con 1.2 millones de pares de (estado oculto, log-score) muestreados del Nemotron-Pretraining-Dataset. La variante MLP utiliza un ancho oculto de $D_h/8$ con activación GELU.
• Desempeño: Los sustitutos logran una correlación $R^2$ promedio de 0.60–0.80 con el oráculo KVzip+, demostrando que los scores costosos pueden aproximarse con precisión desde los estados ocultos.

2.3. Poda y Evicción Adaptativa

A diferencia de KVzip, que impone un presupuesto fijo (por ejemplo, mantener exactamente el 50% de los tokens), KVzap emplea poda basada en umbrales:

• Umbralización (Thresholding): Los tokens con puntuaciones predichas por debajo de un umbral fijo $\tau$ son descartados. Esto permite que la relación de compresión se adapte dinámicamente a la densidad de información del prompt de entrada.
• Ventana Deslizante (Sliding Window): Para preservar el contexto local, los últimos $w=128$ tokens se mantienen siempre, siguiendo el enfoque de StreamingLLM.
• Implementación: La poda ocurre después de la operación de atención. El proceso es computacionalmente eficiente, añadiendo una sobrecarga insignificante (limitada a 1.1% para modelos MLP y 0.02% para modelos lineales).

3. Contribuciones Clave

1. KVzip+: Una métrica de puntuación mejorada para la importancia del caché KV que incorpora la normalización del estado oculto.
2. Aproximación Sustituta: Una demostración de que las puntuaciones de KVzip+ pueden ser efectivamente aproximadas por modelos ligeros entrenados en estados ocultos, eliminando la necesidad de pasos de forward dobles.
3. KVzap: Una nueva técnica de poda que aplica estos sustitutos con evicción basada en umbrales y ventanas deslizantes, logrando un rendimiento de vanguardia mientras cumple con los cuatro criterios prácticos (velocidad, agnosticismo de fase, amigabilidad con la optimización y fidelidad).

4. Resultados Experimentales

Los experimentos se realizaron en Qwen3-8B, Llama-3.1-8B-Instruct y Qwen3-32B a través de tres benchmarks: RULER (recuperación/agregación de contexto largo), LongBench (multitarea de contexto largo) y AIME25 (razonamiento matemático).

• Compresión vs. Precisión: KVzap logra una compresión del caché KV de 2–4× con una pérdida de precisión insignificante.
  • En RULER 4k, KVzap iguala o supera el rendimiento del oráculo KVzip+ y supera a otros 15 métodos de poda (incluyendo Expected Attention y Duo Attention).
  • En LongBench, KVzap mantiene una precisión casi perfecta hasta una compresión de 2–3×.
  • En AIME25, KVzap-MLP preserva la precisión del razonamiento incluso descartando más del 50% del caché KV.
• Adaptabilidad: El enfoque basado en umbrales ajusta automáticamente las relaciones de compresión según la dificultad de la tarea. Por ejemplo, logra una mayor compresión en datos sintéticos de RULER en comparación con los datos de LongBench, que son más densos y del mundo real.
• Eficiencia: El método añade una sobrecarga de cómputo y memoria insignificante. Durante la decodificación, los FLOPs adicionales utilizan ciclos de GPU que de otro modo estarían estancados por las limitaciones de ancho de banda de memoria.

5. Significado y Reivindicaciones

El artículo afirma que KVzap cierra la brecha entre la investigación académica de poda y los motores de inferencia del mundo real. Su importancia radica en:

• Practicidad: Es el primer método que logra una compresión de vanguardia sin sacrificar la velocidad ni requerir cambios arquitectónicos, lo que lo hace compatible con los principales motores de inferencia (vLLM, SGLang, TRT-LLM).
• Fidelidad: Demuestra que los LLM no explotan completamente el caché KV y que los pares no utilizados pueden identificarse mediante estados ocultos sin reentrenar el modelo base de extremo a extremo.
• Escalabilidad: El enfoque es validado en modelos que van desde 8B hasta 32B de parámetros y está diseñado para ser compatible con futuras arquitecturas de atención dispersa.

Los autores reconocen las limitaciones, señalando que KVzap es una adición post-hoc y no un objetivo entrenado de extremo a extremo, y que se requiere más ingeniería para traducir las relaciones de compresión en mejoras de tiempo de ejecución mediante kernels optimizados (por ejemplo, PagedAttention). Sin embargo, postulan que la combinación de simplicidad, alta compresión y rendimiento robusto de KVzap lo convierte en un candidato principal para el despliegue en producción.