FlashAttention and the memory hierarchy
Exact attention can run several times faster, with memory linear in sequence length, by never writing the n × n matrix to slow memory. The trick is tiling plus a way to compute softmax one block at a time.
45 min
After this lesson you can
- Explain why standard attention is limited by memory traffic rather than arithmetic
- Derive the online softmax update (running maximum and running sum) and prove it gives the exact softmax
- Implement tiled attention in NumPy that matches naive attention to within
- Explain recomputation in the backward pass, and what FlashAttention-2 and -3 changed
In this lesson
- Two kinds of memory
- Tiling: never write the big matrix
- The online softmax
- The backward pass: recompute instead of remember
- FlashAttention-2 and -3
- Code: online softmax and tiled attention, checked against the naive version
This lesson is part of Descent Pro
Descent Pro unlocks every lesson on the route, from base camp to the summit. Every plan starts with a free trial.