Understanding FlashAttention Pt 1: Personal Notes

0. 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. The central lesson I take from this framing is that wall-clock speed is not determined by FLOP count alone. An algorithm can perform essentially the same mathematical work, or even recompute intermediate values, and still run faster because it moves far less data to and from high-bandwidth memory. Core distinction. Dense FlashAttention is an exact attention algorithm: it does not replace softmax attention with a low-rank, sparse, kernelized, or approximate formula. “Exact” refers to the mathematical attention computation. Floating-point kernels can still differ by small rounding effects because operations are reordered. The word exact is doing real work here. Exactness is a statement about the mathematical function, not about bitwise reproducibility. The kernel is free to reorder floating-point operations. It is not free to change the function being computed. 0.2 Notation For one attention head, let \[Q \in \mathbb{R}^{N_q \times d}, \quad K \in \mathbb{R}^{N_k \times d}, \quad V \in \mathbb{R}^{N_k \times d_v}.\] For self-attention, typically $N_q = N_k = N$. The scaled score matrix is \[S = \frac{QK^T}{\sqrt{d}} + B,\] where $B$ represents an optional additive mask or bias, and \[P = \mathrm{softmax}\_{\mathrm{row}}(S), \quad O = PV.\] Throughout, HBM refers to large off-chip high-bandwidth GPU memory. On-chip memory is a broad teaching term for much smaller, faster storage such as registers and shared memory/SRAM. Exact hardware details vary by GPU generation. The practical difference I keep coming back to: HBM SRAM slow faster large smaller off-chip on-chip System problem: where do all those intermediate values $(S, P, O)$ live while the GPU computes them? That question — not the arithmetic — is what FlashAttention was built to answer. Contents Part 1: The Fundamental Problem 1.1 What FlashAttention actually optimizes 1.2 The attention equation is not the implementation 1.3 GPU memory hierarchy and why IO matters 1.4 Why materializing $S$ and $P$ is expensive 1.5 Dense arithmetic is still quadratic 1.6 Memory-efficient exact attention predates FlashAttention Part 2: The Mathematical Trick 2.1 Tiling queries, keys, and values 2.2 Softmax is the difficult part of streaming 2.3 Online softmax from first principles 2.4 The blockwise merge recurrence 2.5 A complete numerical example Part 3: Putting the Mathematics onto the GPU 3.1 FlashAttention forward pass 3.2 Why dense FlashAttention is exact 3.3 IO complexity: what the theorem actually says 3.4 Memory complexity: linear auxiliary state, not linear compute 3.5 Causal masking and tile skipping 3.6 Backward pass: recompute instead of save 3.7 Why more FLOPs can still be faster Part 4: Architectural Compatibility 4.1 MHA, MQA, and GQA compatibility 4.2 Variable lengths, local attention, and dropout Part 5: FlashAttention Evolution 5.1 FlashAttention-2: what changed 5.2 FA2 parallelism across sequence tiles 5.3 FA2 work partitioning and non-matmul FLOPs 5.4 FlashAttention-3: the Hopper generation 5.5 FA3 asynchrony: overlap data movement, GEMM, and softmax 5.6 FA3 FP8: performance without pretending precision is free 5.7 FlashAttention-4: the Blackwell generation 5.8 FA4 and asymmetric hardware scaling 5.9 FA4 implementation and current status 5.10 FlashAttention-1 through -4 compared Part 6: Using FlashAttention in Frameworks 6.1 PyTorch scaled-dot-product attention today 6.2 Exactness is not bitwise identity Part 7: FlashAttention vs. Other Techniques 7.1 FlashAttention vs. PagedAttention 7.2 FlashAttention vs. sparse and linear attention Part 8: Training vs. Inference 8.1 Training, prefill, and decode are different regimes Part 9: Practical Engineering 9.1 Common implementation mistakes 9.2 Common misconceptions and the practical mental model Part 1: The Fundamental Problem 1.1 What FlashAttention Actually Optimizes Start with the ordinary attention function: \[O = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}} + B\right) V.\] FlashAttention is IO-aware. My working definition: Minimize data movement between the different levels of GPU memory, rather than just trying to reduce the number of mathematical operations (FLOPs). The speed bottleneck in modern AI hardware is often not how fast the GPU can compute math, but how fast it can read and write data. This is the memory-compute tradeoff. A textbook implementation often makes this look like three large operations: Form the score matrix $S$, apply row-wise softmax to obtain $P$, then multiply by $V$. Mathematically that is fine. On a GPU, however, writing a huge intermediate matrix to HBM and reading it back can be far more expensive than the equation suggests. FlashAttention’s central contribution is to make the algorithm IO-aware. It partitions the computation into tiles that fit in fast on-chip memory, streams blocks of $K$ and $V$, and maintains enough row-wise softmax state to produce the exact output without materializing the full $N \times N$ attention matrix in HBM. What changes: the execution schedule, memory traffic, and stored intermediates. What does not change: the dense scaled-dot-product attention function being evaluated. This distinction is why “FlashAttention is a faster kind of attention” can be misleading. It is better thought of as an algorithm and kernel family for evaluating attention efficiently on accelerators. A model can use causal masking, RoPE, MQA/GQA, or other attention features and still use a FlashAttention implementation underneath. The original paper contrasts this approach with approximate attention methods that reduce arithmetic by changing the mathematical problem. Dense FlashAttention does not make that trade. The same paper also introduced a block-sparse extension, but that sparse extension is a different case because omitting blocks changes which interactions are computed. 1.2 The Attention Equation Is Not the Implementation The equation does not tell you where tensors live. A standard naive attention implementation: Calculate $S$, store $S$. Read $S$, calculate $P$, store $P$. Read $P$, calculate $O$. The problem is the number of HBM round trips. A simple materializing implementation does: \[S \leftarrow QK^T / \sqrt{d}, \quad P \leftarrow \mathrm{softmax}(S), \quad O \leftarrow PV.\] If $S$ is written to HBM after the first matrix multiplication, read for softmax, $P$ is written back, and then $P$ is read again for the $PV$ multiplication, the GPU spends significant time moving an $N^2$ object through memory. FlashAttention avoids repeatedly moving huge intermediates: Calculate small $S$ tile, softmax tile, use tile with $V_j$, discard tile. Calculate next tile. The key questions I ask when looking at any equation: How many operations are required? What data must move between memory levels to perform those operations? FlashAttention reuses the same dependencies. Blocks of $Q$, $K$, and $V$ are brought near the compute units, score tiles are produced and consumed locally, and only compact row-wise statistics plus the output need to persist across tiles. Algorithmic lesson. A computational graph is not a memory schedule. Writing $P = \mathrm{softmax}(QK^T)$ on paper does not require an implementation to store all of $QK^T$ or $P$ in off-chip memory at once. This idea generalizes beyond attention. Fused kernels, tiling, recomputation, and operator scheduling often trade a small amount of extra arithmetic for much less movement of large intermediates. On modern accelerators, that can be the right trade because matrix-multiply throughput has grown much faster than many other parts of the memory and execution hierarchy. FlashAttention-3 and -4 make that hardware dependence increasingly explicit. 1.3 GPU Memory Hierarchy and Why IO Matters GPUs expose a hierarchy rather than one uniform pool of equally fast memory. The names and capacities vary by architecture, but the mental model is: Registers / very local state On-chip shared memory / SRAM HBM / device memory HBM is large, but data must travel to the compute units. On-chip storage is much smaller, but reuse there is much cheaper. The original FlashAttention analysis models this asymmetry using HBM and SRAM and explicitly optimizes the number of transfers between them. My breakdown of each level: HBM: Relatively large Relatively slower to access Physically farther from individual compute operations Stores model weights, $Q/K/V$, activations, large tensors SRAM / shared memory: Much smaller Faster Cheaper reuse cost