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.