Why FlashAttention?
Standard attention is memory-bound at the kernel level: the bottleneck is writing the big S (scores) and P (softmax) matrices to HBM and reading them back, not the matmuls. FlashAttention fuses the whole thing and keeps tiles in on-chip SRAM.
How
- Tiling: process Q/K/V in blocks that fit in SRAM; compute attention block-by-block.
- Online softmax: track a running max and sum so you never materialize the full row of scores: you rescale as you go.
- Recompute in backward: don’t store P; recompute it from stored stats. Trades cheap FLOPs for expensive memory.
Why it matters
- HBM reads/writes drop from to roughly → big wall-clock speedups, especially long context.
- Exact / lossless: unlike the sub-quadratic variants in Attention Optimization.
- Most useful on compute-bound work: LLM prefill and video generation (see Arithmetic Intensity).
Newer versions of FA
- FA1 (2022) introduced the IO-aware approach; FA2 improved work partitioning; FA3 targets Hopper (async, warp specialization, FP8); FA4 targets Blackwell.
Code implementation:
# to be added