FlashAttention-2 paper

The question the second paper asks

FlashAttention made attention’s memory linear, and its speedups were measured against baselines that were themselves inefficient. FlashAttention-2 begins from a measurement: the kernel reached only 25 to 40 percent of the A100’s theoretical peak FLOP/s, while a well-tuned GEMM gets much closer. The algorithm was unchanged, so the loss had to be in how the work was divided.

Three changes

The paper names them precisely. First, non-matmul work: rescaling the accumulator after every tile costs time on units that do not run tensor-core matmuls, so the schedule delays and batches the rescaling. Second, parallelism: the first version split work across batch and heads only, which leaves the GPU idle at small batch, so the query dimension is split across thread blocks even within a single head, which also removes the need for a split-K reduction. Third, within a block, work is partitioned across warps along output rows so that each warp’s accumulator stays in registers, cutting shared-memory traffic.

The result is about twice the first version, 50 to 73 percent of peak on A100, and up to 225 TFLOP/s per GPU when training GPT-style models.

What changed and what did not

The mathematics did not change: tiling with online softmax, exact, same flop count. What changed is the mapping from that algorithm onto blocks and warps. The paper is the clearest published case that on modern accelerators the schedule, not the algorithm, is the remaining lever.

What to read

Read the three subsections of the algorithm section against those numbers, and skim the benchmark tables. Skip the implementation notes tied to a specific CUTLASS version on a first pass.

Afterwards, open a Triton fused-attention kernel and identify which dimension is split across programs, where the rescale happens, and how many warps cooperate on one output tile.

Things that break

Reporting “FlashAttention-2 is twice as fast as FlashAttention” without naming shape, dtype and GPU. The gain is largest where occupancy was the problem, so a benchmark at batch one, short sequence, or on a part where the first version already filled the SMs shows much less. The tell is a speedup that disappears when you rerun at a different batch size.