FastKV: Decoupling of Context Reduction and KV Cache Compression for Prefill-Decoding Acceleration

1. El Problema

Los Modelos de Lenguaje Grande (LLM) con ventanas de contexto extendidas (128K tokens o más) enfrentan dos cuellos de botella críticos durante la inferencia:

• Fase de Prefill (Relleno): El cálculo de atención escala cuadráticamente con la longitud de la entrada, lo que hace que procesar prompts largos sea extremadamente costoso en tiempo de cómputo.
• Fase de Decodificación: La generación de tokens requiere acceder repetidamente a una gran caché de claves y valores (KV Cache). A medida que crece el contexto, el tamaño de esta caché consume mucha memoria GPU y reduce el ancho de banda, degradando la eficiencia.

Limitaciones de los enfoques existentes:

• Métodos enfocados solo en decodificación (ej. StreamingLLM, SnapKV): Reducen el tamaño de la caché KV durante la generación, pero no aceleran la fase de prefill, ya que aún deben procesar todo el contexto inicial.
• Métodos enfocados en prefill (ej. GemFilter, PyramidInfer): Intentan reducir el contexto desde el inicio. Sin embargo, acoplan (vinculan) la reducción de cómputo en prefill con el presupuesto de la caché KV para decodificación. Esto significa que para acelerar el prefill, deben descartar tokens prematuramente, lo que a menudo degrada la precisión porque las capas tempranas pierden tokens críticos que las capas posteriores necesitarían.

2. Metodología: FastKV

FastKV introduce un marco de compresión de caché KV que desacopla la reducción del contexto en el prefill de la retención de la caché KV en la decodificación. Se basa en dos observaciones clave:

1. Dinámica dependiente de la capa: En las primeras capas, el foco de atención es inestable y los tokens críticos cambian rápidamente. En las capas posteriores, el conjunto de tokens importantes se estabiliza y se superpone significativamente entre capas.
2. Uso disperso del contexto: Durante la decodificación, el modelo solo atiende a una pequeña fracción de los tokens prefilled, por lo que no es necesario retener toda la caché KV generada.

Componentes Principales:

A. Propagación Selectiva de Tokens (Token-Selective Propagation - TSP)
FastKV utiliza una estrategia de prefill en dos etapas:

• Etapa 1 (Capas Tempranas): Se procesa el contexto completo (todos los tokens) hasta llegar a una capa dedicada llamada Capa TSP. Esto permite que las capas tempranas capturen dependencias complejas y dinámicas sin perder información.
• Etapa 2 (Capas Tardías): En la Capa TSP, se evalúa la importancia de los tokens basándose en la atención de una ventana de tokens recientes. Solo se propagan hacia las capas siguientes un subconjunto seleccionado de tokens (los más informativos) junto con los tokens de la ventana. Esto reduce la carga computacional para las capas restantes.

B. Retención de KV Desacoplada (KV Retention Decoupled)
FastKV introduce dos hiperparámetros independientes:

1. Tasa TSP: Controla cuántos tokens se propagan después de la capa TSP (reducción de cómputo en prefill).
2. Tasa de Retención KV: Controla cuántas entradas de la caché KV se conservan en cada capa para su uso en la decodificación.

Innovación Clave: A diferencia de métodos anteriores, FastKV permite que la reducción de tokens en el prefill sea independiente de la cantidad de caché KV guardada. Se puede mantener una alta precisión en el prefill (propagando muchos tokens) mientras se comprime agresivamente la caché KV para la decodificación, o viceversa.

3. Contribuciones Clave

• Desacoplamiento: Rompe la relación rígida entre la reducción de cómputo en prefill y el presupuesto de memoria en decodificación, permitiendo una optimización flexible de la precisión y la eficiencia.
• Estrategia de Dos Etapas: Propone procesar el contexto completo en las primeras capas (donde la atención es inestable) y luego aplicar selección de tokens solo cuando la importancia de los tokens se estabiliza en capas posteriores.
• Arquitectura sin sobrecarga: El método no modifica la arquitectura del modelo ni introduce sobrecarga en tiempo de ejecución durante la decodificación; la compresión ocurre durante el prefill.

4. Resultados Experimentales

Los experimentos se realizaron en modelos como LLaMA-3.1-8B-Instruct y Ministral-8B-Instruct con contextos de hasta 128K tokens, utilizando benchmarks como LongBench, RULER y Needle-in-a-Haystack.

• Aceleración:
  • Prefill: Hasta 1.82× más rápido que la línea base de contexto completo.
  • Decodificación: Hasta 2.87× más rápido que la línea base.
• Precisión:
  • FastKV mantiene la precisión casi idéntica a la del contexto completo (caída < 1% en LongBench), superando significativamente a métodos como GemFilter y PyramidInfer que sufren grandes caídas de precisión al comprimir agresivamente.
  • En la prueba "Needle-in-a-Haystack", FastKV obtuvo una puntuación de 99.9 (frente a 99.0 del contexto completo y 95.8 de GemFilter), demostrando una excelente capacidad de recuperación de información.
• Comparación: Logra un equilibrio superior entre velocidad y precisión, siendo la única metodología que acelera simultáneamente el prefill y la decodificación sin sacrificar la calidad del modelo.

5. Significado e Impacto

FastKV representa un avance significativo en la inferencia eficiente de LLMs de largo contexto. Al demostrar que la reducción de tokens y la compresión de caché deben tratarse como problemas separados y gestionarse en diferentes momentos del proceso de inferencia, ofrece una solución práctica para escalar modelos a contextos de millones de tokens.

Su diseño permite:

• Inferencia escalable: Hace viable el uso de ventanas de contexto masivas en hardware limitado.
• Flexibilidad: Los usuarios pueden ajustar las tasas de TSP y retención KV según si priorizan la latencia de inicio (prefill) o el uso de memoria durante la generación.
• Compatibilidad: Es ortogonal a otras técnicas como la cuantización o la atención dispersa, lo que sugiere que puede combinarse con ellas para lograr aún más eficiencia.

En resumen, FastKV resuelve el compromiso tradicional entre velocidad y precisión en el procesamiento de contextos largos, estableciendo un nuevo estándar para la compresión de caché KV en LLMs.