Resumen Técnico: Caché de Prefijos Escasos para la Atención de LLM Híbridos y Recurrentes

1. Declaración del Problema

Los sistemas existentes de atención de Modelos de Lenguaje Grande (LLM) dependen en gran medida de la caché de prefijos densa, donde los estados de Clave/Valor (KV) para cada token en un prefijo compartido se almacenan para evitar el cálculo redundante de prellenado. Este enfoque es estándar para las arquitecturas Transformer. Sin embargo, los Modelos de Espacio de Estados (SSM) y las arquitecturas híbridas (que combinan atención con capas recurrentes) presentan una asimetría estructural: una capa recurrente evoluciona su estado como $h_t = F(h_{t-1}, x_t)$, lo que significa que puede reanudar el cálculo desde un único estado almacenado $h_c$ en lugar de requerir todo el historial de estados intermedios.

La caché densa por token de estados recurrentes a menudo es innecesaria y prohibitivamente costosa. Para modelos como Qwen-3.5, el estado recurrente por cabeza escala como $d_{head}^2$, mientras que el almacenamiento KV de atención escala linealmente. En consecuencia, almacenar el estado recurrente de cada token consume significativamente más memoria que las cachés KV de atención, sin ofrecer ningún beneficio adicional sobre almacenar un conjunto disperso de puntos de control.

El problema central abordado es: Dado un presupuesto de memoria de $M$ ranuras de puntos de control para un prefijo de longitud $N$, ¿dónde deben colocarse los puntos de control para minimizar el costo de recomputación esperado bajo una distribución específica de profundidades de superposición de solicitudes futuras? Esta optimización es distinta de las políticas de admisión/evicción de caché; se centra en qué tokens puntualizar dentro de una sola entrada retenida para maximizar la utilidad para superposiciones parciales de prefijos.

2. Metodología

Formulación Teórica

Los autores formalizan el problema como un problema de mediana k ponderada unidireccional sobre una línea.

• Sea $C = \{c_1, \dots, c_M\}$ el conjunto de posiciones de puntos de control.
• Para una solicitud que superpone el prefijo hasta la profundidad $t$, el punto de control reutilizable más profundo es $\ell(t; C) = \max(\{0\} \cup \{c \in C : c \le t\})$.
• El costo es el número de tokens a recomputar: $r(t; C) = t - \ell(t; C)$.
• El objetivo es minimizar el costo de recomputación esperado $E[r(T; C)]$ dada una distribución de profundidad de superposición $T$ con probabilidades $p_t$.

Enfoque Algorítmico

1. Colocación Equilibrada Óptima (Uniforme/Peor Caso): Los autores demuestran que, bajo una distribución de superposición uniforme o en el escenario de peor caso, el espaciado equilibrado (puntos de control distribuidos uniformemente) es óptimo.
2. Programación Dinámica (DP) Consciente de la Distribución: Para distribuciones de superposición no uniformes y del mundo real, los autores derivan un algoritmo exacto de programación dinámica $O(NM)$.
   • La relación de recurrencia es: $dp[m, j] = \min_{1 \le s \le j} (dp[m-1, s-1] + w(s, j))$, donde $w(s, j)$ es el costo de usar $s$ como el punto de control más profundo para las profundidades $s \dots j$.
   • Esto se resuelve eficientemente utilizando el truco del cascarón convexo monótono, aprovechando la monotonía de las sumas parciales.
3. Manejo de la Deriva de la Distribución: Dado que la verdadera distribución de superposición es desconocida y debe estimarse a partir de datos históricos, el artículo proporciona garantías teóricas (mediante el Lema 1 y el Teorema 3) de que la función de valor es Lipschitz continua con respecto a la distribución de superposición. Esto asegura que el uso de un histograma empírico (o un histograma ponderado exponencialmente para distribuciones con deriva) produce un programa casi óptimo con suboptimalidad acotada.

Estrategia de Implementación

• Puntos de Control: El sistema almacena estados recurrentes exactos en posiciones dispersas. Al producirse un acierto en la caché, el sistema reanuda desde el punto de control almacenado más profundo y recomputa el sufijo exactamente.
• Exactitud: Debido a que la actualización recurrente es determinista, reanudar desde un punto de control y reproducir el sufijo produce salidas idénticas a nivel de bits en comparación con un prellenado completo.
• Integración Híbrida: Para modelos híbridos, el método almacena la caché KV de atención completa (densa) pero aplica la puntuación dispersa solo a las capas recurrentes.

