Optimización de KV Cache con PagedAttention y Cuantización FP8 para Contextos Largos (128K)

1. La Crisis de Memoria en la Era de los Contextos Extensos

Los modelos de lenguaje contemporáneos han expandido sus ventanas de contexto desde los clásicos 4.096 tokens hasta los 128.000 (Llama-3.1, Mistral Large) e incluso 1 o 2 millones de tokens (Gemini 1.5, Claude 3.5 Sonnet). Sin embargo, procesar libros enteros, bases de código completas o transcripciones masivas introduce un obstáculo técnico demoledor: el crecimiento implacable del Key-Value (KV) Cache.

2. Fórmula Matemática de Cálculo del KV Cache

Durante la atención autorregresiva, cada token debe almacenar sus vectores de proyección de Key ($K$) y Value ($V$) en cada una de las capas del modelo para no tener que recalcularlos en cada paso sucesivo. La ecuación analítica para estimar el tamaño del KV Cache en memoria es:

Bytes_por_Token = 2 * n_layers * n_kv_heads * d_head * precision_bytes

Veamos el cálculo exacto para Meta-Llama-3.1-70B:

• Número de capas ($n_{layers}$): 80
• Cabezas KV ($n_{kv_heads}$): 8 (utiliza Grouped-Query Attention con factor 8:1)
• Dimensión por cabeza ($d_{head}$): 128
• Precisión FP16 ($precision_bytes$): 2 bytes

Bytes_por_Token = 2 * 80 * 8 * 128 * 2 = 327.680 bytes ≈ 0.32 MB por token

Ahora multipliquemos este factor por el tamaño del contexto solicitado:

• Contexto 8k tokens: $8.192 imes 0.32 ext{ MB} approx 2.62 ext{ GB}$ por sesión.
• Contexto 32k tokens: $32.768 imes 0.32 ext{ MB} approx 10.48 ext{ GB}$ por sesión.
• Contexto 128k tokens: $131.072 imes 0.32 ext{ MB} approx 41.94 ext{ GB}$ por sesión.

Esto significa que en precisión estándar FP16, ¡una única petición con 128k tokens agota por completo la mitad de la memoria de una NVIDIA A100 de 80 GB antes de que un segundo usuario pueda siquiera conectarse!

3. PagedAttention: El Fin del Desperdicio de Memoria Virtual

El primer pilar para resolver esta crisis fue desarrollado por la Universidad de Berkeley con PagedAttention. Tradicionalmente, los servidores de inferencia reservaban bloques contiguos de memoria VRAM asumiendo el peor caso (la longitud máxima de secuencia configurada). Esto generaba dos problemas devastadores:

1. Fragmentación Interna: Si una solicitud reservaba 32k tokens pero terminaba generando 500 tokens, el 98% restante de la memoria reservada permanecía bloqueada e inutilizada.
2. Fragmentación Externa: Con peticiones de duraciones heterogéneas, la memoria se fragmentaba en huecos disjuntos demasiado pequeños para alojar nuevas secuencias completas.

PagedAttention resolvió esto tratando la VRAM como páginas virtuales de tamaño discreto (típicamente 16 tokens). Las páginas físicas no necesitan ser contiguas en el espacio de memoria; se asignan y vinculan dinámicamente mediante tablas de páginas. Esto elevó la utilización efectiva de la memoria en GPUs de un paupérrimo 20-30% a más del 96% sostenido.

4. Cuantización FP8 del KV Cache: Reducción del 50% de VRAM

Mientras que PagedAttention optimiza la gestión de la memoria, la cuantización del KV Cache ataca directamente el tamaño físico de los tensores almacenados. En lugar de almacenar cada valor en formato FP16 o BF16 de 16 bits, el sistema los convierte en tiempo real a representaciones de 8 bits (FP8).

5. Guía de Configuración en vLLM

Para activar la cuantización FP8 del KV Cache en vLLM, se utiliza el flag --kv-cache-dtype fp8. A continuación se muestra una configuración completa de producción para contextos largos:

vllm serve meta-llama/Meta-Llama-3.1-70B-Instruct \
  --tensor-parallel-size 4 \
  --kv-cache-dtype fp8 \
  --max-model-len 131072 \
  --enable-chunked-prefill \
  --max-num-batched-tokens 8192 \
  --gpu-memory-utilization 0.95 \
  --port 8000

Detalle de los parámetros avanzados para evitar colapsos:

