FlashAttention-2 — Deeper Math

Where FlashAttention-1 left performance on the table

FlashAttention-1 was a large win over a materialised attention matrix: it removed the O(N^2) memory blow-up and the HBM traffic that came with it. But measured against the machine rather than against the baseline, it still left most of the hardware idle.

On an A100 a well-tuned dense GEMM reaches roughly 80–90% of the theoretical FP16 tensor-core peak. FlashAttention-1’s forward pass reached about 124 TFLOP/s — near 40% of the 312 TFLOP/s peak — and the backward pass did worse. Attention is fundamentally two matmuls (QK^T and PV) with a softmax between them, so there is no structural reason it should run at half the efficiency of a GEMM of the same shape. Closing that gap is what FA-2 is about, and the diagnosis was that the loss sat in three separate places: wasted non-matmul arithmetic, idle SMs, and chatty warps.

The real enemy: non-matmul FLOPs

The single most important number here is a ratio. On an A100 the FP16/BF16 tensor cores deliver about 312 TFLOP/s, while ordinary FP32 non-tensor arithmetic — exponentials, comparisons, divisions, row reductions — runs at about 19.5 TFLOP/s. Tensor-core work is roughly 16× cheaper per FLOP than everything else.

The consequence is counter-intuitive: an operation that is a trivial share of the FLOP count can dominate the runtime. If non-matmul work is only 2% of the FLOPs, it still costs 0.02 × 16 = 0.32 — nearly a quarter of total time — against the 98% that is matmul. Attention has an irreducible amount of such work (the exp, the row max, the row sum), but FA-1 was doing meaningfully more than the minimum. FA-2’s first two changes simply delete non-matmul operations that need not be there.

Fix 1: defer the rescaling to the very end

Online softmax keeps a running max m, a running denominator l, and an output accumulator O. FA-1’s inner loop over key/value block j ran:

