Source
Dao et al. (2023), LIT-106 — ARXIV-2307.08691.
What the second version changed
Not the algorithm. FlashAttention-2 computes the same exact attention by the same tiled, never-materialised route (SOTA-087); what it rewrites is how the work is divided.
So it extends SOTA-085 as a later version of the same kernel, not a rival to it: the exactness that lets that practice be unconditional, and the condition in its title about hardware and compiled paths, carry over unchanged, and this practice only says which version to run.
Three changes, and the first is the largest: the inner loop is restructured so that far fewer of the operations are non-matmul. GPUs run matmul on tensor cores at many times the rate of everything else, so a kernel spending a noticeable fraction of its time on rescaling and bookkeeping is leaving most of the hardware idle even when it is fully occupied. The second is parallelising over the sequence-length dimension as well as batch and heads, which matters exactly when batch × heads is too small to fill the GPU — long context, small batch, the regime the kernel exists for. The third is a better split of work between warps inside a block, cutting shared-memory traffic.
The diagnosis, the changes and the recommendation are all LIT-106’s. It measured FlashAttention at only 25–40% of theoretical peak FLOPs/s on an A100 and traced the gap to work partitioning rather than the algorithm; the three changes together gave about 2× and 50–73% of peak. The record’s reading of the paper found no ablation separating them, so “the first is the largest” is an argument from the mechanism rather than a measured share.
So the recommendation is nearly free, with one condition
Same outputs, better occupancy: there is no accuracy argument to have and the upgrade is a kernel swap.
The condition is the usual one for fused kernels (SOTA-114) — the supported paths are specific, and a head dimension or mask type outside them falls back silently. The second version supports a wider set than the first, which is part of why the recommendation is unconditional in the title, and it is still not everything.