FlashAttention y PagedAttention: dos soluciones para dos cuellos de botella distintos

Quien pasa tiempo entre LLM oye hablar de FlashAttention y PagedAttention como si fueran rivales, y la pregunta «¿cuál debería usar?» aparece sin parar en las comunidades de ML. El planteamiento es erróneo. Resuelven dos problemas sin relación entre sí en dos etapas distintas de la vida de un modelo, y una pila de serving seria suele usar ambas a la vez.

Por qué la atención sale cara

La self-attention de un Transformer calcula esto:

Attention(Q, K, V) = softmax(QKᵀ / √d) · V

La fórmula parece inofensiva, pero esconde dos cuellos de botella distintos que dependen de cuándo la ejecutes.

1. Durante el cálculo (entrenamiento y prefill), la matriz de atención crece como O(N²) con la longitud de la secuencia N. Un documento de 10.000 tokens produce una matriz de 100 millones de celdas, que se calcula y se mueve por la memoria de la GPU una y otra vez.
2. Durante la generación (inferencia y decodificación), cada token nuevo obliga al modelo a recordar los vectores Key y Value de todos los tokens anteriores. Ese almacén es la caché KV, y crece con cada token y con cada petición adicional de usuario.

Un restaurante sirve de imagen. El primer problema es la eficiencia con la que el cocinero prepara un solo plato, es decir, el cálculo. El segundo es lo bien que el restaurante gestiona el almacén cuando hay cientos de pedidos en marcha, es decir, la gestión de memoria. FlashAttention acelera el cálculo y PagedAttention hace eficiente el uso de memoria durante el serving. Ten presente esa división, porque todo lo demás se deriva de ella.

FlashAttention: calcular la atención más rápido

El problema

Una GPU tiene dos tipos de memoria. La HBM (high bandwidth memory) es grande, decenas de GB, pero relativamente lenta. La SRAM, la memoria dentro del chip, solo guarda unos pocos MB, pero es rapidísima.

Una implementación estándar construye la matriz de atención completa de N por N y la escribe en la HBM, la vuelve a leer para que el softmax se ejecute en la SRAM, escribe otra vez el resultado y lo lee una vez más para la multiplicación final por V. Todo ese ir y venir entre memoria lenta y rápida es el verdadero cuello de botella. La aritmética en sí no lo es.

Imagina a un cocinero en una cocina pequeña y rápida (SRAM) con todos los ingredientes guardados en un almacén enorme (HBM) a un kilómetro y medio. La atención estándar es el cocinero conduciendo hasta allí, cogiendo un poco, volviendo, cocinando algo, volviendo a guardar las sobras y repitiéndolo en cada paso. El trayecto le hace perder más tiempo que la cocina. FlashAttention es un único viaje bien planificado: traer un lote del tamaño adecuado, hacer todos los pasos en la cocina y volver al almacén solo cuando el plato está terminado.

Cómo funciona

FlashAttention no cambia lo que calcula la atención, sino cómo lo hace. Las técnicas principales son:

• El tiling divide Q, K y V en bloques pequeños que caben en la SRAM.
• El softmax online evita tener que guardar una fila entera en memoria. Mantiene un máximo acumulado y una suma acumulada y actualiza ambos con cada bloque nuevo, de modo que la matriz completa de N por N nunca llega a materializarse.
• Los kernels fusionados ejecutan la multiplicación de matrices, el softmax y la segunda multiplicación dentro de un único kernel de GPU, sin escrituras intermedias en la HBM.
• El recálculo en el backward pass significa que la enorme matriz de atención no se guarda para los gradientes. Se reconstruye sobre la marcha durante la retropropagación, gastando un poco más de cómputo para ahorrar muchísima memoria.

El resultado es exactamente el mismo que el de la atención estándar. FlashAttention no es una aproximación, solo una forma más inteligente de llegar al mismo sitio.

En código

PyTorch usa por sí solo un kernel al estilo de FlashAttention en las GPU compatibles, como la A100 y la H100:

import torch
import torch.nn.functional as F

# PyTorch automatically dispatches to a FlashAttention-style
# kernel on supported GPUs (A100, H100, etc.)
q = torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.float16)  # (batch, heads, seq_len, head_dim)
k = torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.float16)
v = torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.float16)


output = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print(output.shape)  # torch.Size([2, 8, 1024, 64])

También puedes llamar directamente a la biblioteca dedicada flash-attn:

from flash_attn import flash_attn_func
import torch

q = torch.randn(2, 1024, 8, 64, device="cuda", dtype=torch.float16)
k = torch.randn(2, 1024, 8, 64, device="cuda", dtype=torch.float16)
v = torch.randn(2, 1024, 8, 64, device="cuda", dtype=torch.float16)

output = flash_attn_func(q, k, v, causal=True)

Cuándo encaja

• Entrenar Transformers grandes desde cero o hacerles fine-tuning.
• Ventanas de contexto largas de 8k, 32k o 128k tokens, donde el crecimiento N² duele de verdad.
• Cualquier caso, en entrenamiento o en inferencia, en el que importe la velocidad bruta de la atención.

PagedAttention: memoria eficiente para la caché KV

El problema

Este es un problema puramente de serving. En la decodificación autorregresiva, el modelo guarda en caché los tensores Key y Value de todos los tokens anteriores para no tener que recalcularlos, y eso es la caché KV. Los sistemas de serving tradicionales se topan aquí con problemas:

• La longitud de la salida no se conoce de antemano, así que reservan un gran bloque contiguo por petición, por ejemplo para max_tokens = 2048, aunque la respuesta acabe ocupando 50 tokens.
• Eso provoca una fuerte fragmentación interna. Hasta un 60-80 % de la caché KV reservada puede quedarse vacía y desperdiciada.
• No se comparte nada. Con beam search o muestreo en paralelo, donde varias salidas nacen de un mismo prompt, cada rama recibe una copia completa de la caché KV aunque todas empiecen con el mismo texto.

Piensa en un edificio de pisos (la memoria de la GPU) y sus inquilinos (las peticiones). El método antiguo obliga a cada inquilino nuevo a reservar una planta entera por adelantado, por si acaso. La mayoría usa dos o tres habitaciones de veinte, y nadie más puede tocar las vacías. PagedAttention reparte habitaciones sueltas, bloques o páginas de tamaño fijo, allí donde haya hueco y solo cuando hacen falta. Un directorio, la tabla de bloques, registra qué habitaciones son de quién. Los inquilinos emparentados, como las ramas de beam search con un prompt común, pueden incluso compartir habitaciones para las partes que tienen en común. La idea se toma directamente de la memoria virtual paginada que los sistemas operativos usan para los procesos, y de ahí el nombre.

Cómo funciona

1. La caché KV se divide en bloques de tamaño fijo, digamos de 16 tokens cada uno, en lugar de un único array contiguo gigantesco.
2. Cada secuencia mantiene una tabla de bloques que asigna índices lógicos de bloque a bloques físicos de memoria, como una tabla de páginas en un sistema operativo.
3. Los bloques se asignan bajo demanda, solo cuando una secuencia ha crecido lo suficiente como para necesitar uno, así que no hay reservas previas que desperdicien memoria.
4. El uso compartido funciona mediante copy-on-write. Con beam search, muestreo en paralelo o un prompt de sistema común, las secuencias apuntan a los mismos bloques físicos para su prefijo compartido y solo copian cuando su contenido se separa.
5. Un kernel de CUDA propio lee los bloques dispersos a través de la tabla de bloques. Como los bloques ya no son contiguos, ese kernel especializado es la verdadera operación de PagedAttention.

En código

Esto es lo que mueve vLLM, uno de los motores de serving de LLM más populares hoy:

from vllm import LLM, SamplingParams

# vLLM automatically manages KV cache using PagedAttention under the hood
llm = LLM(model="meta-llama/Llama-2-7b-hf")
prompts = [
    "The future of AI is",
    "PagedAttention helps LLM serving by"
]

sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=100)

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    print(output.outputs[0].text)

Como usuario nunca ves PagedAttention directamente. Trabaja en silencio en segundo plano para que la memoria de la GPU se aproveche bien en todas las peticiones simultáneas. Una visión simplificada y conceptual de lo que podría contener una tabla de bloques es esta, y no corresponde a las tripas reales de vLLM:

# Simplified conceptual example (not actual vLLM internals)
block_table = {
    "request_1": [block_3, block_7, block_12],
    "request_2": [block_3, block_7, block_9],
    # request_1 & request_2 share blocks 3 & 7 -> same prompt prefix!
}

Cuándo encaja

• Serving en producción con muchos usuarios simultáneos.
• APIs de chatbots, pipelines de RAG y cualquier sistema con longitudes de salida impredecibles.
• Beam search o muestreo en paralelo, donde compartir memoria ahorra mucho.
• Cualquier configuración en la que quieras sacar el máximo rendimiento de la GPU, porque más peticiones por GPU significa un coste menor por petición.

Usarlas juntas

Las dos se combinan, y en la práctica suele ser así. PagedAttention es la capa de gestión: decide dónde y cómo se guardan los bloques de la caché KV. FlashAttention, o los kernels fusionados de su estilo, es la capa de cálculo: decide lo rápido que se ejecutan el producto de Q por K transpuesta y el softmax. vLLM combina PagedAttention para la memoria con kernels de atención fusionados y optimizados, inspirados en FlashAttention, para las matemáticas, y así consigue a la vez eficiencia de memoria y velocidad de cálculo.

Siguiendo con la imagen de la cocina, PagedAttention elige de qué estante del almacén salen los ingredientes y FlashAttention se asegura de que el cocinero cocine con eficiencia cuando llegan.

¿Cuál necesitas?

• Entrenas un modelo o le haces fine-tuning con secuencias largas: FlashAttention.
• Sirves una API de LLM a cientos o miles de usuarios simultáneos: PagedAttention, mediante vLLM o algo parecido.
• Construyes una pila de serving de nivel de producción: las dos. La mayoría de los frameworks modernos, entre ellos vLLM, TensorRT-LLM y Hugging Face TGI, ya combinan ideas parecidas por dentro.

En resumen

La palabra «Attention» que comparten sus nombres es lo único que hace que se las meta en el mismo saco. Están en capas distintas de la pila. FlashAttention hace que el cálculo de la atención sea más rápido y gaste menos memoria al controlar cómo se mueven los datos entre la HBM y la SRAM. PagedAttention hace que el almacenamiento de la caché KV durante la inferencia sea eficiente y libre de fragmentación al tomar prestada de los sistemas operativos la memoria virtual paginada.

Una optimiza el cálculo y la otra la gestión de memoria para el serving. La mayoría de los sistemas de serving en producción actuales, vLLM incluido, usan las dos para tener modelos que van rápido y escalan de forma barata. Así que cuando alguien pregunte «¿PagedAttention o FlashAttention, cuál es mejor?», la respuesta honesta es que no es una cosa o la otra, sino las dos.