FlashAttention v1/v2/v3/v4: el bibliotecario que nunca despeja la mesa — IO-awareness, async y la asimetría de Blackwell

Fri, 29 May 2026 17:00:00 +0200

Este post complementa los de KV cache: la memoria de trabajo y PagedAttention deep dive. KV cache explica qué se almacena; PagedAttention, cómo se gestiona en memoria; FlashAttention, cómo se ejecuta el cálculo. Son tres capas distintas del mismo problema y se acumulan multiplicativamente.

Estás aquí: DEPLOY

TL;DR

El attention estándar de un transformer tiene un problema estructural en GPUs modernas: cuando se mira con un profiler, no está limitado por compute, está limitado por memoria. La matriz S = QK^T de tamaño N × N —con N la longitud de la secuencia— no cabe en la SRAM rápida del chip y obliga a hacer round trips a HBM que dominan el tiempo total. FlashAttention es la familia de kernels que evita materializar esa matriz haciendo tiling sobre Q, K y V, calculando el softmax bloque a bloque con la versión online y manteniendo todo dentro de SRAM. Cada versión sube el techo de utilización: FA1 (Dao et al. 2022) abolió la matriz N×N y bajó complejidad de memoria a O(N); FA2 (Dao 2023) paralelizó a lo largo de la dimensión de secuencia y redujo los non-matmul FLOPs hasta acercarse al 70 % de utilización en A100; FA3 (Shah, Bikshandi, Zhang, Thakkar, Ramani, Dao 2024) explotó tres mecanismos específicos de Hopper —WGMMA async, TMA y FP8— para llegar a 740 TFLOPS BF16 (75 % del peak) y 1.2 PFLOPS FP8 en H100; FA4 (marzo 2026) rescribió el kernel desde cero para Blackwell, donde el tensor core escala 2.25× pero la SFU (donde corre el exp del softmax) y la SMEM bandwidth no escalan nada — la solución es software-emulated exponential que corre en tensor cores. El resultado en B200: 1605 TFLOPS BF16, 1.3× más rápido que cuDNN 9.13 y 2× más rápido que FA3 en la misma GPU. Este post desmonta el porqué (memory roofline, IO complexity), la analogía maestra del bibliotecario, las matemáticas mínimas y los números reales en H100 y B200.

La analogía: el bibliotecario que nunca despeja la mesa

Una biblioteca grande con dos zonas: una mesa de trabajo muy rápida pero pequeña (el escritorio del bibliotecario, 200 libros caben encima), y una estantería gigantesca que recorre tres pisos donde está todo (50 millones de libros). El bibliotecario tiene que cruzar referencias entre todos los libros de una sala temática y producir un resumen.

Un bibliotecario ingenuo lo hace de la forma directa: trae todos los libros relevantes de la estantería, los apila sobre la mesa de trabajo, y como no caben, deja la mitad en el suelo, en la silla, sobre cajas. Pasa el día corriendo entre el suelo y la mesa, abriendo y cerrando libros, sin poder concentrarse. La mesa de trabajo es rapidísima pero está infrautilizada porque la mayoría del tiempo el bibliotecario está moviendo libros entre el suelo y la mesa. Es lo que hace standard attention: materializa la matriz S = QK^T en HBM y vuelve una y otra vez por trozos.

Un bibliotecario FlashAttention v1 cambia de estrategia: pide un estante a la vez, lo trae a la mesa, lee lo que necesita, anota notas en una libreta compacta —“de este estante me interesa esto, esto y esto, con estos pesos relativos”—, devuelve el estante a su sitio y trae el siguiente. La libreta es lo único que se conserva entre estantes. Nunca apila más de lo que cabe en la mesa. El truco que hace esto posible es la online softmax: en lugar de necesitar todo el contenido a la vez para normalizar, mantiene un running max y un running sum que se actualizan estante a estante.

