El especialista del plato estrella: el backend de atención de vLLM (FlashAttention, FlashInfer y la asimetría prefill/decode)

Mon, 08 Jun 2026 05:40:00 +0200

Sigue la serie por debajo del motor. El post de PagedAttention explicó dónde vive el KV (en bloques paginados). Este explica quién lo lee y cómo: el kernel de atención. Y conecta con FlashAttention v1-v4, que desmontó cómo es ese kernel por dentro; aquí miramos el nivel de arriba —cómo vLLM elige entre varios kernels y por qué necesita más de uno—.

TL;DR

Un forward de un LLM es, en su mayor parte, multiplicaciones de matrices estándar que cualquier librería hace bien. La excepción que decide el rendimiento es la atención, y no basta con tener un kernel bueno: hacen falta dos, porque las dos fases de la inferencia son problemas físicamente opuestos. El prefill procesa el prompt entero: muchas queries contra muchas keys, denso y compute-bound —el terreno del tiling IO-aware de FlashAttention—. El decode genera un token: una sola query contra todo el KV acumulado, flaco y memory-bound —aquí lo único que importa es saturar el ancho de banda de la HBM leyendo el KV paginado—. Por eso vLLM no tiene “el kernel de atención” sino un backend conmutable (FLASH_ATTN, FLASHINFER, TRITON_ATTN…) y una lógica que elige según la GPU: por defecto FA4 en Blackwell (SM100), FA3 en Hopper (SM90), FA2 en lo demás, con FlashInfer como alternativa que compila kernels a medida (JIT) y sabe hacer cascade attention para prefijos compartidos. Este post explica por qué prefill y decode son opuestos (con la intensidad aritmética), cómo el backend lee KV paginado, cómo el motor elige, qué aporta FlashInfer, los 10 knobs y la trampa de fijar un backend a ciegas. Sobre el cluster genérico 4×H100 SXM.

Dónde estás: el especialista, no el pinche

En la cocina, casi todo el trabajo es picar y saltear: operaciones estándar que cualquier pinche competente ejecuta —son las multiplicaciones de matrices de las capas feed-forward y las proyecciones—. Hay un solo plato que no se delega: el plato estrella, el que define al restaurante. Ese plato es la atención, y tiene una particularidad: se cocina de dos maneras radicalmente distintas según el momento del servicio.

Durante el prefill —cuando llega una comanda nueva con su prompt entero— hay que cocinar a lo grande: mucha materia prima de golpe, mucho fuego, una operación intensa que llena los fogones. Durante el decode —cuando una mesa pide “un plato más”— hay que cocinar a la carta: un solo plato, pero hay que ir a la despensa y traer todos los ingredientes que esa mesa ha acumulado durante toda su comida. Uno es un problema de potencia de fuego; el otro, de velocidad de la despensa. No los hace bien el mismo especialista. Por eso vLLM tiene varios, y un jefe que decide cuál entra según la GPU y la fase. Eso es el backend de atención.

Por qué prefill y decode son problemas opuestos

Esta es la idea central, y se demuestra con una sola cuenta: la intensidad aritmética (FLOPs por byte leído). Una operación con intensidad alta está limitada por el cómputo; una con intensidad baja, por la memoria.

Prefill. Atendemos $N$ queries (todo el prompt) contra $N$ keys. La operación $QK^\top$ y la $\text{softmax}\cdot V$ hacen del orden de $N^2 d$ FLOPs y leen del orden de $N d$ datos. La intensidad crece con $N$:

$$I_\text{prefill} \sim \frac{N^2 d}{N d} = N$$

Con $N$ grande (un prompt de miles de tokens), la intensidad es alta: compute-bound. Es donde el tiling de FlashAttention exprime los tensor cores y donde se acerca a los TFLOPS de pico de la GPU.

Decode. Atendemos una query (el token nuevo) contra $L$ keys (todo el KV acumulado). FLOPs del orden de $L d$; bytes leídos del orden de $L d s$ (hay que leer el KV entero de la HBM). La intensidad es:

$$I_\text{decode} \sim \frac{L d}{L d s} = \frac{1}{s} \quad (\approx 0,5 \text{ FLOP/byte en FP16})$$

Constante y diminuta: memory-bound. El kernel de decode no está limitado por cuánto puede calcular la GPU sino por cuán rápido lee el KV de la HBM. Da igual que la H100 tenga 132 SMs ociosos (ver el post de SMs): el cuello es el ancho de banda de 3,35 TB/s, y el kernel de decode existe para no desperdiciar ni uno de esos bytes/s.

La consecuencia de diseño: un kernel optimizado para prefill (tiling denso, máxima ocupación de tensor cores) no es el óptimo para decode (lecturas coalescidas del KV paginado, latencia mínima). Los servidores serios tienen kernels distintos —o un kernel con dos caminos—. En los modelos con MLA (atención latente multi-cabeza), vLLM llega a usar backends separados para prefill y decode, seleccionables de forma independiente (attention backends, vLLM).

El truco del scheduler: prefill y decode en el mismo forward

