FlashAttention v1/v2/v3/v4: the librarian who never clears the desk — IO-awareness, async and Blackwell's asymmetry

· By José Manuel Román
Contents

This post complements KV cache: the working memory and PagedAttention deep dive. The KV cache post explains what is stored; PagedAttention, how it is managed in memory; FlashAttention, how the computation is executed. They are three different layers of the same problem and they stack multiplicatively.

You are here: DEPLOY

You are here: DEPLOY · attention kernel, one layer below PagedAttention and KV cache1 · Data2 · Tune3 · Eval4 · Deploy5 · Observe6 · Retrain

TL;DR

The standard attention of a transformer has a structural problem on modern GPUs: when you look at it with a profiler, it is not compute-bound, it is memory-bound. The S = QK^T matrix of size N × N, with N the sequence length, does not fit in the fast on-chip SRAM and forces round trips to HBM that dominate the total time. FlashAttention is the family of kernels that avoids materialising that matrix by tiling over Q, K and V, computing the softmax block by block with the online version and keeping everything inside SRAM. Each version raises the utilisation ceiling: FA1 (Dao et al. 2022) abolished the N×N matrix and brought memory complexity down to O(N); FA2 (Dao 2023) parallelised along the sequence dimension and cut non-matmul FLOPs to get close to 70 % utilisation on A100; FA3 (Shah, Bikshandi, Zhang, Thakkar, Ramani, Dao 2024) exploited three Hopper-specific mechanisms, async WGMMA, TMA and FP8, to reach 740 TFLOPS BF16 (75 % of peak) and 1.2 PFLOPS FP8 on H100; FA4 (March 2026) rewrote the kernel from scratch for Blackwell, where the tensor core scales 2.25× but the SFU (where the softmax exp runs) and the SMEM bandwidth do not scale at all. The answer is a software-emulated exponential that runs on tensor cores. The result on B200: 1,605 TFLOPS BF16, 1.3× faster than cuDNN 9.13 and 2× faster than FA3 on the same GPU. This post takes apart the why (memory roofline, IO complexity), the master analogy of the librarian, the minimum mathematics and the real numbers on H100 and B200.

The analogy: the librarian who never clears the desk

A large library with two zones: a very fast but small work desk (the librarian’s desk, 200 books fit on top of it), and a giant shelving system spanning three floors where everything lives (50 million books). The librarian has to cross-reference all the books in one thematic room and produce a summary.

A naive librarian does it the direct way: fetches every relevant book from the shelves, piles them on the work desk, and since they do not fit, leaves half of them on the floor, on the chair, on top of boxes. The day is spent running between the floor and the desk, opening and closing books, never able to concentrate. The work desk is extremely fast but underused, because most of the time the librarian is moving books between the floor and the desk. That is what standard attention does: it materialises the S = QK^T matrix in HBM and keeps coming back for it piece by piece.

A FlashAttention v1 librarian changes strategy: asks for one shelf at a time, brings it to the desk, reads what is needed, jots notes in a compact notebook (“from this shelf I care about this, this and this, with these relative weights”), returns the shelf to its place and brings the next one. The notebook is the only thing that carries over between shelves. Nothing is ever piled higher than the desk can hold. The trick that makes this possible is the online softmax: instead of needing all the content at once in order to normalise, it keeps a running max and a running sum that are updated shelf by shelf.

A FlashAttention v2 librarian realises that several topics can be worked in parallel, because the desk is big and the shelves are independent along some axes. Three assistants are put to work, each with a notebook, each covering a different block of the room, and the results are combined at the end.

A FlashAttention v3 librarian manages to mechanise the flow: a conveyor belt with two stations is installed. While station A reads the current shelf and takes notes, station B already has the next shelf in transit from the shelving system. When A finishes, B hands over the new shelf with no waiting. This is the ping-pong producer/consumer pattern: the TMA offloads the data loading and the consumer warps do the work in parallel with the loads. On top of that, the notes are written in lower-precision shorthand (FP8), because the pages that matter have already had an orthogonal pre-treatment that stops them losing precision where it hurts.

A FlashAttention v4 librarian discovers something new about the remodelled library (Blackwell): two much faster conveyor belts have been installed (tensor cores 2.25×), but the shorthand typewriter has not been upgraded (the SFU is unchanged). Now the bottleneck is writing the notes, not fetching the shelves. The fix is elegant: instead of using the shorthand machine, the notes are written with polynomial formulas that the tensor core itself can evaluate (software-emulated exponential). The belt no longer sits idle waiting for the machine.