Un bibliotecario FlashAttention v2 se da cuenta de que puede trabajar varios temas en paralelo porque la mesa es grande y los estantes son independientes en algunos ejes. Pone tres ayudantes con sus libretas, cada uno cubriendo un bloque distinto de la sala, y combinan resultados al final.

Un bibliotecario FlashAttention v3 consigue mecanizar el flujo: instala una cinta transportadora con dos estaciones. Mientras la estación A lee el estante actual y toma notas, la estación B ya tiene el siguiente estante en tránsito desde la estantería. Cuando A termina, B le pasa el estante nuevo sin esperar. Es el ping-pong producer/consumer: la TMA descarga el data loading y los warps consumidores hacen el trabajo en paralelo con las cargas. Encima, las notas se escriben en taquigrafía de menor precisión (FP8) porque las páginas que importan llevan ya un pre-tratamiento ortogonal que no las hace perder precisión donde duele.

Un bibliotecario FlashAttention v4 descubre algo nuevo de su biblioteca remodelada (Blackwell): le han instalado dos cintas transportadoras mucho más rápidas (tensor cores 2.25×), pero la máquina de escribir taquigráfica no la han actualizado (la SFU sigue igual). Ahora el cuello de botella es escribir las notas, no traer los estantes. La solución es elegante: en lugar de usar la máquina taquigráfica, escribe las anotaciones con fórmulas polinómicas que el propio tensor core puede evaluar (software-emulated exponential). La cinta no se queda parada esperando a la máquina.

La analogía se sostiene con cuatro mapeos:

Mesa de trabajo = SRAM por SM (228 KB en H100, 256 KB en B200).
Estantería gigante = HBM (3.35 TB/s en H100, 8 TB/s en B200).
Libreta con max y sum running = stats de la online softmax (m, ℓ).
Cinta transportadora con dos estaciones = pipeline TMA + WGMMA producer/consumer.

Por qué standard attention era el cuello de botella

La intuición ingenua —“attention son matmuls, las GPUs son buenas con matmuls”— es correcta pero incompleta. Hay dos matmuls (QK^T y luego softmax(S) V), y en medio una operación no-matmul (softmax) que requiere materializar la matriz intermedia S de tamaño N × N.

Las GPUs modernas tienen un roofline muy concreto. Para H100 SXM5:

Compute peak: 989 TFLOPS BF16 (tensor core, dense, sin sparsity).
Memory bandwidth: 3.35 TB/s HBM3.
Punto de cruce (arithmetic intensity break-even): 989 × 10¹² / (3.35 × 10¹²) ≈ 295 FLOP/byte.

Cualquier operación que no consiga ejecutar 295 operaciones por cada byte que mueve desde HBM está memory-bound: el tensor core se queda esperando datos. Para B200 el ratio es similar (~281 FLOP/byte) porque tanto compute como bandwidth subieron.

Standard attention materializando S lee la matriz dos veces (una para softmax, otra para multiplicar por V) y la escribe una vez. Para Llama 3 70B con head dim d = 128 y contexto N = 128K:

Matriz S por head por capa: N × N × 2 bytes = 128K × 128K × 2 = 34.36 GB.
80 capas × 64 Q-heads → tráfico HBM agregado (si se materializara serialmente) del orden de TBs, prohibitivo en transient.
Aunque no se materialice todo a la vez, los round trips dominan el tiempo: arithmetic intensity efectiva muy por debajo de 295 → operación memory-bound → tensor core infrautilizado a ~25 %.

FlashAttention no cambia las matemáticas del attention, cambia el orden de las operaciones para no materializar S en HBM. Su IO complexity es Θ(N²·d²/M) con M = tamaño SRAM por SM, frente a Θ(N·d + N²) de standard attention. Con d = 128 y M = 228 KB: factor de reducción aproximado M/d² ≈ 14× menos tráfico HBM. Eso es lo que mueve la operación de memory-bound a casi compute-bound.

El truco que hizo posible todo: online softmax

