Cómo PagedAttention y RadixAttention optimizan el manejo del KV cache en LLMs

El problema real detrás del rendimiento en producción

Los grandes modelos de lenguaje (LLMs) generan texto token por token. Para cada nuevo token, el modelo debe atender a todas las representaciones previas mediante las matrices de keys (K) y values (V). Recalcular K y V en cada paso sería impráctico, por eso los sistemas de serving mantienen un KV cache en memoria que evita trabajo redundante.

Sin embargo, ese mismo KV cache se convierte en el factor que más limita la capacidad de servicio: a medida que crece la ventana de contexto, la memoria requerida aumenta linealmente con la longitud de la secuencia. En modelos de contexto largo, el KV cache suele ser el mayor consumidor dinámico de memoria GPU y define cuántas solicitudes pueden procesarse en paralelo, la latencia y el throughput.

Un ejemplo ilustrativo del artículo original: en un modelo Llama-3 clase 8B con 32 capas, 8 cabezas KV, dimensión de cabeza de 128 y precisión FP16, cada token ocupa aproximadamente 128 KiB de KV cache. Un contexto de 100.000 tokens requiere casi 12.8 GiB solo para el KV cache, antes de considerar batching o múltiples solicitudes concurrentes.

Dos cuellos de botella fundamentales

Al llenar la memoria GPU con tensores KV, los sistemas de serving enfrentan dos problemas distintos:

• Fragmentación de memoria: asignaciones ineficientes dejan espacios inútiles en la GPU, reduciendo la concurrencia.
• Computación redundante: prefijos idénticos de prompts se vuelven a calcular y almacenar una y otra vez, aunque su estado KV ya existe.

Cada problema requiere una solución enfocada: PagedAttention para la asignación de memoria y RadixAttention para la reutilización de prefijos. Juntas, estas técnicas forman la base del serving de LLMs de alto rendimiento.

PagedAttention: resolver la sobreasignación y la fragmentación

Durante mucho tiempo, los motores de serving reservaron un gran bloque contiguo de memoria GPU para el KV cache de cada solicitud, a menudo cercano al máximo largo de contexto del modelo. Como no se podía anticipar cuántos tokens produciría una respuesta, la mayor parte del espacio quedaba sin usar, reduciendo drásticamente el número de secuencias servibles al mismo tiempo.

La asignación contigua tradicional genera dos tipos de fragmentación:

• Fragmentación interna: una petición reserva miles de ranuras de token pero solo genera unas pocas, dejando memoria sin usar dentro de la reserva.
• Fragmentación externa: al finalizar solicitudes de distintas longitudes quedan huecos dispersos en la memoria que ya no forman un bloque contiguo utilizable.

PagedAttention propone una idea simple y poderosa: asignar memoria KV solo cuando se necesita, dividiendo la secuencia en bloques lógicos de tamaño fijo (por ejemplo, 16 o 32 tokens) en lugar de reservar un buffer contiguo para toda la posible secuencia.

Cómo funciona PagedAttention

1. División en bloques: cada secuencia se fragmenta en bloques lógicos de tamaño fijo. Estos bloques pueden almacenarse físicamente en cualquier lugar de la GPU.

2. Tabla de bloques para traducción de direcciones: cada petición mantiene una tabla que mapea IDs de bloques lógicos a sus ubicaciones físicas. Durante la atención, el kernel consulta esta tabla para reunir las keys y values necesarias, haciendo que la secuencia parezca continua incluso si sus datos están físicamente dispersos.

3. Crecimiento on‑demand: conforme avanza la generación, se asignan nuevos bloques solo cuando el bloque anterior se llena. Una petición que genere 60 tokens solo ocupará los bloques necesarios para esos 60 tokens, evitando reservas para tokens que quizás nunca se produzcan.

Una ventaja adicional es la posibilidad de compartir bloques: si varias solicitudes comienzan con el mismo prompt, pueden referenciar los mismos bloques físicos KV en lugar de almacenar tensores duplicados. Cuando las solicitudes divergen, el sistema puede manejar la separación de estados en el punto necesario sin almacenar duplicados desde el inicio.

RadixAttention: atacar la recomputación por prefijos idénticos

El segundo problema crítico es la recomputación. En escenarios de chat o APIs que manejan plantillas y prompts similares, es común que múltiples interacciones compartan prefijos idénticos. Sin una estrategia de reutilización, esos prefijos se vuelven a codificar y a almacenar en caché repetidamente, consumiendo CPU/GPU y memoria innecesariamente.

RadixAttention fue propuesta como la técnica que permite una reutilización eficiente de prefijos: identificar cuando una nueva solicitud contiene un prefijo cuyo estado KV ya está disponible y reutilizar ese estado en lugar de volver a computarlo.

Cómo aborda RadixAttention el problema

El enfoque consiste en mantener una estructura de referencia para estados KV asociados a prefijos de prompts y en proporcionar búsquedas rápidas para detectar coincidencias parciales o completas. Al reconocer un prefijo ya procesado, el motor de serving puede enlazar la nueva solicitud con el KV cache existente y continuar la generación desde ese punto, evitando el prefetcheo y recálculo de K y V.

Esta optimización reduce la carga computacional y la latencia inicial en escenarios con alto grado de repetición de prompts, algo habitual en aplicaciones empresariales, asistentes virtuales y flujos de trabajo basados en plantillas.

PagedAttention vs. RadixAttention: diferencias y complementariedad

Es clave entender que PagedAttention y RadixAttention no compiten: atacan cuellos de botella distintos y se complementan:

• PagedAttention optimiza cómo y cuándo se reserva memoria KV, reduciendo la fragmentación interna y externa y permitiendo un crecimiento incremental y compartido de la caché.
• RadixAttention optimiza el trabajo computacional evitando recomputar estados KV para prefijos ya existentes, lo que reduce uso de GPU y mejora latencias.

Combinadas, estas técnicas permiten servir modelos con ventanas de contexto mucho más grandes y con mejor aprovechamiento de memoria y cómputo. En la práctica, eso se traduce en mayor concurrencia, menor coste por petición y respuestas más rápidas.

Relevancia para América Latina

En la región, donde los presupuestos de infraestructura y la disponibilidad de GPUs de alta memoria pueden ser limitados, las eficiencias que introducen PagedAttention y RadixAttention son especialmente valiosas. Permiten a empresas y equipos de investigación obtener más rendimiento por GPU, reducir costos de nube y escalar servicios conversacionales con menos inversión en hardware.

Además, muchas aplicaciones latinoamericanas (soporte al cliente, automatización de procesos, asistentes en español) se benefician del hecho de que los prompts y turnos de conversación tienden a repetir patrones y plantillas, lo que hace que la reutilización de prefijos sea una ganancia directa en eficiencia.

Conclusión

Mientras que técnicas como cuantización, pruning o kernels de atención más rápidos son importantes, en producción el verdadero cuello de botella suele ser cómo se gestiona el KV cache. PagedAttention y RadixAttention abordan, respectivamente, la asignación eficiente de memoria y la reutilización de prefijos, y su combinación es clave para un serving de LLMs más rápido y económico.

Para equipos en América Latina que despliegan modelos conversacionales o de contexto largo, adoptar estas ideas puede significar mayor concurrencia, menor latencia y un uso más racional del presupuesto en infraestructura.

Si están evaluando optimizaciones para producción, enfoquen primero en cómo se almacena y comparte el KV cache: ahí suelen estar los mayores retornos por esfuerzo.