The analogy holds up with four mappings:

  • Work desk = SRAM per SM (228 KB on H100, 256 KB on B200).
  • Giant shelving system = HBM (3.35 TB/s on H100, 8 TB/s on B200).
  • Notebook with running max and sum = online softmax stats (m, ℓ).
  • Conveyor belt with two stations = TMA + WGMMA producer/consumer pipeline.

Why standard attention was the bottleneck

The naive intuition, “attention is matmuls, and GPUs are good at matmuls”, is correct but incomplete. There are two matmuls (QK^T and then softmax(S) V), and in between a non-matmul operation (softmax) that requires materialising the intermediate matrix S of size N × N.

Modern GPUs have a very concrete roofline. For H100 SXM5:

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

Any operation that fails to perform 295 operations for every byte it moves from HBM is memory-bound: the tensor core sits waiting for data. On B200 the ratio is similar (~281 FLOP/byte) because both compute and bandwidth went up.

Standard attention materialising S reads the matrix twice (once for the softmax, once to multiply by V) and writes it once. For Llama 3 70B with head dim d = 128 and context N = 128K:

  • Matrix S per head per layer: N × N × 2 bytes = 128K × 128K × 2 = 34.36 GB.
  • 80 layers × 64 Q-heads → aggregate HBM traffic (if it were materialised serially) on the order of TBs, prohibitive in transient.
  • Even if it is not all materialised at once, the round trips dominate the time: effective arithmetic intensity well below 295 → memory-bound operation → tensor core underused at around 25 %.

FlashAttention does not change the mathematics of attention, it changes the order of the operations so that S is never materialised in HBM. Its IO complexity is Θ(N²·d²/M) with M = SRAM size per SM, against Θ(N·d + N²) for standard attention. With d = 128 and M = 228 KB: an approximate reduction factor of M/d² ≈ 14× less HBM traffic. That is what moves the operation from memory-bound to nearly compute-bound.

The trick that made it all possible: online softmax

Without online softmax there is no FlashAttention. The idea comes from Milakov and Gimelshein, 2018 (the paper “Online normalizer calculation for softmax”, arXiv:1805.02867), and it allows softmax([x_1, ..., x_N]) to be computed in one incremental pass without needing to know the global maximum before starting.

The standard softmax is:

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

The online trick: keep a running maximum m^{(t)} and a running sum ℓ^{(t)}. When a new block of values arrives with local maximum 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{new}} e^{x_j - m^{(t+1)}}$$

And the accumulated partial outputs are rescaled by the same factor e^{m^{(t)} - m^{(t+1)}}. At the end, dividing by the final ℓ gives exactly the same result as the standard softmax. It is mathematically exact, not an approximation.

This is what allows K to be traversed block by block without materialising the whole S matrix. Each block updates the stats and the accumulated outputs, and is then discarded. The desk never fills up.

The four versions (May 2026)

FA1 (2022)FA2 (2023)FA3 (2024)FA4 (2026)
Target GPUA100 / AmpereA100 / H100H100 / HopperB200 / Blackwell
Core ideaTiling + online softmaxSequence parallelism + work partitioningAsync WGMMA + TMA + FP8Polynomial exp + 2-CTA tensor cores
MemoryO(N) (vs O(N²))samesamesame
Typical peak util~25 % A100~70 % A100, ~35 % H10075 % H100 BF16, 60 % H100 FP871 % B200 BF16
Effective TFLOPS—225 TFLOPS A100 BF16740 H100 BF16, 1,200 H100 FP81,605 B200 BF16
Speedup vs previous2-4× standard2× FA11.5-2× FA2 (BF16), 2.6× (FP8)2× FA3 on B200
PaperarXiv:2205.14135arXiv:2307.08691arXiv:2407.08608arXiv:2603.05451

FA1 — changing the order of the operations

Three ideas combined: tiling of Q/K/V into blocks that fit in SRAM, online softmax over those blocks, and recomputation in the backward pass. The N×N matrix S is not stored, only the stats (m, ℓ), and in the backward pass S is recomputed block by block from Q, K and the stats. Result: 7.6× speedup on GPT-2 versus PyTorch standard attention, memory O(N).

FA2 — parallelising along the sequence

