Notizie IA Logo

AITalk

Noticias y análisis sobre Inteligencia Artificial

PagedAttention y RadixAttention: hablemos de KV Cache

ResearchGenerative AIApplications

pagedattention-radixattention.jpg

En nuestro artículo anterior en AITalk habíais comenzado con una cifra incómoda: Llama-3.1-70B, en precisión BF16, acumula aproximadamente 0,31 megabytes de caché por cada token procesado, lo que significa que un contexto de 128.000 tokens cuesta ya 40 gigabytes de memoria GPU, antes de haber cargado siquiera una sola solicitud concurrente. Habíais leído sobre tres respuestas a este problema: TurboQuant, OSCAR y EpiCache, tres formas distintas de reducir el tamaño de esos datos, actuando sobre el número de bits por valor o sobre qué partes de la conversación vale la pena conservar en memoria.

Sin embargo, comprimir los datos solo resuelve la mitad del problema. Es como conseguir empaquetar los expedientes de un archivo en cajas más delgadas sin cambiar el modo en que el archivo está organizado: si las estanterías se siguen asignando en bloque a cada expediente independientemente de lo voluminoso que sea realmente, y si cada nuevo expediente idéntico al anterior se vuelve a copiar desde cero en lugar de recuperarse, la eficiencia ganada en el peso de cada carpeta se pierde igualmente por el camino. Este artículo aborda las dos técnicas que han afrontado precisamente este segundo problema: no cuánto pesa un token en la caché, sino cómo se asigna y cómo se reutiliza. Se llaman PagedAttention y RadixAttention, y no son alternativas a la compresión: operan en un nivel diferente y, en los sistemas de producción más maduros, conviven con ella.

Dos problemas independientes, no uno solo

Incluso disponiendo de una caché ya comprimida, un motor de inferencia se topa con dos ineficiencias que no tienen nada que ver la una con la otra. La primera atañe a la asignación de memoria. El método tradicional, el adoptado por los primeros sistemas de procesamiento, reservaba para cada solicitud entrante un bloque contiguo de memoria GPU dimensionado según la longitud máxima de contexto que el modelo podía alcanzar teóricamente. El problema radica en que el sistema no tiene forma de saber de antemano cuál será la longitud real de la respuesta, por lo que el margen de seguridad terminaba casi siempre estando sobredimensionado. De ello se derivan dos formas de desperdicio bien diferenciadas. Existe la fragmentación interna, cuando una solicitud reserva espacio para miles de tokens futuros pero solo genera unas pocas decenas, dejando sin usar la mayor parte de la memoria asignada. Y existe la fragmentación externa, cuando solicitudes de distinta longitud terminan en momentos diferentes dejando huecos dispersos en la memoria GPU: la suma del espacio libre puede ser abundante, pero ningún fragmento es lo suficientemente grande y contiguo como para albergar una nueva solicitud voluminosa. El resultado práctico, medido por los investigadores que abordaron el problema por primera vez, era que los sistemas de entonces desperdiciaban entre el 60 y el 80 por ciento de la memoria asignada, una cifra que vale la pena tener presente al evaluar la magnitud de los avances posteriores.

El segundo problema está conceptualmente alejado del primero y atañe al cálculo, no a la memoria. En las cargas de trabajo reales, las solicitudes rara vez son independientes entre sí. Miles de usuarios que consultan al mismo asistente comparten el mismo prompt de sistema, una conversación de varios turnos vuelve a presentar en cada intercambio todo el historial anterior, y un agente que razona en bucle continúa añadiendo nuevos pasos a un contexto que en su mayor parte permanece idéntico al del turno anterior. En todos estos casos, el motor de inferencia se ve obligado a recalcular desde cero, durante la fase de prefill, los mismos tensores de clave y valor exactamente que ya había producido instantes antes para otra solicitud, o para el mismo usuario en el turno anterior. Es un trabajo puramente redundante que, en escenarios de chat prolongado o en sistemas RAG, puede representar la parte más pesada del tiempo de respuesta percibido por el usuario.