3. Contribuciones Clave

1. Optimalidad Teórica: Se demostró que el espaciado equilibrado es óptimo bajo condiciones de superposición uniforme y de peor caso (Teorema 1).
2. Algoritmo Exacto: Se formuló la caché de prefijos dispersos como un problema de mediana k ponderada y se proporcionó una solución exacta de programación dinámica $O(NM)$ (Teorema 2).
3. Garantías de Robustez: Se demostró que el programa óptimo es estable bajo la especificación incorrecta de la distribución. La suboptimalidad está acotada por la distancia de variación total entre las distribuciones de superposición empírica y verdadera (Teorema 3), y se extiende a escenarios con deriva de distribución utilizando el reponderado exponencial (Teorema 4).
4. Validación Empírica: Se validó el método en conjuntos de datos del mundo real (QuALITY, NarrativeQA, Prompts del Sistema), mostrando mejoras consistentes sobre las líneas base de presupuesto fijo.

4. Resultados Experimentales

Conjuntos de Datos y Configuración

• Conjuntos de Datos:
  • QuALITY: Documento único con múltiples preguntas (alta superposición cerca de la longitud completa).
  • NarrativeQA: Estructura similar con amplia superposición multimodal.
  • Prompts del Sistema: Prompts del sistema reales combinados con consultas de usuario (superposición concentrada en longitudes cortas).
• Líneas Base: Comparado contra "Sin Caché", "Solo KV", "Equilibrado" (espaciado uniforme), "Bloque" (tamaño de bloque fijo), "Logarítmico" y estrategias espaciadas en "$\sqrt{L}$".
• Hardware: Prototipo ejecutándose en NVIDIA RTX 2080 Super con el modelo Qwen-3.5-0.8B.

Hallazgos

• Ahorro de Tokens: El programa óptimo de DP domina consistentemente a las líneas base deterministas (Equilibrado, Bloque, Logarítmico) en términos de factores de reducción del trabajo recurrente.
• Ventaja de Bajo Presupuesto: La brecha de rendimiento es mayor en presupuestos de puntos de control bajos donde la distribución de superposición es más no uniforme. Por ejemplo, en Prompts del Sistema (donde las superposiciones son cortas), el método de DP coloca puntos de control densamente cerca del pico de divergencia, logrando el mismo tiempo de ejecución que la caché de bloques con menos puntos de control.
• Tiempo Real: En QuALITY y Prompts del Sistema, el programa derivado de DP produce aceleraciones medibles en tiempo real en grupos de capas representativos. La reducción del tiempo de ejecución se correlaciona linealmente con el número de tokens ahorrados.
• Sobrecarga: La sobrecarga de cargar puntos de control desde la CPU a la GPU es pequeña en relación con los ahorros obtenidos por la reducción de la recomputación, resultando en una mejora neta del tiempo de ejecución.

5. Significado y Afirmaciones

El artículo afirma que la caché de prefijos dispersos desbloquea un nuevo punto de diseño para la atención de LLM híbridos y recurrentes que anteriormente no estaba disponible.

• Compensación Memoria/Cálculo: Ofrece una compensación entre el uso de memoria y la recomputación que la caché KV densa no puede proporcionar, específicamente para arquitecturas donde el estado puede extraerse y restaurarse exactamente.
• Exactitud: A diferencia de las técnicas de compresión KV aproximadas, este método preserva salidas exactas sin alterar los núcleos de cálculo recurrente.
• Aplicabilidad: Es más beneficioso para cargas de trabajo donde las solicitudes comparten un prefijo sustancial pero no idéntico (por ejemplo, RAG con diferentes preguntas, prompts del sistema variables). No ofrece beneficios adicionales para cargas de trabajo de chat estrictamente de solo anexión (donde el último estado es suficiente), pero no incurre en sobrecarga en tales casos.
• Trabajo Futuro: Los autores reconocen desafíos de ingeniería, específicamente la necesidad de APIs eficientes de extracción/restauración de estado en tiempos de ejecución de producción y la integración de esta estrategia de colocación con políticas de admisión/evicción (por ejemplo, Marconi et al.). También señalan que la teoría actual de prefijo único podría extenderse a árboles de prefijos generales.