GoKawiil0. Introduction
How IO-Aware Attention Makes Transformers Faster Without Approximating Attention
The mechanism, in three words: Tiling + Online Softmax + Recomputation. Everything in this handbook is elaboration on that summary.
A technical handbook on exact tiled attention: GPU memory traffic, online softmax, forward and backward passes, IO complexity, the evolution from FlashAttention-1 through FlashAttention-4, and current framework behavior.
0.1 How to Read This Handbook
This handbook was inspired by this tweet. Before the fix, here is what the standard attention implementation looks like. Load $Q, K, V \in \mathbb{R}^{N \times d}$ in HBM, then:
Read $Q, K$ from HBM, compute $S$, write $S$ to HBM. Read $S$ from HBM, compute $P$, write $P$ to HBM. Read $P, V$ by blocks from HBM, compute $O$, write $O$ to HBM. Return $O$.
What stands out to me is the number of round trips to HBM. Every intermediate value — $S$, $P$, $O$ — has to be written out and read back. That is the problem FlashAttention is solving.
The handbook itself frames the subject as easiest to understand when three different questions are kept separate:
What mathematical function is being computed? For dense attention, the target remains ordinary scaled dot-product attention. How much arithmetic does that function require? Dense all-pairs query-key scoring remains quadratic in sequence length. How does the implementation move data through the GPU memory hierarchy? This is where FlashAttention changes the algorithmic execution dramatically.
... continue reading