First "Deep Dive" post — a step down from the Foundations series into one specific, widely-used piece of software. FlashAttention is a good subject for that: it's the attention implementation almost every modern LLM training and inference stack actually uses, and its three major versions (2022, 2023, 2024) are a clean case study in what "optimizing a kernel" actually looks like once the algorithm itself is already settled.

It's tempting to describe v2 and v3 as "the same idea, but faster." That's not really what happened. Each version fixed a specific, different bottleneck the previous one didn't address:

Timeline showing FlashAttention, FlashAttention-2, and FlashAttention-3, each labeled with the year, its core fix, and its headline result

The problem all three versions share

Standard attention computes softmax(QK^T)V by materializing the full N×N attention matrix — writing it to GPU high-bandwidth memory (HBM), then reading it back for the softmax and the second matrix multiply. For long sequences this is the dominant cost, not the matrix multiplies themselves: attention is memory-bandwidth-bound, not compute-bound. Every version of FlashAttention is attacking that same underlying fact from a different angle.

FlashAttention (2022): don't materialize the matrix at all

The original paper's contribution is IO-awareness: restructure the computation so the N×N intermediate matrix never touches HBM in the first place. Two ideas make that possible:

  • Tiling — process Q, K, V in blocks small enough to fit in on-chip SRAM.
  • Online softmax — since softmax normally needs the full row to compute a sum, FlashAttention keeps a running max and running sum per row and rescales as new blocks arrive, so the softmax can be computed incrementally without ever holding the whole row at once.
# one step of the streaming softmax used inside FlashAttention
m_new = max(m_running, max(scores_block))
correction = exp(m_running - m_new)
p = exp(scores_block - m_new)
l_new = correction * l_running + sum(p)
acc_new = correction * acc_running + p @ values_block

The result is exact attention (not an approximation) that's both faster and far more memory-efficient — O(N) memory instead of O(N²) — purely by changing where the intermediate results live, not what's being computed.

FlashAttention-2 (2023): stop wasting cycles on everything that isn't a matmul

FlashAttention-1 was IO-aware, but it still left GPU throughput on the table in ways that had nothing to do with HBM traffic. FlashAttention-2 targeted three separate inefficiencies:

  • Fewer non-matmul FLOPs — rebalanced the algorithm to spend less time on the rescaling arithmetic around each block, since non-matmul operations are far more expensive per-FLOP on a GPU than Tensor Core matrix multiplies.
  • Loop order swapped — the outer loop iterates over Q blocks and the inner loop over K/V blocks (reversed from v1), which reduces how often partial results need to be written back and re-read.
  • Parallelization over sequence length, not just batch size and attention heads — important because for a single long sequence with a small batch, batch × heads alone isn't enough work to keep every SM on the GPU busy.

The result: roughly 2x faster than the original FlashAttention, reaching 50-73% of theoretical peak FLOPs/s on an A100 — and up to 72% model FLOPs utilization in actual end-to-end training runs.

FlashAttention-3 (2024): stop waiting, period

FlashAttention-2's remaining inefficiency wasn't algorithmic — it was that computation and data movement were still happening sequentially: load a block, then compute on it, then load the next block. On Hopper, FlashAttention-3 overlaps the two, using hardware this series has already covered:

  • Warp specialization — some warps in a thread block ("producers") do nothing but issue asynchronous data movement via TMA, while other warps ("consumers") run the WGMMA matrix multiplies, pipelined against each other instead of run in sequence.
  • Interleaved matmul and softmax — block-wise matrix multiplies and softmax rescaling are scheduled to overlap rather than strictly alternate.
  • FP8 with incoherent processing — block-level quantization designed to keep FP8 accuracy competitive with full precision even when attention scores have outlier values, rather than a naive per-tensor quantization that degrades badly on outliers.

The result on H100: 1.5-2.0x faster than FlashAttention-2 in BF16 (up to 840 TFLOPs/s, ~85% utilization), and up to 1.3 PFLOPs/s using FP8 — a version that only makes sense on the hardware described in the GPU architecture history post and built with the exact tools covered in What is CUTLASS and CuTe Layouts: WGMMA, TMA, and warp-level pipelining aren't FlashAttention-3-specific inventions — they're Hopper primitives, and this paper is what using them for attention specifically looks like.

The pattern

v1 fixed a memory-traffic problem. v2 fixed a GPU-occupancy and instruction-mix problem. v3 fixed a sequencing problem — computation and data movement waiting on each other when the hardware can do both at once. None of these are the same bottleneck, which is exactly why the versions aren't interchangeable: FlashAttention-3's techniques don't help on pre-Hopper GPUs that lack TMA and WGMMA, and FlashAttention-2's loop reordering doesn't fix a problem that only shows up when data movement and compute are already overlapped. Each paper solved the bottleneck its predecessor's fix made visible next.

References

Timeline diagram above is original artwork made for this post.