Estas dos ineficiencias surgieron de necesidades distintas y han recibido soluciones diferentes. PagedAttention aborda la primera; RadixAttention aborda la segunda.

PagedAttention: la memoria como un sistema operativo

La idea fundamental de PagedAttention, introducida por el grupo creador de vLLM, toma prestado casi literalmente un concepto con décadas de antigüedad: la memoria virtual paginada de los sistemas operativos. Quien haya examinado alguna vez el funcionamiento interno de un ordenador sabe que los programas no obtienen un bloque físico contiguo de memoria RAM, sino la ilusión de contigüidad construida sobre páginas dispersas gestionadas mediante una tabla de traducción. PagedAttention aplica exactamente la misma lógica a la caché de atención.

El mecanismo se articula en cuatro pasos. En primer lugar, la caché de cada secuencia se divide en bloques lógicos de tamaño fijo, normalmente de dieciséis o treinta y dos tokens cada uno, en lugar de tratarse como un único contenedor monolítico. En segundo lugar, el sistema mantiene una tabla de páginas, un mapa que traduce cada bloque lógico a su ubicación física real dispersa en la memoria de la GPU: durante el cálculo de la atención, el motor consulta esta tabla para reunir los bloques necesarios, y para el modelo la secuencia aparece continua aunque físicamente no lo sea en absoluto. El tercer paso es el crecimiento bajo demanda: la memoria ya no se reserva en bloque al principio, sino que se asigna bloque a bloque a medida que avanza la generación, de modo que una respuesta de sesenta tokens ocupa solo el espacio de sesenta tokens, y no el previsto para un contexto máximo que tal vez nunca se alcance. El cuarto elemento, quizás el más elegante, es la compartición de bloques mediante copy-on-write: si varias solicitudes comienzan con el mismo prompt, como ocurre a diario con un prompt de sistema compartido por miles de usuarios, los bloques físicos correspondientes se referencian en común en lugar de duplicarse, y solo en el momento preciso en que una solicitud diverge de las demás se copia efectivamente el bloque afectado.

El resultado de esta arquitectura es que la fragmentación cae casi a cero, con márgenes de desperdicio que en los benchmarks publicados bajan del 4 por ciento en comparación con el 60-80 por ciento de los sistemas anteriores, lo que permite atender a un número mucho mayor de solicitudes concurrentes sobre el mismo hardware. Resulta interesante señalar que PagedAttention no altera en modo alguno el algoritmo de atención ni modifica las salidas del modelo: su innovación es puramente arquitectónica, orientada a dónde y cómo se ubican físicamente los datos, y no a lo que calcula el modelo. immagine1.jpg Imagen obtenida del artículo en arxiv.org

RadixAttention: un árbol que no olvida

Si PagedAttention resuelve el dónde, RadixAttention resuelve el qué. La idea, introducida por el equipo de SGLang, parte de una observación tan sencilla como desatendida: en los sistemas tradicionales, una vez completada una solicitud, su caché simplemente se descartaba, como si cada conversación naciera y muriera sin dejar rastro útil para las siguientes. RadixAttention invierte esta lógica transformando la caché en una estructura persistente y de búsqueda, un radix tree, es decir, un trie comprimido en el que cada prefijo único de tokens se almacena una sola vez, y diferentes solicitudes que comparten el mismo comienzo recorren la misma rama del árbol, bifurcándose únicamente en el punto exacto en que sus tokens empiezan a diferir.