• --enable-chunked-prefill: Impide que una petición con un documento enorme bloquee a todos los demás usuarios. Divide la fase de prefill en paquetes de tokens (por ejemplo, 8.192) y los intercala con los pasos de decodificación en curso.
• --max-num-batched-tokens 8192: Limita el número total de tokens procesados en una única pasada de GPU para evitar picos de latencia inter-token (ITL).

6. Diagnóstico y Gotchas de Memoria en Contextos Masivos

1. El error ValueError: No available memory for the cache blocks

Este fallo se manifiesta durante el arranque de vLLM cuando los pesos estáticos del modelo ocupan tanta VRAM que no queda espacio suficiente para asignar el espacio mínimo reservado para los bloques de PagedAttention. Por ejemplo, en una tarjeta de 24 GB intentando servir un modelo de 14B parámetros en FP16.

Solución: Reduce el tamaño máximo de contexto con --max-model-len 16384 o incrementa --gpu-memory-utilization a 0.95. Si el problema persiste, cuantiza los pesos del modelo a AWQ o FP8 antes de servirlo.

2. Desincronización de escalas dinámicas en FP8 E4M3

Si la GPU no implementa soporte nativo de hardware para FP8 (como ocurre en las tarjetas de arquitectura Ampere RTX 3090 o A100), el runtime de PyTorch emulará las operaciones FP8 mediante conversiones flotantes en software. Esto provoca que el rendimiento decaiga dramáticamente, volviéndose 3x más lento que FP16 nativo.

Solución: En GPUs de arquitectura Ampere (SM80), utiliza en su lugar --kv-cache-dtype int8 o mantén el KV Cache en FP16 y aplica cuantización en los pesos estáticos.

3. Latencia disparada por recomputación de prefijos (Prefill Bottleneck)

Cuando múltiples usuarios envían documentos de más de 32k tokens de forma simultánea, el tiempo necesario para procesar el prefijo inicial (Time-to-First-Token) puede dispararse a más de 10 segundos, degradando la experiencia de los usuarios existentes.

Solución: Habilita --enable-chunked-prefill con un tamaño de lote de prefijo controlado (e.g. --max-num-batched-tokens 4096). Esto fracciona el prompt en bloques concurrentes y asegura que los tokens de decodificación sigan fluyendo a tiempo constante.

7. Checklist y Monitorización en Producción

• Monitoriza constantemente la métrica de Prometheus vllm:gpu_cache_usage_factor. Si se mantiene por encima de 0.90, escala el clúster o reduce el max-model-len.
• Verifica que el test de aguja en el pajar (Needle-in-a-Haystack) mantenga un score del 100% en todas las profundidades de contexto tras activar FP8.

8. Calibración Estática vs Escalado Dinámico en Cuantización FP8

Al comprimir el KV Cache al formato FP8 (E4M3), existen dos metodologías para mapear los rangos dinámicos de los tensores flotantes:

1. Escalado Dinámico (Dynamic Scaling): La GPU calcula el valor máximo absoluto ($MaxAbs$) de cada bloque de atención en tiempo real durante la inferencia y aplica el factor de escala sobre la marcha. Ofrece la máxima precisión matemática sin necesidad de archivos auxiliares, pero añade una pequeña sobrecarga de cálculo en los núcleos CUDA.
2. Calibración Estática (Offline Calibration): Se ejecuta un proceso de muestreo previo con un dataset representativo para calcular los factores de escala fijos de cada capa, guardándolos en un archivo kv_cache_scales.json. Esta aproximación elimina todo el cálculo al vuelo, maximizando la velocidad pura en servidores de gran volumen.

9. Integración de FlashInfer: Aceleración de Atención Cuantizada

A partir de vLLM v0.6, el motor incorpora soporte para FlashInfer, una biblioteca de kernels CUDA de alto rendimiento especializada en atención heterogénea y KV Cache cuantizado. Para forzar el uso del backend FlashInfer en lugar de FlashAttention tradicional:

export VLLM_ATTENTION_BACKEND=FLASHINFER
vllm serve meta-llama/Meta-Llama-3.1-70B-Instruct-FP8 \
  --kv-cache-dtype fp8 \
  --max-model-len 32768

FlashInfer destaca especialmente en la fase de decodificación con lotes irregulares (ragged batching), ofreciendo entre un 10% y un 20% menor latencia por token que los kernels de referencia.