Resumen Técnico: Spherical KV

1. Planteamiento del Problema

La inferencia de contexto largo en los Modelos de Lenguaje de Gran Escala (LLMs) está cada vez más limitada por el caché de Clave-Valor (KV) en lugar de por el rendimiento aritmético. A medida que la longitud del contexto ($T$) crece, la huella de memoria residente escala linealmente con $T$, el tamaño de lote ($B$), las capas ($L$) y las cabezas ($H$). En entornos de servicio de producción que utilizan memoria paginada (por ejemplo, PagedAttention), el cuello de botella dominante se desplaza de las Operaciones de Punto Flotante (FLOPs) al tráfico de Memoria de Alto Ancho de Banda (HBM). Específicamente, la decodificación se ve limitada por el flujo repetido de páginas KV históricas desde la memoria del dispositivo hacia las unidades de cómputo.

Las soluciones existentes intentan mitigar esto mediante tres palancas:

1. Retención/Evicción: Almacenar menos tokens (por ejemplo, StreamingLLM, H2O).
2. Cuantización/Compresión: Almacenar menos bits (por ejemplo, KVQuant, KIVI).
3. Offloading: Almacenar el KV en otro lugar (por ejemplo, CPU/Almacenamiento).

Sin embargo, estos métodos a menudo no logran abordar completamente el cuello de botella de la ruta crítica. Muchas técnicas de compresión requieren reconstruir vectores densos a partir de representaciones comprimidas antes de computar los logits de atención (una ruta de "descomprimir-luego-producto-punto"). Esta reconstrucción reintroduce el tráfico de HBM y evita la fusión de kernels, anulando los ahorros teóricos de memoria en sustratos de servicio paginados realistas. Además, las políticas de evicción simples suelen carecer de la granularidad necesaria para preservar información crítica (por ejemplo, evidencia recuperada) mientras se descartan estados de baja utilidad.

2. Metodología: Spherical KV

El artículo propone Spherical KV, un método de inferencia de contexto largo que trata la asignación de KV como un problema de tasa–distorsión fundamentado en la geometría de la atención. Consta de dos componentes ortogonales diseñados para trabajar dentro de un sustrato de servicio irregular y paginado: Atención de Dominio Angular (ADA) y Retención de Tasa–Distorsión (RDR).

2.1 Atención de Dominio Angular (ADA)

ADA reemplaza la primitiva estándar de "descomprimir-luego-producto-punto" con un enfoque nativo de cómputo que opera directamente sobre códigos esféricos compactos.

• Parametrización Esférica: Las claves se descomponen en un radio escalar ($r = \|k\|$) y ángulos direccionales ($\phi = \Phi(\hat{k})$).
• Similitud en el Dominio Comprimido: En lugar de almacenar vectores densos, ADA almacena radios cuantizados ($\tilde{r}$) y códigos angulares empaquetados ($c_\theta$). Durante el bucle caliente de decodificación, el kernel de atención computa los logits directamente desde estos códigos:
  $$\hat{\ell}(q, k) \approx \|q\| \cdot \tilde{r} \cdot f_b(c_\theta^q, c_\theta^k)$$
  donde $f_b$ es una operación de kernel fusionado que computa la similitud de coseno a partir de los códigos angulares.
• Invariante de No Reconstrucción: La restricción de diseño crítica es que el kernel de decodificación nunca materializa un vector de clave denso ($\tilde{k} \in \mathbb{R}^d$) en HBM. Esto asegura que las reducciones en los bytes de KV almacenados se traduzcan directamente en una reducción del tráfico de HBM por token.

2.2 Retención de Tasa–Distorsión (RDR)

RDR es un mecanismo de control que optimiza conjuntamente qué estados retener y con qué precisión (nivel) almacenarlos, sujeto a un presupuesto estricto de memoria.

• Optimización Conjunta: A diferencia de los métodos que tratan la evicción y la cuantización por separado, RDR resuelve un problema de asignación de tasa–distorsión discreto. Para cada estado de token/cabeza/capa $i$, selecciona una decisión de residencia binaria ($z_i \in \{0, 1\}$) y un nivel de precisión ($b_i \in \mathcal{B}$).
• Sustituto de Distorsión: El controlador utiliza un sustituto de distorsión mecanístico derivado del análisis de error de ADA. Estima que el error de logit está dominado por los errores de aproximación angular ($\epsilon_\theta$) escalados por las magnitudes de la consulta y la clave.
• Restricciones de Protección: Para evitar fallos catastróficos en información crítica (por ejemplo, evidencia recuperada, intervalos de seguridad, valores atípicos), RDR impone restricciones de protección estrictas ($I_{prot}(i)=1 \implies z_i=1, b_i=b_{max}$).
• Páginas Homogéneas por Niveles: El controlador agrupa los estados retenidos en páginas donde todos los elementos comparten el mismo nivel de precisión. Esto permite lecturas de streaming de paso único, guiadas por encabezados, por parte del kernel de decodificación, evitando recolecciones irregulares.

