FlashAttention-2: Where the Second 2× Actually Comes From

Core insights. FlashAttention-1 already removed the O(N²) memory traffic; it was still leaving roughly half the GPU’s compute on the table. FlashAttention-2’s gains come from three changes, none of which alter the math of attention: (1) deferring softmax rescaling so the inner loop is almost pure matmul, (2) parallelizing across the sequence dimension so long-sequence/small-batch workloads fill the GPU, and (3) splitting work across warps by Q instead of K/V, which removes shared-memory round-trips from the inner loop. The lesson generalizes: once an algorithm is IO-optimal, the next bottleneck is how work maps onto the hardware’s compute hierarchy.

Method

Standard attention materializes the score matrix $S = QK^\top \in \mathbb{R}^{N \times N}$ in HBM. FlashAttention tiles $Q, K, V$ into blocks that fit in SRAM and computes softmax online: as each block of keys is processed, a running row-max $m$ and running normalizer $\ell$ are maintained, and the partial output is rescaled so the final result is exact, not approximate.

For query block $i$ , iterating over key blocks $j$ :

m^{(j)} = \max\!\left(m^{(j-1)},\ \mathrm{rowmax}\!\left(S^{(j)}\right)\right)

\ell^{(j)} = e^{\,m^{(j-1)} - m^{(j)}}\,\ell^{(j-1)} + \mathrm{rowsum}\!\left(e^{\,S^{(j)} - m^{(j)}}\right)

\tilde{O}^{(j)} = \mathrm{diag}\!\left(e^{\,m^{(j-1)} - m^{(j)}}\right)\tilde{O}^{(j-1)} + e^{\,S^{(j)} - m^{(j)}}\,V_j

FlashAttention-1 rescaled $\tilde O$ by $\mathrm{diag}(\ell)^{-1}$ inside the loop. FlashAttention-2’s first change is to keep the output unnormalized through the whole loop and apply

O = \mathrm{diag}\!\left(\ell^{(\text{last})}\right)^{-1} \tilde{O}^{(\text{last})}

once at the end. Division and exponentials run on the GPU’s non-matmul units, which have ~16× lower throughput than tensor cores on an A100; every non-matmul FLOP removed from the inner loop counts. Only the logsumexp $L = m + \log \ell$ is stored for the backward pass, instead of both $m$ and $\ell$ .

Parallelism over sequence length

FlashAttention-1 assigned one thread block per (batch × head). With batch 1 and long context — exactly the inference-relevant regime — that leaves most of the GPU’s SMs idle. FlashAttention-2 additionally parallelizes the forward pass across query blocks (and the backward pass across key blocks), so occupancy stays high even at batch 1.

Warp partitioning: split-Q, not split-K

Within a thread block, FA-1 split $K$ and $V$ across 4–8 warps (“split-K”), which forces each warp to write partial results to shared memory, synchronize, and re-read for the final reduction. FA-2 splits $Q$ across warps instead: each warp owns a slice of the output rows outright, and no inter-warp communication is needed in the inner loop.

Claims & evidence

Claim	Evidence in paper	Verdict
~2× faster than FlashAttention-1	Fwd+bwd benchmarks on A100, head dims 64/128, seq 512–16k (§4.1, Fig. 4–7)	verified
Reaches 50–73% of A100 peak FLOPs/s in the forward pass	Same benchmark suite; peak ≈ 312 TFLOPs/s FP16	verified
Up to 225 TFLOPs/s per A100 (72% MFU) in end-to-end GPT training	GPT-style models 1.3B/2.7B, seq 2k/8k, 8×A100 (§4.2)	partial — end-to-end MFU depends heavily on the rest of the training stack; the number replicates with their Megatron config, not universally
Benefits carry to H100	Brief mention; no Hopper-specific tuning (no TMA/WGMMA use) in the paper	partial — later FlashAttention-3 (2024) shows Hopper needed a redesign to get there

Benchmarks

Forward+backward, A100 80GB SXM, FP16, from the paper’s Fig. 4–6 (representative points, seq len 2k):

Implementation	Head dim 64, no mask	Head dim 128, causal
PyTorch standard attention	~25 TFLOPs/s	~30 TFLOPs/s
FlashAttention-1	~110 TFLOPs/s	~85 TFLOPs/s
Triton FlashAttention	~95 TFLOPs/s	—
FlashAttention-2	~190 TFLOPs/s	~160 TFLOPs/s

Numbers read from published plots; treat as ±10%. The consistent pattern — FA-2 ≈ 2× FA-1, larger gaps at long sequence and causal masking — is what matters, and it replicates.

Limitations & open questions

The paper targets Ampere. On Hopper (H100), the same ideas required a third redesign (FlashAttention-3: warp specialization, TMA, FP8) to approach peak — work partitioning is architecture-specific, and each generation reopens the problem.
Head dimensions are capped at 256, and irregular masks (beyond causal) fall back to slower paths.
The 72% end-to-end MFU claim bundles attention improvements with a tuned Megatron stack; attention alone does not determine MFU at that level.

Reproduction notes

Reproduced the core speedup claim with flash-attn 2.x on A100 80GB: forward+backward at seq 4096, head dim 128, batch 8 lands within 10% of the paper’s curve, and ~1.9× over flash-attn 1.0.9 on the same shapes. The benchmarks/benchmark_flash_attention.py script in the repo reproduces the paper’s plots directly. One footgun: comparing against PyTorch scaled_dot_product_attention with the memory-efficient backend (not the math backend) shrinks the headline gap — always check which backend SDPA dispatched to.