PolyKV: A Shared Asymmetrically-Compressed KV Cache Pool for Multi-Agent LLM Inference

1. Enunciado del Problema

La inferencia de Modelos de Lenguaje Grande (LLM) está dominada por cuellos de botella de memoria causados por la memoria caché de claves y valores (KV), que almacena los estados de atención para cada token en el contexto. Este problema se ve exacerbado en sistemas de inferencia multiagente, donde $N$ agentes que procesan el mismo contexto de documento compartido requieren tradicionalmente $N$ memorias caché KV independientes y de precisión completa. Esto resulta en una escalabilidad de memoria lineal ($O(N)$), lo que hace ineficiente ejecutar múltiples agentes en contextos largos.

Las soluciones existentes abordan esto de forma aislada:

• Compresión KV: Reduce la memoria por caché mediante cuantización, pero opera sobre cachés aisladas por solicitud.
• Compartición Multiagente: Reutiliza prefijos comunes entre agentes, pero típicamente utiliza memorias caché de precisión completa, sin aprovechar la compresión.

PolyKV busca resolver la intersección de estos problemas: crear un único grupo de memoria caché KV compartido y comprimido con pérdida que pueda ser accedido simultáneamente por múltiples agentes concurrentes.

2. Metodología

PolyKV introduce una arquitectura novedosa que combina cuantización asimétrica, compresión basada en rotación y una abstracción de memoria compartida.

A. Arquitectura Central

• SharedKVPool: Una abstracción de escritura única, lectura múltiple. Calcula el estado KV completo para un contexto de documento compartido una sola vez, lo comprime y lo almacena como un único objeto (complejidad de memoria $O(1)$ en relación con la cantidad de agentes).
• PooledAgent: Agentes de inferencia individuales que reciben una referencia al SharedKVPool. En el momento de la inferencia, cada agente inyecta los tensores KV descomprimidos en su propia instancia fresca de DynamicCache, capa por capa, eludiendo la acumulación estándar por paso de avance. Esto elimina la sobrecarga de copia por agente.

B. Esquema de Compresión Asimétrica (TurboQuant MSE)

PolyKV emplea una estrategia de cuantización asimétrica adaptada a la sensibilidad de las Claves (K) y los Valores (V):

1. Claves (K): Cuantizadas a int8 (q8_0) utilizando cuantización lineal por tensor. Esto preserva la estabilidad de las puntuaciones de atención softmax, que son altamente sensibles a la cuantización de bajo bit.
2. Valores (V): Comprimidos utilizando TurboQuant MSE a precisión de 3 bits. Esto implica:
   • Rotación: Aplicar una Transformada Rápida de Walsh-Hadamard (FWHT) para rotar el vector de entrada, distribuir los valores atípicos y hacer que la distribución converja hacia una Gaussiana $N(0, 1)$.
   • Cuantización: Mapear las coordenadas rotadas a centroides Lloyd-Max precalculados (optimizados para $N(0, 1)$).
   • De-cuantización: Recuperar los valores de los centroides y aplicar la FWHT inversa.

Ratio de Compresión:
Con Claves a 8 bits y Valores a 3 bits (comparado con una línea base de bfloat16 de 16 bits), el ratio de compresión teórico es:
$$\text{Ratio} = \frac{16}{(8 + 3)/2} = \frac{16}{5.5} \approx 2.91\times$$

3. Contribuciones Clave

• Arquitectura de Grupo Compartido: La primera implementación de un modelo de memoria KV comprimido de escritura única, lectura múltiple para agentes concurrentes, eliminando la escalabilidad de memoria $O(N)$ de los sistemas multiagente anteriores.
• TurboQuant Asimétrico: Una combinación específica de Claves int8 y Valores rotados por FWHT de 3 bits que logra un ratio de compresión estable de 2.91× a través de diferentes arquitecturas de modelos.
• Validación Trans-modelo: Demostró estabilidad tanto en SmolLM2-1.7B como en Llama-3-8B (incluyendo arquitectura GQA), probando que el método es agnóstico a la arquitectura.
• Escalabilidad: Escaló con éxito a 15 agentes concurrentes compartiendo un único contexto, con ahorros de memoria que crecen linealmente con la cantidad de agentes (hasta una reducción del 97.7%).
• Hallazgo de Inversión de PPL: Descubrió que, para contextos coherentes y largos, la memoria caché comprimida puede generar en realidad una perplejidad menor que la línea base de precisión completa, sugiriendo que el ruido de cuantización actúa como regularización implícita.

4. Resultados Experimentales

El sistema fue evaluado en dos modelos, tres longitudes de contexto (de 600 a 7,194 tokens) y hasta 15 agentes.

Eficiencia de Memoria

• Llama-3-8B (15 agentes, contexto de 4K): Redujo la memoria de la caché KV de 19.8 GB (precisión completa) a 0.45 GB.
• Reducción: 97.7% de reducción de memoria manteniendo una alta calidad de salida.
• Escalabilidad: El uso de memoria para el grupo permanece constante ($O(1)$) independientemente del número de agentes (ej. 0.116 GB para 3 agentes vs. 10 agentes en contexto de 2K).

Métricas de Calidad

• Perplejidad (PPL):
  • Estabilidad: El delta de PPL es invariante al número de agentes.
  • Tendencia de Longitud de Contexto: El delta de PPL mejora a medida que aumenta la longitud del contexto.
    • En 2K tokens: +1.59% de degradación.
    • En 4K tokens: +0.57% de degradación.
    • Inversión: En 1,851 tokens coherentes en SmolLM2-1.7B, la memoria caché comprimida logró un delta de PPL de -0.26% (mejor que la línea base).
• Similitud Semántica (BERTScore):
  • Las puntuaciones F1 medias oscilaron entre 0.928 y 0.970 en todas las configuraciones.
  • La superposición de tokens fue perfecta (1.000) en las pruebas de SmolLM2.
  • La ligera degradación de PPL no se tradujo en pérdida semántica; las variaciones de redacción preservaron la equivalencia factual.

5. Significado e Hipótesis

Significado:
PolyKV es el primer sistema que combina un único grupo KV compartido y comprimido con pérdida con acceso de múltiples lectores concurrentes. Cambia fundamentalmente la economía de la inferencia multiagente de LLM, permitiendo una concurrencia masiva en contextos compartidos con una huella de memoria mínima, haciendo viable ejecutar enjambres de agentes complejos en hardware de nivel de consumidor (ej. GPUs T4 de Kaggle).

La Hipótesis de Regularización:
Los autores hipotetizan que el ruido de FWHT de TurboQuant actúa como regularización implícita. En documentos largos y coherentes, las cabezas de atención a menudo atienden a correlaciones redundantes y espurias (sumideros de atención). El ruido de cuantización introducido por la rotación FWHT interrumpe estos patrones espurios, efectivamente "desruidando" el mecanismo de atención. Esto explica por qué la PPL mejora (o se degrada menos) a medida que aumentan la longitud del contexto y la coherencia léxica.

6. Limitaciones y Trabajo Futuro

• Escala: Los experimentos se limitaron a modelos de 1.7B y 8B; el comportamiento en 70B+ es desconocido.
• Métricas: No se reportaron el tiempo hasta el primer token (TTFT) ni el rendimiento (throughput).
• Validación: El mecanismo de inversión de PPL requiere estudios de ablación controlados para confirmar la causalidad.
• Futuro: Los planes incluyen probar en modelos más grandes (Qwen2.5-7B, 70B), estandarizar los protocolos de evaluación WikiText-2 y medir la latencia de extremo a extremo.