2.3 Integración del Sistema

Spherical KV opera sobre un sustrato de KV irregular y paginado estándar (compatible con los diseños tipo vLLM). El flujo de trabajo mapea las claves densas a códigos esféricos, aplica RDR para filtrarlos y categorizarlos por niveles, y los almacena en un formato comprimido. El kernel de decodificación consume estos códigos directamente, manteniendo la invariante de "no reconstrucción densa".

3. Contribuciones Clave

1. Atención de Dominio Angular (ADA): Una primitiva a nivel de kernel que computa los logits de atención directamente desde códigos esféricos compactos, eliminando la penalización de tráfico de HBM asociada con la reconstrucción de claves densas.
2. Retención de Tasa–Distorsión (RDR): Un controlador unificado que asigna conjuntamente la residencia y los niveles de precisión bajo un presupuesto fijo, priorizando estados de alta utilidad y protegiendo aquellos que son frágiles.
3. Evaluación Orientada al Despliegue: El método se evalúa no en cachés contiguas de juguete, sino en un sustrato de servicio irregular y paginado realista, midiendo el rendimiento (throughput) de extremo a extremo, el tráfico de HBM y la estabilidad de comportamiento.
4. Desplazamiento de la Frontera de Pareto: El artículo demuestra que Spherical KV desplaza la frontera de memoria–calidad–rendimiento, logrando un mayor rendimiento con menores huellas de memoria en comparación con el KV denso y bases de referencia sólidas.

4. Resultados Experimentales

Los autores evaluaron Spherical KV en Llama-3.1-8B-Instruct, Qwen2.5-14B-Instruct y gpt-oss-20b a través de longitudes de contexto de 8K, 32K y 128K. Las cargas de trabajo incluyeron modelado de lenguaje (PG-19), QA con alta carga de recuperación (LongBench: HotpotQA, 2WikiMultiHopQA) y ejecuciones de agentes de herramientas (AgentBench, ToolBench).

• Eficiencia de Rendimiento y Memoria: Bajo restricciones de calidad igualadas (iso-quality), Spherical KV mejoró el rendimiento de decodificación en un 1.55×–1.72× mientras reducía los bytes de KV residentes por token en un 24–42%.
• Escalado de Contexto: Las ganancias fueron más pronunciadas en el régimen de 128K, donde la residencia de KV y el tráfico de HBM dominan el rendimiento del sistema.
• Eficiencia Realizada por Kernel: El método redujo con éxito los bytes de HBM por token generado, confirmando que los ahorros de memoria no se vieron compensados por los costos de reconstrucción.
• Estabilidad: El sistema demostró una degradación suave bajo presupuestos más estrictos y preservó las trayectorias de razonamiento y las secuencias de llamadas a herramientas en cargas de trabajo agénticas, evitando los fallos abruptos observados en las bases de referencia de evicción estricta.

5. Significado y Reivindicaciones

El artículo afirma que Spherical KV aborda el cuello de botella fundamental de la inferencia de contexto largo: el conflicto entre la capacidad de memoria y el ancho de banda de memoria. Al tratar la compresión de KV como un problema nativo de cómputo (donde la representación comprimida es directamente ejecutable por el kernel) en lugar de solo un problema de almacenamiento, el método logra un cambio genuino en la frontera operativa.

Los autores enfatizan que:

• La compresión debe ser nativa de cómputo: Almacenar simplemente menos bits es insuficiente si la ruta de decodificación requiere reconstrucción densa.
• La memoria debe presupuestarse donde importa: El control de tasa–distorsión permite al sistema gastar precisión en estados de alto impacto mientras ahorra en aquellos de baja utilidad.
• La evaluación debe ser fiel al despliegue: Las ganancias deben sobrevivir a diseños paginados realistas, sobrecargas de metadatos y restricciones de kernel para ser consideradas reales.

Spherical KV se presenta como un mecanismo para reducir la residencia de KV y el tráfico de HBM preservando la eficiencia de decodificación y la estabilidad del comportamiento, apuntando específicamente a los regímenes limitados por memoria de los servicios de LLM modernos.