S  = Q_i K_j^T / sqrt(d)        # [B_r, B_c]
m' = max(m, rowmax(S));   P = exp(S - m')
l' = exp(m - m')*l + rowsum(P)
O' = diag(l')^-1 * ( diag(l)*exp(m - m')*O + P V_j )   <-- divides EVERY step

That diag(l')^-1 normalises the whole [B_r, d] accumulator on every inner iteration, though only the final value is ever read. FA-2 keeps an unnormalised accumulator instead:

O~' = diag(exp(m - m'))*O~ + P V_j     # no division
O   = diag(l_final)^-1 * O~_final      <-- divides ONCE, after the loop

With T_c = N / B_c inner iterations, one rescale replaces T_c of them, and the result is identical because division distributes over the accumulation.

Fix 2: store the logsumexp, not m and l

The backward pass must reconstruct the softmax probabilities from whatever the forward pass saved. FA-1 saved two tensors per row, the running max m and the running sum l. FA-2 saves one combined statistic, the logsumexp L = m + log(l), so that P_ij = exp(S_ij - L_i) yields the normalised probability directly.

Small change, two payoffs. It halves the softmax statistics written to and read back from HBM — real saving when the row count is batch × heads × N — and it collapses the backward pass’s ‘subtract the max, then divide by the sum’ into a single subtraction inside the exponent, removing another per-element division from the hot loop.

Fix 3: swap the loop order so queries are outermost

FA-1 put key/value blocks on the outer loop and query blocks on the inner loop. That is natural if you picture streaming K and V through SRAM, but it has an ugly consequence: a single output row block O_i is touched by every outer iteration, so its accumulator and statistics must repeatedly go out to HBM and come back, or be synchronised across thread blocks.

FA-2 inverts it: queries outermost, keys and values innermost. Each outer iteration now owns one row block O_i for its entire lifetime. The accumulator and the running m and l live in registers and SRAM from the first key block to the last, and reach HBM exactly once. Crucially, different query blocks never touch the same output memory, so they are fully independent — which is what makes the next fix possible.

Fix 4: parallelise over the sequence length

FA-1 launched one thread block per (batch, head) pair. That is ample parallelism at batch 64 with 32 heads. It is a disaster in exactly the regime long-context work lives in — take batch = 1, heads = 32, N = 8192:

FA-1 blocks = batch * heads = 1 * 32 = 32   vs   108 SMs on an A100
occupancy   = 32 / 108 ≈ 30%            → two thirds of the GPU idle

FA-2 adds a third axis: T_r = N / B_r = 8192 / 128 = 64 query blocks
FA-2 blocks = 1 * 32 * 64 = 2048            → fully occupied

Because Fix 3 made query blocks independent, this extra parallel axis is free: no atomics, no cross-block reduction, no extra passes. For short sequences with big batches it changes little, but for the long-sequence, small-batch shapes that motivated FlashAttention in the first place it is the difference between using 30% of the GPU and all of it.

Fix 5: split Q across warps, not K

Inside a thread block, work is divided among warps, and the two versions divide it differently. FA-1 used split-K: all warps hold the same Q_i tile and each takes a slice of K and V. Every warp then produces a partial contribution to the same output rows, so the partials go to shared memory, the block hits a barrier, and one warp reduces them.

FA-2 uses split-Q: K and V are shared by all warps and each warp takes a slice of the query rows. Each warp now computes a disjoint slice of O_i from start to finish — no partial sums, no shared-memory round trip, no __syncthreads() in the inner loop. It is the same principle as Fix 3 one level down the hierarchy: partition along the axis the output is indexed by, and communication disappears.

Causal masking: skip the blocks, do not mask them

With causal attention, query i may only attend to keys j ≤ i, so about half the score matrix is discarded. A naive kernel still computes those scores and then writes -∞ over them, paying full price for work it throws away.

Because FA-2 tiles the computation and knows each tile’s block coordinates, it classifies every tile up front. A block entirely above the diagonal is skipped outright — no QK^T, no exp, no PV. A block entirely below needs no masking logic at all. Only the O(N/B) tiles straddling the diagonal need per-element masks. For long causal sequences this approaches the ideal ~2× saving, and it feeds back into scheduling: query block i = 0 has one tile of work while block i = T_r - 1 has T_r, so the assignment of blocks to SMs must account for the imbalance.

A worked example: where the 2× actually comes from

It is tempting to credit the rescaling fix, but the arithmetic says otherwise. Take B_r = B_c = 128, d = 64, and count one inner iteration:

matmul (QK^T and PV) = 2 * (2 * B_r * B_c * d) ≈ 4.19 M FLOP
non-matmul floor     ≈ exp + rowmax + rowsum on [128,128] ≈ 49 K
FA-1 extra rescale   ≈ 2 * B_r * d                        ≈ 16 K

time ∝ matmul/312 + nonmatmul/19.5
FA-1: 0.01343 + 0.00333 = 0.01676     (non-matmul = 20% of the time)
FA-2: 0.01343 + 0.00251 = 0.01594     (−5%)

Fix 1 alone buys roughly 5% — real, but nowhere near 2×. The heavy lifting comes from the scheduling changes, because an idle SM contributes zero and raising occupancy from 30% to 100% is a multiplicative win rather than a percentage one. Together the changes took the A100 forward pass from ~124 to about 230 TFLOP/s (~73% of peak), with end-to-end GPT training reaching roughly 72% model FLOPs utilisation.

What FA-2 does not fix

Three limits are worth stating plainly. First, FA-2 does nothing for single-token decoding. Autoregressive generation issues one query against a long KV cache, so T_r = 1 and the sequence-parallel axis FA-2 added collapses. That workload is memory-bandwidth bound, not compute bound, and needs the opposite split — parallelise over keys and merge the partial softmaxes, which is what Flash-Decoding does.

Second, it is hardware-specific: the warp partitioning and register budgeting are tuned for Ampere-class tensor cores, head dimensions much above 128–256 spill registers and lose the benefit, and on Hopper the asynchronous copy engines and warp-group matmuls move the optimum again — hence FlashAttention-3. Third, it is still exact attention with O(N^2) compute. FA-2 makes the quadratic term cheap; it does not make it go away.

What transfers to CPUs and small models

None of the CUDA specifics apply to a CPU, but three underlying principles transfer almost verbatim, and they are worth internalising if you tune SLM inference on commodity hardware.

Keep the cheap units busy. The 312-vs-19.5 gap has a CPU analogue in AVX-512 or AMX matmul throughput versus scalar transcendentals; an exp or a division per element in an inner loop can cost more than the GEMM it decorates, which is why vectorised exp approximations matter so much. Defer reductions. Accumulate unnormalised and divide once — the trick works for any streaming softmax, on any hardware. Partition along the output axis. Give each thread a disjoint block of output rows and synchronisation vanishes; split along the reduction axis instead and you have bought yourself a merge step. FA-2 is, at heart, a well-executed lesson in all three.