FA1 parallelised only over batch × heads. With a small batch (1-4) and models with few heads or aggressive GQA, the GPU was left with idle SMs. FA2 also parallelises along the sequence dimension: different SMs process different stretches of Q at the same time. It also rewrites the algorithm to minimise the non-matmul operations (softmax rescaling), because those do not go through tensor cores. And it improves work partitioning between warps (split-Q instead of split-K reduces shared memory traffic). Result: around 2× over FA1 on H100 and A100, 225 TFLOPS on A100 (72 % MFU). On H100 it stays around 30-35 % of BF16 peak because it does not exploit async WGMMA.

FA3 — the Hopper moment

This is where FlashAttention stops being an algorithm and becomes a Hopper-specific artefact. Three pillars:

  1. Async WGMMA: Hopper’s new tensor core instructions let a warpgroup fire a GEMM and have the rest of the warpgroup do something else (the softmax, for instance) while the tensor core keeps working. This is the trick that unlocks matmul/softmax overlap.
  2. TMA (Tensor Memory Accelerator): dedicated hardware for copying tiles between HBM and SRAM. It frees the SM from computing indices and predicating out-of-bounds accesses, which used to consume SM cycles. It is the equivalent of hiring warehouse staff: the librarian no longer has to carry the books personally.
  3. FP8 with block quantisation + incoherent processing: quantising Q and K to FP8 doubles tensor core throughput. The precision loss is mitigated with two tricks: a per-tile scale (64×d) instead of a whole-tensor scale, and a pre-multiplication by a random orthogonal Hadamard-based matrix that “spreads” the outliers before quantising. Documented result: numerical error 2.6× lower than the FP8 baseline.

These three pillars combine with producer/consumer warp specialisation (producer warps do TMA loads; consumer warps do WGMMA + softmax) and ping-pong scheduling with two warpgroups taking turns so that there are never pipeline bubbles. When WG1 does softmax, WG2 does GEMM; then they swap.

Numbers: 740 TFLOPS BF16 on H100 (75 % of the 989 peak), 1.2 PFLOPS FP8 (60 % of the 1978 FP8 dense peak). For sequences ≥ 1K it beats cuDNN. Speedup over FA2: 1.5-2× BF16, 2.6× FP8.

FA4 — Blackwell’s asymmetry

Blackwell scaled everything unevenly:

  • Tensor core BF16 throughput: 1 PFLOP H100 → 2.25 PFLOPS B200 (2.25×).
  • SFU count (where the softmax exp runs): unchanged.
  • Shared memory bandwidth: unchanged.

In other words, if FA3 runs as-is on a B200 with no changes, the matmul goes twice as fast but the softmax stays exactly the same, and that stalls the pipeline. It was only a matter of time before someone fixed the imbalance.

FA4 (March 2026, same team of Dao + Princeton + Together AI + Meta + NVIDIA + Colfax) is a ground-up rewrite with three ideas:

  1. Software-emulated exponential: a polynomial approximation of exp that runs on the tensor core instead of on the SFU. It loses a little precision (carefully bounded and compensated by the rest of the kernel) but keeps the conveyor belt moving.
  2. Conditional softmax rescaling: avoids rescaling accumulators when the running max does not change significantly. A “lazy” style optimisation: it only pays the cost when the cost is needed.
  3. 2-CTA tensor core: two CTAs (Cooperative Thread Arrays) cooperate to feed the tensor cores with larger tiles. It gets more out of Blackwell’s new capabilities.

Written in CuTeDSL (NVIDIA’s CUTLASS Python DSL, not CUDA C++ directly). Result on B200 BF16: 1,605 TFLOPS (71 % of the 2250 peak). 1.3× over cuDNN 9.13. 2.7× over Triton. 2× over FA3 run as-is on B200 (which was the previous baseline). It is the first attention kernel to pass 1 PFLOPS.

Note: there is recurring confusion about “FP4 attention”. Blackwell’s NVFP4/MXFP4 extensions apply to weights, not to attention. FA4 can be combined with NVFP4 weights, but the attention computation itself is still BF16 or FP8 depending on the configuration. FP4 quantisation of QK^T does exist in some proprietary kernels (Fireworks AI’s FireAttention V4 combines it) but it is not standard practice.