Imaginad a tres usuarios que inician una conversación con el mismo mensaje de sistema, seguido cada uno de una pregunta diferente. En lugar de conservar tres copias casi idénticas del mismo prefijo, el árbol lo mantiene una sola vez y crea tres hojas separadas únicamente para la última parte, la verdaderamente específica de cada pregunta. Cuando llega una nueva solicitud, el motor recorre el árbol token a token buscando el prefijo más largo ya presente, recupera inmediatamente los tensores de clave y valor correspondientes sin recalcularlos y, finalmente, calcula solo el sufijo no visto, insertando la nueva ruta en el árbol para que pueda reutilizarse en el futuro. Un aspecto a menudo subestimado es que la estructura no puede crecer indefinidamente: cuando la memoria de la GPU se satura, RadixAttention aplica una política de desalojo basada en el uso menos reciente (least recently used) que elimina primero las hojas menos utilizadas recientemente, protegiendo al mismo tiempo los nodos internos compartidos de los que dependen varias solicitudes, a fin de preservar precisamente los prefijos que hacen rentable todo el mecanismo.

El beneficio principal no atañe al tamaño de la memoria en sí, sino al Time to First Token, es decir, cuánto tiempo transcurre antes de que el usuario vea aparecer la primera palabra de la respuesta: si mil novecientos de cada dos mil tokens están ya presentes en el árbol, el modelo solo necesita procesar los cien restantes, lo que repercute directamente en la latencia percibida, algo especialmente evidente en chatbots, sistemas RAG y flujos de trabajo de agentes donde los contextos crecen de forma incremental en lugar de reescribirse desde cero cada vez. Para quienes conozcáis la narrativa de videojuegos de nicho, este mecanismo recuerda a la estructura de 80 Days, el juego narrativo de Inkle en el que miles de rutas de viaje comparten tramos enteros de trama común y solo se separan en los puntos de elección real: el motor del juego no reescribe desde cero cada ramificación posible, sino que la construye injertándola en un tronco narrativo compartido. immagine2.jpg https://arxiv.org/html/2312.07104v1

Dos capas, no dos rivales

Un error común, alimentado en parte por la manera en que a menudo se presentan ambos proyectos en contraposición, es pensar en PagedAttention y RadixAttention como tecnologías competidoras entre las que hay que elegir. En realidad responden a preguntas distintas. PagedAttention decide dónde residen físicamente los bloques de caché en la memoria GPU; RadixAttention decide si esos bloques existen ya y pueden reutilizarse. El radix tree, en un sistema que integra ambas técnicas, apunta simplemente a bloques que a su vez están gestionados por el asignador paginado: son dos capas superpuestas de la misma pila, no dos caminos alternativos. immagine3.jpg

Volviendo al ejemplo de los tres usuarios con el mismo prompt de sistema, sin RadixAttention el motor recalcularía ese prefijo tres veces, y sin PagedAttention cada una de esas tres solicitudes reservaría igualmente una porción sobredimensionada de memoria. Combinando ambas técnicas, el prefijo se calcula una sola vez, se almacena en bloques paginados compactos y se reutiliza en cada solicitud que lo comparte.

Conviene precisar que RadixAttention no es la única vía hacia la reutilización automática de prefijos. También vLLM, el proyecto que introdujo PagedAttention, admite actualmente el prefix caching automático, aunque con una estructura de datos diferente: en lugar de un radix tree, utiliza un mecanismo de hash encadenado en el que cada bloque completado recibe un hash calculado a partir del hash del bloque padre, de los tokens contenidos en el propio bloque y de eventuales metadatos adicionales. Puesto que cada hash depende del anterior, toda la cadena representa de forma unívoca el prefijo que conduce a ese bloque y, si otra solicitud produce exactamente la misma secuencia de tokens, genera la misma cadena de hash, encontrando así de inmediato los bloques ya calculados. En la práctica, para la mayoría de las aplicaciones, la diferencia entre ambas implementaciones sigue siendo más arquitectónica que funcional: ambos motores llegan al mismo resultado, evitar el prefill de los prefijos ya vistos, aunque con estructuras de datos diferentes en la base.

Más allá de las dos técnicas: hacia dónde va la gestión de la caché

