05.07 · Concept
FlashAttention
Explain FlashAttention as an IO result rather than a maths one: tiling the computation so the full attention score matrix is never written to HBM, and what FA2 and FA3 changed on Ampere and Hopper.
FlashAttention is exact attention reorganised around memory traffic: compute attention in tiles, maintain softmax state online, and avoid writing the full score and probability matrices to HBM. FA2 mainly improved work partitioning for Ampere-style GPUs, while FA3 retargeted the inner loop to Hopper features such as TMA, WGMMA and FP8.
What this lesson answers
- why FlashAttention is an IO optimisation
- how FlashAttention avoids materialising attention scores
- what changed between FlashAttention FA2 and FA3
Notes
FlashAttention is an exact attention algorithm whose optimization target is HBM traffic, not the definition of attention: it computes while never materializing the score matrix or probability matrix in HBM. The mechanism is block streaming with an online softmax. For a query block and key/value block , the kernel forms in SRAM/registers, updates the running row maximum and normalization , and rescales the partial output .
References
Common questions
- Is FlashAttention an approximation of normal attention?
- No. FlashAttention computes the same dense attention result, but changes how the computation is scheduled. Instead of forming the whole score matrix in HBM, it streams blocks through on-chip memory and updates the softmax normalisation as it goes. The gain comes from avoiding expensive intermediate memory traffic, not from changing the attention formula.
- Why does FlashAttention help more during prefill than decode?
- Prefill has many query tokens, so naive attention can create a large quadratic score or probability intermediate. FlashAttention avoids that spill. Single-token decode mostly reads the existing KV cache for one new query, so the dominant cost is linear cache bandwidth and scheduling overhead, not writing a huge score matrix.
- What is the practical difference between FA2 and FA3?
- FA2 kept the same exact attention idea but improved parallelism, occupancy and tensor-core use, especially for Ampere-era hardware. FA3 is more Hopper-specific: it uses asynchronous data movement, warpgroup matrix instructions and related pipeline changes so memory movement, matrix work and softmax bookkeeping overlap more effectively.
Short definition: what is FlashAttention?