Resumen Técnico: OSCAR - Rotación Consciente de la Covarianza Espectral Offline para Cuantización de Cache KV de 2 bits

1. Declaración del Problema

La inferencia de contexto largo en Modelos de Lenguaje Grandes (LLM) está cada vez más limitada por la huella de memoria del cache de Clave-Valor (KV). A medida que crecen la longitud del contexto, el tamaño del lote y la profundidad del modelo, el cache consume memoria GPU significativa, limitando el rendimiento y la escalabilidad. Aunque comprimir el cache KV a precisión ultra-baja (por ejemplo, INT2) ofrece una reducción de memoria de 8×, enfrenta un desafío crítico de precisión: las activaciones KV contienen valores atípicos severos a nivel de canal.

Bajo cuantización de pocos bits, estos valores atípicos dominan la escala de cuantización, forzando a la mayoría de los valores "normales" a unos pocos niveles de cuantización efectivos y degradando severamente la calidad de la atención. Los métodos existentes basados en rotación (por ejemplo, transformadas de Hadamard) intentan suavizar estos valores atípicos redistribuyendo la energía a través de los canales. Sin embargo, estos métodos son ciegos a los datos; suavizan la distribución cruda del cache pero no logran alinearse con las direcciones específicas que los mecanismos de atención aguas abajo realmente consumen. En INT2, donde solo existen cuatro niveles de cuantización, esta falta de alineación provoca que los métodos de rotación ingenuos colapsen, resultando en una precisión cercana a cero en tareas de razonamiento.

2. Metodología: OSCAR

Los autores proponen OSCAR (Rotación Consciente de la Covarianza Espectral Offline), un método que estima estructuras de covarianza conscientes de la atención de forma offline para derivar rotaciones fijas y umbrales de recorte. A diferencia de trabajos anteriores que minimizan el error de reconstrucción euclidiano crudo ($\|K - \hat{K}\|_F$), OSCAR minimiza la distorsión en las puntuaciones y salidas de atención aguas abajo.

2.1 Motivación Teórica

El artículo argumenta que el objetivo de rotación óptimo debe derivarse de las estadísticas de la propia computación de la atención:

• Para Claves ($K$): La distorsión en los logits de atención está controlada por la covarianza de la consulta $Q^\top Q$, no solo por la covarianza de la clave $K^\top K$.
• Para Valores ($V$): La distorsión en la salida de la atención está ponderada por las puntuaciones de atención $S$, lo que conduce a un objetivo de covarianza de $V^\top S^\top S V$.

2.2 Diseño del Algoritmo

OSCAR opera en dos fases:

A. Fase de Calibración Offline:

1. Estimación de Covarianza: Utilizando un pequeño conjunto de datos de calibración, el método calcula covarianzas objetivo empíricas:
   • $C_Q = \frac{1}{N} \sum q_n^\top q_n$ (Objetivo de clave consciente de la consulta).
   • $C_S = \frac{1}{N} V^\top S^\top S V$ (Objetivo de valor consciente de la puntuación).
2. Derivación de Rotación: El método realiza una descomposición espectral sobre estas matrices para encontrar las bases propias óptimas ($U_Q$ y $U_S$).
3. Construcción de Rotación Compuesta: Las matrices de rotación finales ($R_K, R_V$) se construyen como un producto de tres componentes:
   • $U_Q/U_S$: La base propia consciente de la atención que alinea los canales con las direcciones de importancia.
   • $H_{Had}$: Una transformada de Hadamard que iguala la métrica de importancia entre canales y dispersa la energía de los valores atípicos.
   • $P_{br}$: Una permutación de inversión de bits que entrelaza canales de alta y baja varianza para garantizar rangos dinámicos equilibrados dentro de los grupos de cuantización.
   • Fórmula: $R_K = U_Q H_{Had} P_{br}$ y $R_V = U_S H_{Had} P_{br}$.
4. Umbrales de Recorte: Los umbrales de recorte por capa se estiman a partir de las activaciones de calibración rotadas para controlar los valores atípicos.

B. Fase de Servicio Online:

