Resumen Técnico: RDKV - Asignación de Bits de Tasa-Distorsión para Evicción y Cuantización Conjunta de la Caché KV

1. Declaración del Problema

Los Modelos de Lenguaje Grande (LLM) enfrentan un cuello de botella crítico durante la inferencia de contexto largo: el ancho de banda y la capacidad de memoria. El tamaño de la caché de Clave-Valor (KV) crece linealmente con la longitud de la secuencia, requiriendo la relectura de decenas de gigabytes de datos desde la Memoria de Alto Ancho de Banda (HBM) fuera del chip hacia la memoria dentro del chip en cada paso de decodificación. Esto convierte la inferencia en un proceso limitado por la memoria en lugar de por el cómputo.

Las estrategias de compresión existentes generalmente caen en dos categorías aisladas:

1. Evicción: Puntuación de tokens o canales y descarte total de los de menor rango (mantener o descartar binario). Esto no deja ningún nivel de precisión para tokens de importancia moderada.
2. Cuantización: Retención de todos los tokens pero reduciendo su ancho de bits. Esto retiene incluso tokens negligentes, desperdiciando memoria.

Aunque híbridos recientes intentan combinar estas técnicas, a menudo dependen de puntuaciones de enrutamiento heurísticas o ratios preasignados determinados mediante búsqueda fuera de línea, fallando en tratar la evicción y la cuantización como un problema de optimización unificado.

2. Metodología: RDKV

El artículo propone RDKV (Compresión de caché KV de Tasa-Distorsión), un marco que formula la compresión de la caché KV como un problema de tasa-distorsión. En esta visión, la evicción y la cuantización no son acciones separadas, sino dos extremos de una única curva de asignación de bits.

2.1 Formulación Teórica

RDKV deriva un peso de distorsión para cada unidad de caché (token en la caché V, canal en la caché K) basado en la perturbación inducida por su eliminación o cuantización en el cálculo de atención.

• Caché V (Tokens): El peso de distorsión $w_t$ para un token es la distancia de variación total de la distribución de atención después de la evicción. Se calcula como la suma de las puntuaciones de atención recibidas por el token a través de todas las posiciones de consulta: $w_t = \sum_{\tau} a_{\tau,t}$.
• Caché K (Canales): El peso de distorsión $w_c$ para un canal es la norma espectral de la desviación de logits causada por poner ese canal en cero. Se calcula como el producto de las normas de las columnas correspondientes de consulta y clave: $w_c = \frac{1}{\sqrt{d}} \|Q_{:,c}\|_2 \|K_{:,c}\|_2$.

Estos pesos sirven como coeficientes de costo en una función objetivo de tasa-distorsión:
$$\min \sum_{u} w_u \sigma_u 2^{-b_u} \quad \text{sujeto a} \quad \sum_{u} b_u \leq B$$
donde $b_u$ es el ancho de bits asignado a la unidad $u$, y $B$ es el presupuesto total de bits.

2.2 Pipeline de Asignación de Bits

La asignación se resuelve mediante llenado inverso de agua (relajación Lagrangiana) adaptado para anchos de bits discretos soportados por hardware $\mathcal{B} = \{0, 2, 4, 8, 16\}$.

1. Etapa 1 (Ponderación): Calcular los pesos de tokens ($w_t$) y los pesos de canales ($w_c$) una vez después de la etapa de prellenado utilizando la caché completa sin comprimir.
2. Etapa 2 (Asignación de Tokens): Resolver un Problema de la Mochila de Elección Múltiple (MCKP) para la caché V. Los tokens se asignan anchos de bits desde $\{0, 2, 4, 8, 16\}$. Un ancho de bits de 0 corresponde a la evicción.
3. Etapa 3 (Asignación de Canales): Resolver el MCKP para la caché K, pero solo para los tokens retenidos en la Etapa 2. Los canales se asignan anchos de bits independientemente del ancho de bits V del token, permitiendo que un solo token tenga componentes K y V de precisión mixta.

2.3 Decodificación Empaquetada Eficiente (TriZone)

Para lograr ahorros reales de memoria y latencia, RDKV introduce una disposición de caché TriZone que fusiona la dequantización en el kernel de atención, evitando la necesidad de desempaquetar datos a FP16 antes del procesamiento.