Aquí cierra el círculo con el post del scheduler. Como vLLM V1 mezcla en cada step peticiones en prefill y en decode, un mismo forward tiene que atender las dos cosas. El backend recibe metadatos que le dicen, para cada secuencia del batch, cuántas queries trae y cuánto KV tiene que leer, y aplica el camino que toca a cada una. Por eso el backend de atención y el scheduler están acoplados: el primero tiene que digerir el batch heterogéneo que el segundo arma.

Cómo lee el backend el KV paginado

El kernel no recibe un tensor de KV contiguo: recibe la block table del block manager y hace un gather sobre los bloques físicos. Esto impone una restricción real al backend: tiene que soportar el layout paginado y el block_size de vLLM. No todos los kernels del mundo lo hacen; los que vLLM integra (FlashAttention, FlashInfer, Triton) están adaptados a leer KV en bloques de tamaño fijo dispersos por la HBM. Es la razón de que no puedas enchufar cualquier kernel de atención de un paper: tiene que hablar el idioma de la despensa por casilleros.

Los backends y cómo elige el motor

vLLM expone una abstracción de backend con varias implementaciones (deepwiki vLLM):

FLASH_ATTN — la familia FlashAttention. Por defecto se elige la versión según la arquitectura: FA4 en SM100 (Blackwell), FA3 en SM90 (Hopper), FA2 en el resto, configurable con flash_attn_version.
FLASHINFER — motor de atención con compilación JIT y kernels especializables; fuerte en KV heterogéneo y prefijos compartidos.
TRITON_ATTN — escrito en Triton, portable y sin depender de binarios CUDA precompilados (Triton backend deep dive, vLLM, mar-2026).
Backends específicos para MLA y para hardware no-NVIDIA.

La selección es automática salvo que la fuerces con VLLM_ATTENTION_BACKEND. La heurística prueba FlashAttention primero; en Blackwell (SM100) el orden de respaldo para MLA es TRT-LLM Ragged → FlashInfer → otros; en otras GPUs solo se considera FlashAttention para el camino principal (attention backends, vLLM). La decisión depende de: arquitectura (SM), dtype (FP16/BF16/FP8), dimensión de cabeza, y si la carga necesita una feature que solo un backend tiene (cascade attention, ciertos soft caps, FP8 en KV).

Qué aporta FlashInfer: JIT y cascade attention

FlashInfer no compite con FlashAttention en “ser un poco más rápido”; ataca un problema distinto: la heterogeneidad del KV en servicio real (FlashInfer, arXiv 2501.01005). Dos ideas:

Compilación JIT. En lugar de un kernel monolítico, FlashInfer genera kernels a medida para la variante de atención, la forma del problema y el layout del KV que tengas, inyectando functors (transformaciones de query/key/logits, máscaras). Especializa en vez de generalizar.

Cascade attention. Aquí está la joya para servicio con prefijos compartidos. Si $R$ peticiones comparten un prefijo de $P$ tokens (un system prompt común), la atención ingenua leería ese prefijo $R$ veces. La cascade attention lo calcula una vez contra el prefijo compartido y luego combina con el sufijo propio de cada petición:

$$\text{lecturas: } \underbrace{R \cdot (P + s_i)}{\text{ingenua}} ;\longrightarrow; \underbrace{P + \textstyle\sum_i s_i}{\text{cascade}}$$

Con $R=50$ peticiones y un prefijo $P=1000$, eso es leer 50.000 tokens de prefijo frente a 1.000. Es el complemento natural del prefix caching: el block manager comparte la memoria del prefijo, y la cascade attention comparte el cómputo de atender sobre él.

Las matemáticas que importan: cuándo cambiar de backend te da algo

El backend solo mueve la aguja donde la atención es el cuello. En decode memory-bound, un kernel que aprovecha mejor el ancho de banda de HBM da una mejora real; en prefill compute-bound con secuencias largas, FA3/FA4 acercándose al pico de tensor cores da una mejora real. Pero si tu cuello está en otra capa —el launch overhead, el scheduler mal dimensionado, el cold start— cambiar de backend no toca esa parte. La regla, otra vez: medir el régimen antes de optimizar.

Los 10 knobs

#	Knob	Qué controla	Coste / riesgo
1	`VLLM_ATTENTION_BACKEND`	forzar backend	mismatch con hardware/feature
2	`flash_attn_version` (2/3/4)	versión de FA	versión no soportada en tu SM
3	habilitar FlashInfer	JIT + cascade	tiempo de compilación JIT inicial
4	cascade attention	reuso de cómputo de prefijo	solo ayuda con prefijo muy compartido
5	`kv_cache_dtype` (FP8)	soporte FP8 en el kernel	no todos los backends/SM lo soportan
6	`block_size`	layout que el kernel debe leer	coherencia con PagedAttention
7	backend de prefill MLA	kernel de la fase densa	solo modelos MLA
8	backend de decode MLA	kernel de la fase flaca	solo modelos MLA
9	soft cap / sliding window	features que limitan backends	menos opciones de kernel
10	head_dim / variante	qué kernels son elegibles	modelos exóticos sin soporte

Cómo se conecta con el resto del stack