Sin online softmax no hay FlashAttention. La idea es de Milakov y Gimelshein, 2018 (paper “Online normalizer calculation for softmax”, arXiv:1805.02867), y permite calcular softmax([x_1, ..., x_N]) en un pase incremental sin necesitar conocer el máximo global antes de empezar.

El softmax estándar es:

$$\text{softmax}(x_i) = \frac{e^{x_i - m}}{\sum_{j} e^{x_j - m}}, \quad m = \max_j x_j$$

El truco online: mantén un máximo running m^{(t)} y una suma running ℓ^{(t)}. Cuando llega un bloque nuevo de valores con máximo local m_{\text{new}}:

$$m^{(t+1)} = \max(m^{(t)}, m_{\text{new}})$$

$$\ell^{(t+1)} = e^{m^{(t)} - m^{(t+1)}} \cdot \ell^{(t)} + \sum_{j \in \text{nuevo}} e^{x_j - m^{(t+1)}}$$

Y los outputs parciales acumulados también se reescalan por el mismo factor e^{m^{(t)} - m^{(t+1)}}. Al final, dividir por ℓ final da exactamente lo mismo que el softmax estándar. Es matemáticamente exacto, no es una aproximación.

Esto es lo que permite recorrer K bloque a bloque sin materializar la matriz S entera. Cada bloque actualiza los stats y los outputs acumulados, y se descarta. La mesa nunca se llena.

Las cuatro versiones (mayo 2026)

	FA1 (2022)	FA2 (2023)	FA3 (2024)	FA4 (2026)
GPU target	A100 / Ampere	A100 / H100	H100 / Hopper	B200 / Blackwell
Idea central	Tiling + online softmax	Sequence parallelism + work partitioning	Async WGMMA + TMA + FP8	Polynomial exp + 2-CTA tensor cores
Memoria	O(N) (vs O(N²))	igual	igual	igual
Peak util típica	~25 % A100	~70 % A100, ~35 % H100	75 % H100 BF16, 60 % H100 FP8	71 % B200 BF16
TFLOPS efectivos	—	225 TFLOPS A100 BF16	740 H100 BF16, 1200 H100 FP8	1605 B200 BF16
Speedup vs anterior	2-4× standard	2× FA1	1.5-2× FA2 (BF16), 2.6× (FP8)	2× FA3 en B200
Paper	arXiv:2205.14135	arXiv:2307.08691	arXiv:2407.08608	arXiv:2603.05451

FA1 — el cambio de orden de las operaciones

Tres ideas combinadas: tiling de Q/K/V en bloques que caben en SRAM, online softmax sobre esos bloques, y recomputation en backward —no se guarda la matriz S de tamaño N×N, solo los stats (m, ℓ), y en backward se recomputa S bloque a bloque a partir de Q, K y los stats—. Resultado: 7.6× speedup en GPT-2 vs PyTorch standard attention, memoria O(N).

FA2 — paralelizar en serie

FA1 paralelizaba solo en batch × heads. Con batch pequeño (1-4) y modelos con pocos heads o GQA agresivo, la GPU se quedaba con SMs ociosos. FA2 paraleliza también en la dimensión de secuencia: distintos SMs procesan distintos tramos de Q al mismo tiempo. Además reescribe el algoritmo para minimizar las operaciones no-matmul (rescaling del softmax) porque esas no pasan por tensor cores. Y mejora el work partitioning entre warps (split-Q en lugar de split-K reduce tráfico de shared memory). Resultado: ~2× sobre FA1 en H100 y A100, 225 TFLOPS en A100 (72 % MFU). En H100 se queda en torno al 30-35 % del peak BF16 porque no aprovecha WGMMA async.

FA3 — el momento Hopper

Aquí FlashAttention deja de ser un algoritmo y se convierte en una pieza específica de Hopper. Tres pilares:

WGMMA async: las instrucciones nuevas de tensor core de Hopper permiten que un warpgroup dispare un GEMM y el resto del warpgroup haga otra cosa (la softmax, por ejemplo) mientras el tensor core sigue trabajando. Es el truco que destraba el solapamiento matmul/softmax.
TMA (Tensor Memory Accelerator): hardware dedicado a copiar tiles entre HBM y SRAM. Libera al SM del trabajo de calcular índices y predicar out-of-bounds, que antes ocupaba ciclos del propio SM. Es el equivalente a contratar mozos de almacén: el bibliotecario deja de tener que cargar libros él mismo.
FP8 con block quantization + incoherent processing: cuantizar Q y K a FP8 dobla el throughput del tensor core. La pérdida de precisión se mitiga con dos trucos: una escala por tile (64×d) en lugar de por tensor entero, y una pre-multiplicación por una matriz ortogonal aleatoria basada en Hadamard que “esparce” los outliers antes de cuantizar. Resultado documentado: error numérico 2.6× menor que FP8 baseline.

Estos tres pilares se combinan con warp specialization producer/consumer (warps productores hacen TMA loads; warps consumidores hacen WGMMA + softmax) y un ping-pong scheduling con dos warpgroups que se turnan para que nunca haya pipeline bubbles. Cuando WG1 hace softmax, WG2 hace GEMM; luego se intercambian.

Números: 740 TFLOPS BF16 en H100 (75 % del peak 989), 1.2 PFLOPS FP8 (60 % del peak 1978 FP8 dense). Para secuencias ≥ 1K supera a cuDNN. Speedup sobre FA2: 1.5-2× BF16, 2.6× FP8.

FA4 — la asimetría de Blackwell

Blackwell escaló todo de forma desigual:

Tensor core BF16 throughput: 1 PFLOP H100 → 2.25 PFLOPS B200 (2.25×).
SFU count (donde corre exp del softmax): sin cambios.
Shared memory bandwidth: sin cambios.

Es decir, si FA3 corre tal cual en B200 sin tocarlo, el matmul va el doble de rápido pero el softmax queda exactamente igual, y eso bloquea el pipeline. Era cuestión de tiempo que alguien resolviera el desequilibrio.

FA4 (marzo 2026, mismo equipo Dao + Princeton + Together AI + Meta + NVIDIA + Colfax) hace un rewrite ground-up con tres ideas:

Software-emulated exponential: aproximación polinómica del exp que se ejecuta en el tensor core en lugar de en la SFU. Pierde un poquito de precisión (cuidadosamente acotada y compensada por el resto del kernel) pero deja la cinta transportadora moviéndose.
Conditional softmax rescaling: evita rescalar acumuladores cuando el running max no cambia significativamente. Optimización de tipo “lazy”: solo paga el coste cuando hace falta.
2-CTA tensor core: dos CTAs (Cooperative Thread Arrays) colaboran para alimentar los tensor cores con tiles más grandes. Saca más juego de las capacidades nuevas de Blackwell.

Escrito en CuTeDSL (Python DSL de NVIDIA CUTLASS, no CUDA C++ directo). Resultado en B200 BF16: 1605 TFLOPS (71 % del peak 2250). 1.3× sobre cuDNN 9.13. 2.7× sobre Triton. 2× sobre FA3 ejecutado tal cual en B200 (que era el baseline anterior). Es el primer kernel de attention que pasa de 1 PFLOPS.

Nota: hay confusión recurrente con “FP4 attention”. Las extensiones NVFP4/MXFP4 de Blackwell se aplican a pesos, no a attention. FA4 puede combinarse con weights NVFP4, pero el cómputo de attention en sí sigue siendo BF16 o FP8 según configuración. La cuantización a FP4 de QK^T existe en algunos kernels propietarios (FireAttention V4 de Fireworks AI lo combina) pero no es la práctica estándar.

Implementaciones y librerías en 2026

