TL;DR: Version 1 removed the N by N score matrix from HBM with tiling and online softmax, moving attention from memory-bound to compute-bound. Version 2 then attacked what was left: it cut non-matmul work by deferring the softmax division to the epilogue, added parallelism over the sequence dimension so long-context and small-batch shapes still fill the GPU, and split the query block across warps so warps stop exchanging partial results through shared memory. That took it from roughly a third of peak to roughly two thirds. Version 3 is Hopper-specific: producer warps issue TMA loads while consumer warpgroups issue wgmma instructions, and the softmax of one tile is scheduled to overlap the matmul of another, because exponentials run on units that are about sixty times slower per operation than the tensor cores. It reaches around 75% of bf16 peak and about 1.2 PFLOPS in FP8, as reported by the papers. Each version is the answer to the profile the previous one produced.
How to approach it
Frame it as three profiles rather than three papers: after each version, run the kernel and ask what is now the limit. Say what version 1 left on the table, then what version 2 left, then what version 3 does about it. Do the non-matmul arithmetic out loud, because that single number explains both version 2 and version 3. Close with what is portable and what is Hopper-only.
A strong answer
A typical situation: a team upgrades from A100 to H100, keeps their FlashAttention 2 kernel, and sees attention throughput rise by much less than the 3x the spec sheet promised. The kernel is correct and well written; it is simply not using the two Hopper features that the newer generation of attention kernels is built around.
The number that explains most of the story is the cost of a non-matmul operation. Tensor cores and the units that compute exponentials are not in the same league:
H100 SXM, bf16 matmul on tensor cores: 989 TFLOPS
non-matmul math (exp, reciprocal) on the SFUs: roughly 16 TFLOP-equivalents per second
ratio: a non-matmul operation costs about 60x what a matmul operation costs
attention per head at sequence N, head dim d = 128:
matmul FLOPs = 4 x N^2 x d (QKt and PV, two FLOPs per multiply-add)
softmax ops ~ 5 x N^2 (subtract max, exp, add to sum, scale, divide)
FLOP ratio = 4d / 5 = 102, so softmax is about 1% of the arithmetic
time ratio = (5 N^2 / 16e12) / (4 N^2 d / 989e12) = 0.31 x 1.93 = 0.60
sanity: 1% of the operations, 60% of the time of the matmuls if the two do not overlap.
That is why version 2 removes softmax work and version 3 hides what is left.
The three versions, and the bottleneck each attacked:
| What it fixed | The mechanism | Reported result | |
|---|---|---|---|
| FlashAttention (2022) | HBM traffic for the N by N scores | tiling plus online softmax, scores never leave SRAM | 2 to 4x over a materializing kernel; about 25 to 40% of A100 peak |
| FlashAttention 2 (2023) | non-matmul work, idle SMs, warp chatter | keep the accumulator unnormalized and divide by the row sum once in the epilogue; parallelize over sequence blocks as well as batch and heads; give each warp its own slice of the query block instead of splitting keys | up to about 70% of A100 bf16 peak |
| FlashAttention 3 (2024) | tensor cores idle during softmax, and loads on the critical path | warp specialization: producer warps issue TMA copies while consumer warpgroups issue wgmma; two-stage scheduling so the softmax of tile j overlaps the matmul of tile j+1; FP8 with block quantization | about 75% of H100 bf16 peak, roughly 1.2 PFLOPS in FP8 |
The three version 2 changes are worth separating, because interviewers ask which one mattered. Deferring the division is arithmetic removal: the accumulator stays unnormalized through the loop and is divided by the row sum once at the end, so an N by N division becomes an N by d division. The sequence-dimension parallelism matters at the shape teams care about now: batch 1 with 8 heads gives the older scheme 8 independent work units, leaving most of an H100's 132 SMs idle, while splitting the query dimension into blocks creates hundreds. The warp partitioning removes a shared-memory round trip: splitting keys leaves every warp with a partial row sum to reduce, while splitting queries gives each warp whole rows and nothing to exchange.
Version 3 is a scheduling change more than an algorithmic one. Tensor Cores and Matrix Units explains why the tensor cores need operands delivered continuously to hit their rate, and CUTLASS and Tensor Core Kernels is where the producer and consumer pattern comes from. Hopper adds two hardware pieces that make it expressible: TMA, which moves a whole tile with one instruction and no per-element address arithmetic, and wgmma, which issues a matmul across a warpgroup of four warps with operands read directly from shared memory. Splitting warps into producers that only load and consumers that only compute means the loads for the next tile are already in flight while the current tile is being multiplied, and the pingpong schedule puts one warpgroup's softmax next to another warpgroup's GEMM so the 60% figure above becomes overlap rather than added time.
FP8 in version 3 is not simply a cast. Attention logits carry outliers that a per-tensor scale handles badly, so the kernel quantizes in blocks and first applies a random orthogonal transform that spreads outlier energy across dimensions. Numerics: FP32, BF16, FP8 and FP4 covers why the accumulator stays in fp32 regardless.
Knowing which kernel your stack dispatched to is the practical half. Divide the attention FLOPs (4 times N squared times d per head, summed over heads and layers) by the measured attention time from an Nsight Systems timeline and compare against the 989 TFLOPS peak: near 250 TFLOPS behaves like version 1, near 600 like version 2, near 740 like version 3. Compute that before anyone argues about which library version is installed.
The reversal condition: version 3 is Hopper-only in the parts that matter. TMA, wgmma and warpgroup scheduling do not exist on Ampere, so on an A100 the best available kernel is still version 2 and the answer to "should we upgrade the kernel" is no. On Blackwell the same reasoning repeats with a new instruction set, which is the real lesson: these kernels are re-derived per architecture, and a team that treats attention as a solved library call will be one generation behind on every migration.
What interviewers probe next
- "Which version 2 change gave the most?" It depends on shape, and say so with the shape: at large batch the non-matmul removal dominates, at batch 1 with long context the sequence-dimension parallelism does, because the old scheme simply could not fill the SMs.
- "What is warp specialization actually doing?" Dedicating some warps to memory movement and others to math, so the two proceed at once instead of every warp alternating between them and stalling on its own loads.
- "Why does FP8 need more than a cast here?" Attention logits have heavy outliers; per-tensor scaling either clips them or wastes range on everything else, so version 3 quantizes per block and decorrelates first.
- "Is FlashAttention 3 always faster than 2 on an H100?" For long sequences yes, by a wide margin. At very short sequences the tile pipeline never fills and the two converge.
Common mistakes
- Describing the versions as "the same idea, more optimized", which misses that each targeted a different limiter.
- Thinking softmax is negligible because it is 1% of the FLOPs.
- Claiming version 3 speedups on an A100, where its two key instructions do not exist.
- Forgetting that version 2's parallelism change is what makes small-batch long-context serving viable.
Key takeaways
- A non-matmul operation costs roughly 60x a matmul operation on H100, so softmax at 1% of the FLOPs is 60% of the matmul time unless it is removed or overlapped.
- Version 1 fixed HBM traffic, version 2 fixed non-matmul work and SM occupancy, version 3 fixed the serialization between softmax and matmul.
- Version 2's three changes: unnormalized accumulator with one epilogue division, parallelism over sequence blocks, queries split across warps rather than keys.
- Version 3 needs TMA and wgmma, so it is Hopper and later; on Ampere version 2 remains the kernel.