Con FlashAttention. El post de FA explica el kernel por dentro (tiling, online softmax, FA1-4); este es el nivel de arriba —cómo vLLM elige entre kernels y por qué necesita más de uno—.

Con PagedAttention. El backend lee el KV que el block manager coloca en bloques; tiene que hablar el idioma del block table.

Con el scheduler. El scheduler arma batches mixtos prefill+decode; el backend tiene que atender los dos regímenes en un solo forward.

Con los CUDA graphs. Los kernels de atención se capturan en los CUDA graphs; un backend que lanza muchos kernels pequeños se beneficia más de la captura.

Con el prefix caching. La cascade attention es el lado cómputo de lo que el prefix caching hace en memoria.

Con FP8. Atender sobre KV en FP8 requiere que el backend tenga el camino FP8; no todos lo tienen en toda arquitectura.

Trampas y cosas que no son lo que parecen

“FlashInfer siempre es más rápido que FlashAttention.” No. FlashInfer gana cuando su especialización (cascade, KV heterogéneo, una variante de atención concreta) aplica a tu carga; en prefill denso clásico, FA3/FA4 suele ir igual o mejor. Depende del régimen, no hay un ganador universal.

“Un buen kernel de atención sirve para todo.” El error de fondo de este post. Prefill y decode son compute-bound y memory-bound respectivamente; un kernel ajustado a uno desperdicia en el otro. Por eso existen caminos separados (y backends separados en MLA).

“El decode es compute-bound porque la GPU está al 100%.” El nvidia-smi al 100% engaña (ver el post de SMs): el decode es memory-bound, la GPU está moviendo KV, no calculando. Optimizar el cómputo del decode es pulir lo que no es el cuello.

“Fijo VLLM_ATTENTION_BACKEND y me olvido.” Fijar un backend a mano puede dejarte en uno subóptimo cuando cambias de GPU o de versión, o forzar un fallback lento si tu hardware no soporta lo que pediste. La autoselección suele acertar; fíjalo solo con una medida que lo justifique.

“La cascade attention siempre ayuda.” Solo con prefijo muy compartido entre muchas peticiones concurrentes. Si cada petición tiene su propio contexto, no hay nada que compartir y el overhead de organizar la cascada no se amortiza.

“El backend de atención es el cuello, por eso voy lento.” Casi siempre el cuello está más arriba (lanzamiento, scheduling, memoria) o más abajo (ancho de banda). El backend importa donde la atención domina; mídelo con nsys/DCGM antes de cambiarlo.

Conclusión

De todo lo que hace un LLM al generar texto, casi todo son multiplicaciones de matrices que cualquier librería resuelve. El rendimiento se juega en un solo kernel —la atención— y la sorpresa es que ni siquiera es un kernel: son dos problemas opuestos disfrazados del mismo nombre. El prefill quiere fuego —cómputo denso sobre miles de tokens— y el decode quiere despensa rápida —leer todo el KV de un token con el mínimo desperdicio de ancho de banda—. Por eso vLLM no eligió un kernel ganador sino una abstracción que conmuta: FlashAttention afinado a cada arquitectura para el caso general, FlashInfer compilando a medida cuando hay heterogeneidad o prefijos que compartir, Triton para portabilidad. El jefe de cocina no cocina el plato estrella de una sola manera: mira quién pide y en qué momento del servicio, y manda al especialista que toca. La lección para quien tunea es la de siempre en esta serie: antes de cambiar de especialista, asegúrate de que el plato estrella es de verdad lo que te está frenando.

Ver también

FlashAttention v1/v2/v3/v4 — el kernel por dentro; este post es el nivel de arriba (cómo se elige entre kernels).
PagedAttention y el block manager — el KV paginado que el backend lee vía block table.
El pase: el scheduler step de vLLM — el batch mixto prefill+decode que el backend digiere en un forward.
SM, CUDA streams y CUDA graphs — por qué el nvidia-smi al 100% no significa compute-bound, y dónde se capturan los kernels de atención.
Prefix cache hit rate engineering — el lado memoria de lo que la cascade attention hace en cómputo.
FP8 end-to-end: pesos y KV — el camino FP8 que el backend necesita soportar.
Continuous batching — por qué un forward tiene que atender prefill y decode a la vez.
KV cache: la memoria de trabajo — el dato que el kernel de decode lee entero en cada paso.

Referencias

vLLM, Attention Backends (selección, FA2/3/4 por arquitectura, MLA): https://docs.vllm.ai/en/latest/design/attention_backends/.
vLLM / DeepWiki, FlashAttention and FlashInfer: https://deepwiki.com/vllm-project/vllm/8.2-flashattention-and-flashinfer.
vLLM, Triton Attention Backend Deep Dive (mar-2026): https://vllm.ai/blog/2026-03-04-vllm-triton-backend-deep-dive.
Z. Ye et al., FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving (arXiv 2501.01005): https://arxiv.org/abs/2501.01005.
T. Dao, FlashAttention-2 / FlashAttention-3 (kernel IO-aware, async Hopper): https://github.com/Dao-AILab/flash-attention.

Cascade-Attention on lo0 — Blog Técnico