PagedAttention: gestionar la KV cache como memoria de un sistema operativo

El desperdicio que no ves en el nvidia-smi

Montas un servidor de inferencia, cargas un modelo de 13.000 millones de parámetros en una GPU con 40 GB y esperas atender un buen puñado de peticiones a la vez. En la práctica te caben cuatro. La GPU dice que está llena, pero si midieras qué hay dentro te llevarías una sorpresa: más de la mitad de esa memoria está reservada y vacía.

El culpable no es el modelo. Es cómo guardas la KV cache, esa tabla de claves y valores que crece token a token durante la generación y que cada petición arrastra consigo. Servir un LLM es, en buena medida, un problema de gestión de memoria. Y resulta que ese problema ya lo resolvieron los sistemas operativos hace cincuenta años.

Por qué la reserva ingenua tira la mitad de la VRAM

La forma obvia de guardar la KV cache es reservar un bloque contiguo de memoria por petición, del tamaño del contexto máximo que podría llegar a ocupar. Si tu modelo admite 4.096 tokens, reservas espacio para 4.096 aunque la respuesta se quede en 200.

De ahí salen tres fugas:

• Reserva por si acaso. Apartas memoria para la longitud máxima porque no sabes cuánto generará el modelo. La mayoría de respuestas ni se acercan, así que ese hueco queda muerto durante toda la petición.
• Fragmentación interna. El último trozo del bloque casi nunca se llena justo hasta el borde. Cada petición deja una esquina inutilizable.
• Fragmentación externa. Al reservar y liberar bloques contiguos de tamaños distintos, la memoria se llena de huecos que no encajan con la siguiente petición. Tienes VRAM libre, pero repartida en pedazos inservibles.

El resultado medido en el paper original de vLLM: los sistemas anteriores desperdiciaban entre el 60 % y el 80 % de la memoria dedicada a la KV cache. Estás pagando una GPU entera para usar un tercio.

La idea prestada: paginación

Un sistema operativo se enfrenta al mismo dilema con la RAM y lo resolvió con memoria virtual paginada. En vez de dar a cada proceso un bloque contiguo, parte la memoria en páginas de tamaño fijo y mantiene una tabla que traduce las direcciones que ve el proceso a las páginas físicas reales, que pueden estar desperdigadas por donde sea.

PagedAttention aplica exactamente ese truco a la KV cache. Divide la caché de cada secuencia en bloques de tamaño fijo, cada uno con espacio para un número pequeño y constante de tokens (por ejemplo 16). Los bloques no tienen por qué ser contiguos en la VRAM. Una tabla de bloques por secuencia lleva la cuenta de qué bloque físico corresponde a cada tramo lógico de la conversación.

La secuencia cree que su KV cache es un array continuo. Por debajo, vive troceada y repartida.

Cómo funciona por dentro

El bloque es la pieza clave. Cada uno guarda las claves y los valores de unos pocos tokens consecutivos. Cuando una secuencia arranca, no se le reserva nada más allá de lo que necesita para su prompt. Conforme el decode va generando tokens y se llena el bloque actual, el gestor le entrega el siguiente bloque libre del pool, esté donde esté físicamente.

El kernel de atención se adapta para leer esa memoria no contigua. En lugar de recorrer un array lineal, consulta la tabla de bloques para saber dónde vive cada tramo antes de calcular la atención sobre él:

# Esquema mental de PagedAttention (no es la implementación real,
# que vive en un kernel CUDA por rendimiento).

BLOCK_SIZE = 16  # tokens por bloque

class SequenceCache:
    def __init__(self):
        # La tabla de bloques: índice lógico -> id de bloque físico.
        # La secuencia ve un contexto continuo; por debajo está troceado.
        self.block_table = []

    def append_token(self, allocator, k, v):
        # ¿El último bloque está lleno? Pide uno nuevo, no antes.
        if not self.block_table or allocator.is_full(self.block_table[-1]):
            self.block_table.append(allocator.alloc())  # bloque de cualquier hueco libre
        allocator.write(self.block_table[-1], k, v)

def attention(query, seq: SequenceCache, allocator):
    # Recorre los bloques físicos siguiendo la tabla, no un array lineal.
    scores = []
    for block_id in seq.block_table:
        k_block, v_block = allocator.read(block_id)
        scores.append(dot(query, k_block))  # atención sobre este tramo
    return softmax_weighted_sum(scores, seq, allocator)

Como todos los bloques miden lo mismo, cualquier hueco libre sirve para cualquier secuencia. La fragmentación externa desaparece por construcción. Y el único desperdicio interno posible es el del último bloque a medio llenar: unos pocos tokens por secuencia, no miles.

El regalo extra: compartir memoria

Trocear la caché abre una puerta que la reserva contigua tenía cerrada a cal y canto: compartir bloques entre secuencias.

Piensa en el muestreo paralelo, cuando pides varias respuestas al mismo prompt. Con el enfoque ingenuo, duplicas la KV cache del prompt tantas veces como respuestas quieras. Con PagedAttention, todas las secuencias apuntan a los mismos bloques físicos del prompt compartido. Solo cuando una empieza a divergir del resto se copia el bloque afectado: es el mismo copy-on-write que usa tu sistema operativo al hacer fork.

Lo mismo aplica a la búsqueda por haces (beam search) o a un prefijo de sistema común a miles de peticiones. El ahorro no es cosmético: en escenarios de muestreo paralelo, compartir prefijos recorta la memoria de la KV cache de forma notable, y esa memoria liberada se traduce directamente en más peticiones simultáneas.

Qué significa esto para lo que pagas

La cadena es corta y va toda en la misma dirección. Menos desperdicio de memoria significa más peticiones dentro de la GPU a la vez. Más peticiones simultáneas significan lotes más grandes en el decode, que es justo donde el batching continuo reparte el coste de leer los pesos entre muchos usuarios. Más throughput por GPU significa menos GPUs para el mismo tráfico.

No es teoría de laboratorio: es la razón por la que vLLM se comió el mundo de la inferencia y por la que la técnica acabó copiada en casi todos los motores serios que vinieron después.

Un par de cosas que conviene tener claras antes de tocar nada:

• No acelera un token suelto. PagedAttention no baja la latencia de una única petición aislada; el decode sigue siendo secuencial y limitado por memoria. Lo que mejora es cuánta gente cabe a la vez, o sea el throughput, no el TPOT de un usuario solo.
• El tamaño de bloque es un compromiso. Bloques pequeños afinan el desperdicio interno pero engordan la tabla y añaden trabajo de gestión. Bloques grandes hacen lo contrario. Los valores por defecto del motor suelen ser razonables; no los toques sin medir.

En resumen

Servir un LLM con eficiencia es, más que nada, gestionar bien la KV cache, y la reserva contigua ingenua desperdicia la mayor parte de la VRAM en huecos vacíos y fragmentación. PagedAttention roba la solución a los sistemas operativos: trocea la caché en bloques de tamaño fijo, los reparte bajo demanda por donde haya sitio y los mapea con una tabla de bloques. Adiós a la fragmentación externa, y de paso ganas poder compartir prefijos entre secuencias casi gratis.

Si estás dimensionando GPUs para inferencia y las cuentas no salen, antes de firmar la factura mira cuánta de tu VRAM está reservada y vacía. Probablemente el problema no sea que te falte hardware, sino que lo estás repartiendo como si la paginación no se hubiera inventado.

Fuentes

• Efficient Memory Management for Large Language Model Serving with PagedAttention — el paper original de vLLM (Kwon et al., 2023).