FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré2022NeurIPS 2022
Read it on arxiv.org(opens in a new tab)Why this one
Read this after Attention Is All You Need, then read FlashAttention-2 when you want the production kernel story. The useful idea is not a new kind of attention, it is a refusal to pretend FLOPs are the whole cost model. Standard attention materializes a huge score matrix, moves it through slow GPU memory, then throws most of the intermediate structure away. FlashAttention keeps the math exact but changes the schedule: tile the work, stream through blocks, and keep the hot state in SRAM. What people get wrong is treating it as a clever optimization trick. It is a systems paper wearing ML clothes, and the lesson transfers to every inference bottleneck: if your algorithm ignores the memory hierarchy, the hardware will grade you anyway. Worth your evening because it gives you a concrete mental model for why kernels, layouts, and IO decide model speed.
What to take away
- Exact attention can get faster by changing data movement, not the formula.
- The avoided cost is materializing and rereading the full attention matrix in HBM.
- Tiling works because softmax can be computed online while preserving numerical correctness.
Reads with
- Attention Is All You Need
defines the attention math whose memory traffic flashattention reschedules
- FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
turns the io aware idea into a faster production grade kernel
- Ring Attention with Blockwise Transformers for Near-Infinite Context
extends blockwise exact attention scheduling across devices for long context