Dao-AILab/flash-attention (repo canónico): soporta SM 8.0 (Ampere) con FA2, SM 9.0 (Hopper) con FA3, SM 10.0 (Blackwell datacenter B100/B200/B300) con FA4. La versión consumer Blackwell (5090, SM 12.0) tiene soporte parcial al cierre de este post.
FlashInfer (flashinfer-ai/flashinfer, arXiv:2501.01005): engine de attention orientado a serving (no a training). Su contribución conceptual es el Block-Sparse Row (BSR), una abstracción unificada que cubre paged KV cache, radix tree de prefix caching y máscaras de árbol de speculative decoding. Internamente puede llamar a kernels FA2/FA3, cuDNN, CUTLASS, o trtllm-gen FMHA según el caso. JIT compila variantes específicas en runtime. Integrado en vLLM, SGLang, TensorRT-LLM.
vLLM (mayo 2026): selección automática del backend según GPU. Default FA4 en SM 10.0+, FA3 en SM 9.0, FA2 en resto. Fallbacks en Blackwell: TRT-LLM Ragged → FlashInfer → TokenSpeed MLA. Para FP8 KV cache en B200, FlashInfer es competitivo.
SGLang: usa FlashInfer como backend de attention; RadixAttention es la capa de prefix caching encima (un radix tree del KV cache).
TensorRT-LLM: kernels fused propios (trtllm-gen FMHA). XQA es la optimización propia de NVIDIA para GQA en decode.
PyTorch SDPA y FlexAttention: torch.nn.functional.scaled_dot_product_attention selecciona backend automático. FlexAttention (nuevo) permite definir custom masks declarativamente y compila a kernels que pueden usar FA4 como backend.
xFormers: sigue vivo pero residual. PyTorch SDPA built-in cubre la mayoría de casos.

Casos donde FlashAttention no ayuda

Contextos muy cortos (N < 512): el overhead de tiling y kernel launch no compensa; cuDNN puede ganar.
Custom masks no estándar: FA solo trae causal, sliding window y ALiBi de serie. Para máscaras arbitrarias hace falta FlexAttention o variantes JIT de FlashInfer.
Head dim no estándar: FA optimiza para d = 64, 128, 256. Dimensiones extrañas (d = 96, d = 192) caen en paths lentos.
GQA/MQA con ratios extremos: soportado nativo, pero el speedup vs MHA puro depende del ratio Q-heads : KV-heads.
Cross-attention: soportado pero menos optimizado; el caso self-attention es donde más ganancia hay.
FP8 sin block quantization ni incoherent processing: pierde varios puntos en benchmarks. Si tu serving framework no implementa los dos trucos del FA3 paper, FP8 attention puede ser una mala idea.

Implicaciones en hardware on-premise

En una RTX 4090 (24 GB, Ada Lovelace, SM 8.9)

La 4090 es Ada Lovelace, no Hopper. No corre FA3 ni FA4; corre FA2. Eso significa: ~70 % de utilización en attention BF16 (~250 TFLOPS efectivos sobre el peak de 330 TFLOPS BF16 de la 4090). No es trágico —FA2 ya es muy bueno comparado con standard attention— pero el techo está claramente por debajo del de un H100. Para deploys consumer en 4090 con Llama 3 8B BF16 o cualquier 14B-32B INT4 AWQ, FA2 es lo que vas a estar usando, y es perfectamente razonable.

En un cluster genérico 4×H100 SXM (320 GB, NVLink, FP8 nativo)

Aquí FA3 brilla y es lo que vLLM/SGLang/TRT-LLM van a seleccionar por defecto. Dos configuraciones comunes:

Llama 3 70B FP8 con FA3 FP8 attention: 1.2 PFLOPS pico en la GPU, throughput agregado del cluster en el orden de 8000-12000 tokens/s en batch medio dependiendo de TP y contexto. Para que el FP8 attention dé su pleno rendimiento es crucial usar las técnicas de block quantization + incoherent processing del paper FA3 (vienen activadas en vLLM por defecto).
DeepSeek-V3 671B FP8 + MLA con FlashInfer: DeepSeek usa Multi-head Latent Attention (MLA), una variante distinta del attention estándar. FlashInfer tiene kernels específicos (FlashMLA). El stack típico es vLLM/SGLang + FlashInfer + FlashMLA + FA3 fallback para las capas no-MLA.