• Diseño de Precisión Mixta: El sistema mantiene un cache KV paginado donde los tokens "sumidero" (iniciales) y los tokens "recientes" se mantienen en precisión completa BF16. El contexto histórico intermedio se almacena en INT2 después de aplicar la rotación fija de OSCAR y el recorte.
• Integración de Kernel: Los autores desarrollaron un kernel fusionado personalizado en Triton para el marco de servicio SGLang. Este kernel maneja el desempaquetado, la de-cuantización, la rotación inversa y la computación de atención para el historial INT2, garantizando compatibilidad con la atención paginada y la caché de prefijos sin introducir tráfico de memoria adicional.

3. Contribuciones Clave

1. Identificación del Objetivo Faltante: El artículo identifica que las rotaciones genéricas dispersan valores atípicos pero fallan en INT2 porque no se alinean con la atención aguas abajo. El objetivo de rotación debe ser inducido por las estadísticas de atención ($Q^\top Q$ y $V^\top S^\top S V$), no por la reconstrucción cruda del cache.
2. Marco OSCAR: Un sistema de cuantización de cache KV en INT2 totalmente desplegable que utiliza calibración offline para derivar rotaciones conscientes de la atención. El análisis teórico muestra que estos objetivos son óptimos bajo una suposición de sustituto de error congelado.
3. Sistema Listo para Producción: Un kernel de atención INT2 personalizado integrado en SGLang que admite diseños de cache KV paginados y caché de prefijos. Logra una integración fluida con marcos de servicio modernos (SGLang, vLLM) mientras mantiene un alto rendimiento.

4. Resultados Experimentales

Los autores evaluaron OSCAR en modelos de razonamiento (Qwen3-4B, Qwen3-8B, Qwen3-32B) y un modelo de vanguardia (GLM-4.7, 358 mil millones de parámetros) en cinco benchmarks (AIME25, GPQA, HumanEval, LiveCodeBench, MATH500) y recuperación de contexto largo (RULER-NIAH hasta 128k).

• Preservación de Precisión:

  • En Qwen3-4B-Thinking y Qwen3-8B, OSCAR reduce la brecha de precisión respecto a BF16 a 3.78 y 1.42 puntos, respectivamente. En contraste, la rotación ingenua INT2 colapsa a una precisión casi nula.
  • En Qwen3-32B y GLM-4.7, OSCAR permanece efectivamente a la par con BF16 (dentro de 0.02–0.27 puntos), mientras que otras bases de referencia INT2 sufren una degradación significativa.
  • En recuperación de contexto largo (RULER-NIAH), OSCAR mantiene un rendimiento robusto hasta 128k tokens, mientras que los métodos INT2 ingenuos y solo Hadamard colapsan a longitudes mucho más cortas (por ejemplo, 16k).

• Rendimiento del Sistema:

  • Memoria: Reduce la huella de memoria del cache KV en aproximadamente 8× en comparación con BF16.
  • Rendimiento: Mejora el rendimiento hasta en 7× en tamaños de lote grandes bajo el mismo presupuesto de memoria.
  • Latencia: Acelera la decodificación con tamaño de lote 1 hasta en 3× sobre BF16 debido a la reducción de la sobrecarga de ancho de banda de memoria.
  • Escalabilidad: Demuestra la capacidad de manejar hasta 28 solicitudes concurrentes con entradas de 100k tokens en una sola H100, un régimen donde las bases de referencia BF16 e INT4 se quedan sin memoria.

5. Significado y Afirmaciones

El artículo afirma que OSCAR cierra exitosamente la brecha entre la promesa teórica del cache KV de ultra-bajos bits (INT2) y la implementabilidad práctica. Al desplazar el objetivo de rotación de la reconstrucción de datos crudos a la covarianza consciente de la atención, OSCAR recupera la precisión perdida por la cuantización ingenua mientras retiene los beneficios masivos de memoria y rendimiento del almacenamiento de 2 bits.

Los autores enfatizan que su enfoque no es meramente una mejora algorítmica, sino una solución a nivel de sistema. Al proporcionar un kernel personalizado que se integra con la atención paginada y la caché de prefijos, demuestran que la cuantización INT2 puede utilizarse en marcos de servicio de producción (como SGLang) sin sacrificar las características de eficiencia (como la reutilización de prefijos) en las que se basa el servicio moderno de LLM. El trabajo sugiere que para el razonamiento de contexto largo, la clave de una cuantización exitosa radica en alinear la estrategia de compresión con las demandas estadísticas específicas del mecanismo de atención.