PagedAttention y RadixAttention abordaron los dos problemas fundacionales, pero los contextos continúan creciendo y, con ellos, han surgido necesidades para las que ninguna de las dos técnicas, por sí sola, estuvo concebida inicialmente.

La primera dirección es la memoria jerárquica. Tratar la GPU como el único nivel de caché disponible resulta insostenible cuando los contextos superan los cientos de miles de tokens: proyectos como LMCache y Mooncake organizan ya la caché en múltiples niveles, situando los bloques más activos en la memoria de alto ancho de banda de la GPU, los de uso reciente en la RAM del host y los más antiguos pero potencialmente útiles derivados a almacenamiento distribuido o discos NVMe, recuperándolos automáticamente cuando vuelven a ser necesarios. Para quienes hayáis leído a Zafón, es una arquitectura que recuerda de cerca al Cementerio de los Libros Olvidados: nada se destruye definitivamente, sino que los volúmenes menos consultados se desplazan a estanterías cada vez más remotas, listos para ser rescatados cuando alguien los necesite de nuevo.

La segunda dirección es el enrutamiento consciente de la caché. En una infraestructura distribuida en múltiples nodos, un balanceador de carga normal que distribuya las solicitudes en modo puramente round-robin corre el riesgo de enviar dos turnos de la misma conversación a dos GPU distintas, anulando cualquier ventaja del prefix caching aun cuando exista el prefijo compartido. Sistemas como el router de llm-d abordan este problema dirigiendo las solicitudes hacia la réplica que ya posee la caché del prefijo requerido, equilibrando no solo la carga, sino también la localidad de los datos.

La tercera dirección, menos discutida pero relevante para quienes operan infraestructuras multiinquilino (multi-tenant), se refiere a la seguridad. El prefix caching introduce un posible canal lateral: si un prefijo se sirve de forma visiblemente más rápida por estar ya presente en la caché, un usuario malintencionado que observe tiempos de respuesta inusualmente bajos podría deducir que otro usuario envió previamente el mismo prompt, sin que el contenido de la respuesta llegue a exponerse directamente. La contramedida más extendida es el cache salting, que incluye en el hash utilizado para identificar los bloques un elemento específico del inquilino, de modo que solicitudes idénticas procedentes de clientes distintos generen claves de caché diferentes, eliminando la posibilidad de un acierto cruzado entre inquilinos pero manteniendo intactos los beneficios de latencia dentro del mismo inquilino.

Compresión y gestión, dos caras de la misma moneda

Volviendo al punto de partida de nuestro artículo anterior, la visión completa de la gestión moderna de la KV cache solo emerge al unir ambos frentes. La compresión —de la que TurboQuant, OSCAR y EpiCache son expresiones recientes— reduce el tamaño del problema actuando sobre cuántos bits se necesitan para representar cada token y sobre qué partes del historial conversacional vale la pena conservar. La gestión —de la que PagedAttention y RadixAttention son los pilares fundamentales— optimiza su estructura, decidiendo dónde residen físicamente esos datos y si pueden reutilizarse en lugar de recalcularse. Un sistema de producción de vanguardia integra hoy en día ambos enfoques: comprime lo que debe permanecer en memoria y organiza de manera eficiente el espacio y el cálculo en torno a esa memoria comprimida.

Para quienes diseñéis o evaluéis una infraestructura de inferencia, esto significa que la elección ya no consiste únicamente en qué modelo servir, sino en cómo servirlo: dominar el funcionamiento de PagedAttention y RadixAttention resulta hoy tan crucial como conocer la arquitectura del modelo que esa caché está respaldando.


Nota técnica: los datos citados en este artículo proceden de los artículos originales de vLLM y SGLang, así como de la documentación de sus respectivos proyectos de código abierto; no han sido reproducidos de forma independiente por esta redacción y, como siempre en este campo, las cifras de laboratorio no garantizan automáticamente el mismo rendimiento en escenarios de producción con cargas de trabajo reales.