Si la infraestructura es Blackwell (B200/B300, que algunos clusters empiezan a recibir en 2026), FA4 es la opción correcta y debería estar habilitado por defecto en vLLM 0.16+ y SGLang 0.5.11+.

Lo que no hemos cubierto

MLA (Multi-head Latent Attention) de DeepSeek y los kernels FlashMLA específicos: optimizan compresión de KV cache pero requieren kernels distintos.
Flexible masking y los casos de uso de FlexAttention (PyTorch 2.5+): cómo declarar máscaras arbitrarias sin pagar el coste de un kernel custom.
Asistencia de hardware para sparse attention (NVIDIA sparse tensor cores 2:4) y por qué attention sparse no ha consolidado como techo más alto que FA dense.
FA en backward de fine-tuning: el post se centra en inferencia, pero FA3/FA4 también pasan por backward y son lo que hace viable entrenar modelos con contextos largos en H100/B200.

Ver también

KV cache: la memoria de trabajo que sostiene la inferencia LLM — el KV cache es lo que FlashAttention recorre y multiplica contra Q en cada iteración; entender uno requiere entender el otro.
PagedAttention deep dive — PagedAttention organiza el KV cache en bloques físicos no contiguos; FlashAttention es el kernel que itera sobre esos bloques. Capas distintas del mismo problema.
Quantization para inferencia LLM — FP8 attention de FA3 con block quantization y FP4 weights de Blackwell se acumulan; este post da el marco de cuantización general.
Disaggregated serving: prefill y decode en pods especializados — prefill es compute-bound y se beneficia mucho de FA3/FA4 FP8; decode es memory-bound y se beneficia menos pero igualmente. La separación deja optimizar el kernel por fase.
Speculative decoding: el secretario que adelanta lo que va a decir el jefe — speculative produce más tokens por forward pass; FlashAttention hace cada forward pass más barato. Palancas multiplicativas.
El pipeline LLMOps de seis etapas — el mapa maestro donde Deploy es la etapa 4.

Referencias

Dao, T., Fu, D., Ermon, S., Rudra, A., Ré, C. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022. https://arxiv.org/abs/2205.14135
Dao, T. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. 2023. https://arxiv.org/abs/2307.08691
Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., Dao, T. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. NeurIPS 2024. https://arxiv.org/abs/2407.08608
Dao, T. et al. FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling. 2026. https://arxiv.org/abs/2603.05451
Milakov, M., Gimelshein, N. Online normalizer calculation for softmax. 2018. https://arxiv.org/abs/1805.02867
Ye, Z. et al. FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. MLSys 2025. https://arxiv.org/abs/2501.01005
Tri Dao FA3 blog: https://tridao.me/blog/2024/flash3/
Tri Dao FA4 blog: https://tridao.me/blog/2026/flash4/
PyTorch FlashAttention-3 announcement: https://pytorch.org/blog/flashattention-3/
PyTorch FlexAttention + FA4: https://pytorch.org/blog/flexattention-flashattention-4-fast-and-flexible/
Together AI FA4 blog: https://www.together.ai/blog/flashattention-4
Colfax Research FA3: https://research.colfax-intl.com/flashattention-3-fast-and-accurate-attention-with-asynchrony-and-low-precision/
Colfax Research FA4: https://research.colfax-intl.com/flashattention-4-algorithm-and-kernel-pipelining-co-design-for-asymmetric-hardware-scaling/
Repo Dao-AILab/flash-attention: https://github.com/Dao-AILab/flash-attention
Repo flashinfer-ai/flashinfer: https://github.com/flashinfer-ai/flashinfer
vLLM attention backends: https://docs.vllm.ai/en/latest/design/attention_backends/
NVIDIA Hopper Architecture in Depth: https://developer.nvidia.com/blog/nvidia-hopper-architecture-in-depth/

Fa3 on lo0 — Blog Técnico