Implementations and libraries in 2026

  • Dao-AILab/flash-attention (the canonical repo): supports SM 8.0 (Ampere) with FA2, SM 9.0 (Hopper) with FA3, SM 10.0 (Blackwell datacenter B100/B200/B300) with FA4. The consumer Blackwell version (5090, SM 12.0) has partial support as this post goes out.
  • FlashInfer (flashinfer-ai/flashinfer, arXiv:2501.01005): an attention engine aimed at serving (not training). Its conceptual contribution is the Block-Sparse Row (BSR), a unified abstraction covering paged KV cache, the radix tree of prefix caching and the tree masks of speculative decoding. Internally it can call FA2/FA3, cuDNN, CUTLASS or trtllm-gen FMHA kernels depending on the case. It JIT-compiles specific variants at runtime. Integrated into vLLM, SGLang and TensorRT-LLM.
  • vLLM (May 2026): automatic backend selection by GPU. Default FA4 on SM 10.0+, FA3 on SM 9.0, FA2 on the rest. Fallbacks on Blackwell: TRT-LLM Ragged → FlashInfer → TokenSpeed MLA. For FP8 KV cache on B200, FlashInfer is competitive.
  • SGLang: uses FlashInfer as its attention backend; RadixAttention is the prefix caching layer on top (a radix tree over the KV cache).
  • TensorRT-LLM: its own fused kernels (trtllm-gen FMHA). XQA is NVIDIA’s own optimisation for GQA in decode.
  • PyTorch SDPA and FlexAttention: torch.nn.functional.scaled_dot_product_attention selects the backend automatically. FlexAttention (new) lets you define custom masks declaratively and compiles to kernels that can use FA4 as a backend.
  • xFormers: still alive but residual. The built-in PyTorch SDPA covers most cases.

Cases where FlashAttention does not help

  • Very short contexts (N < 512): the tiling and kernel launch overhead does not pay off; cuDNN can win.
  • Non-standard custom masks: FA ships causal, sliding window and ALiBi only. Arbitrary masks need FlexAttention or the JIT variants of FlashInfer.
  • Non-standard head dim: FA optimises for d = 64, 128, 256. Odd dimensions (d = 96, d = 192) fall into slow paths.
  • GQA/MQA with extreme ratios: natively supported, but the speedup versus pure MHA depends on the Q-heads : KV-heads ratio.
  • Cross-attention: supported but less optimised; self-attention is where the gain is largest.
  • FP8 without block quantisation or incoherent processing: it loses several points in benchmarks. If your serving framework does not implement the two tricks from the FA3 paper, FP8 attention can be a bad idea.

Implications on on-premise hardware

On an RTX 4090 (24 GB, Ada Lovelace, SM 8.9)

The 4090 is Ada Lovelace, not Hopper. It runs neither FA3 nor FA4; it runs FA2. That means around 70 % utilisation in BF16 attention (~250 effective TFLOPS against the 4090’s 330 TFLOPS BF16 peak). It is not a disaster, FA2 is already very good compared with standard attention, but the ceiling is clearly below that of an H100. For consumer deploys on a 4090 with Llama 3 8B BF16 or any 14B-32B INT4 AWQ, FA2 is what you will be using, and it is perfectly reasonable.

This is where FA3 shines and it is what vLLM/SGLang/TRT-LLM will select by default. Two common configurations:

  • Llama 3 70B FP8 with FA3 FP8 attention: 1.2 PFLOPS peak on the GPU, aggregate cluster throughput on the order of 8,000-12,000 tokens/s at medium batch depending on TP and context. For FP8 attention to deliver its full performance it is crucial to use the block quantisation + incoherent processing techniques from the FA3 paper (they are enabled by default in vLLM).
  • DeepSeek-V3 671B FP8 + MLA with FlashInfer: DeepSeek uses Multi-head Latent Attention (MLA), a different variant of standard attention. FlashInfer has specific kernels (FlashMLA). The typical stack is vLLM/SGLang + FlashInfer + FlashMLA + an FA3 fallback for the non-MLA layers.

If the infrastructure is Blackwell (B200/B300, which some clusters start receiving in 2026), FA4 is the right option and it should be enabled by default in vLLM 0.16+ and SGLang 0.5.11+.

What we have not covered

  • MLA (Multi-head Latent Attention) from DeepSeek and the specific FlashMLA kernels: they optimise KV cache compression but require different kernels.
  • Flexible masking and the FlexAttention use cases (PyTorch 2.5+): how to declare arbitrary masks without paying the cost of a custom kernel.
  • Hardware assistance for sparse attention (NVIDIA 2:4 sparse tensor cores) and why sparse attention has not established itself as a higher ceiling than dense FA.
  • FA in the backward pass of fine-tuning: this post focuses on inference, but FA3/FA4 also go through the backward pass and they are what makes training models with long contexts viable on H100/B200.

See also

References