• Zona A: Flujo de caché antigua empaquetado que contiene filas K y filas V retenidas con anchos de bits $\{2, 4, 8\}$. Estas se agrupan por ancho de bits para permitir rutas de dequantización uniformes.
• Zona B: Filas V en FP16 para tokens asignados a precisión de 16 bits.
• Zona C: Filas K y V en FP16 para tokens recién generados.

El kernel de atención calcula logits y salidas combinando estas tres zonas, con la dequantización fusionada directamente en la multiplicación de matrices, eliminando la sobrecarga adicional de ancho de banda de memoria.

3. Contribuciones Clave

1. Formulación de Tasa-Distorsión: El artículo plantea la compresión de la caché KV como un problema unificado de tasa-distorsión, demostrando que la solución óptima mezcla inherentemente unidades de tasa cero (evicción) y unidades de tasa finita (cuantización). Esto elimina la necesidad de ratios heurísticos o búsquedas fuera de línea.
2. Asignación Discreta mediante Llenado Inverso de Agua: Los autores instan la asignación como un problema de mochila discreto resoluble mediante relajación Lagrangiana, asignando anchos de bits óptimos desde $\{0, 2, 4, 8, 16\}$ basados en pesos de distorsión derivados de la dinámica de atención.
3. Disposición Empaquetada TriZone: Una implementación consciente del hardware que fusiona la dequantización en el kernel de atención, convirtiendo los ahorros teóricos de ancho de bits en reducciones reales de la memoria pico y la latencia de decodificación.

4. Resultados Experimentales

Los experimentos se realizaron en LLaMA-3.1-8B-Instruct, Mistral-7B, Qwen3-4B, LLaMA-2-13B y Qwen2.5-72B a través de LongBench, RULER e InfiniteBench.

• Precisión: RDKV supera a la mejor línea base evaluada (Ada-SnapKV) en un 9.1% en promedio. En LongBench, recupera el 97.81% de la precisión de caché completa con solo un 2.48% de retención de caché.
• Rendimiento de Contexto Largo:
  • RULER: Con un contexto de 128K, RDKV alcanza una precisión del 66.95% en comparación con el 58.43% de la línea base más fuerte (AdaKV), una brecha de 8.5 puntos.
  • InfiniteBench: RDKV ocupa el primer lugar entre los métodos de compresión con una puntuación promedio de 39.46, superando a AdaKV por 1.40 puntos.
  • Aguja en un Pajareo: RDKV mantiene una precisión de recuperación casi uniforme (promedio 0.99) a través de las profundidades, mientras que las líneas base solo de evicción fallan significativamente en profundidades intermedias.
• Eficiencia:
  • Latencia: Con un contexto de 128K, RDKV logra una aceleración de decodificación de 4.5× en comparación con la decodificación FlashAttention-2 de caché completa.
  • Memoria: Logra una reducción de memoria pico de 1.9× en un contexto de 128K y ejecuta con éxito un contexto de 256K en una sola GPU A100 de 64GB, donde FullKV se queda sin memoria (OOM).

5. Significado y Afirmaciones

El artículo afirma que tratar la evicción y la cuantización conjuntamente dentro de un único marco de optimización es superior a los enfoques por etapas o basados en ratios. Al derivar puntuaciones de asignación directamente de la distorsión inducida en el cálculo de atención, RDKV evita la "trampa binaria" de los métodos solo de evicción (que descartan tokens moderados) y el "desperdicio" de los métodos solo de cuantización (que retienen tokens negligentes).

Los autores enfatizan que su método:

• Requiere ningún reentrenamiento del modelo.
• Es agnóstico a la arquitectura, desempeñándose consistentemente a través de diferentes tamaños y familias de modelos.
• Proporciona un compromiso óptimo de Pareto entre latencia y precisión, superando estrictamente a las líneas base en ambas métricas en longitudes de contexto largas.
• Demuestra que el "óptimo continuo" de la teoría de tasa-distorsión, cuando se discretiza para hardware, produce naturalmente una estrategia de bits mixtos que supera a los híbridos actuales del estado del arte.

El artículo concluye que este enfoque conjunto mejora la precisión sin sacrificar la eficiencia, habilitando la inferencia de contexto largo en hardware con recursos limitados.