Flash Attention: Why the NxN Matrix Is a Memory Problem, Not a Compute One

S1 The bus, not the math: Tensor Cores only compute on data in on-chip SRAM; HBM is slow off-chip memory. The NxN attention score matrix is too big for SRAM, so it lives in HBM. The bottleneck is memory bandwidth, not FLOPs. S2 The four round-trips, named: 1. Load Q, K, V from HBM into SRAM. 2. Compute S = QK^T in SRAM, write the full NxN back to HBM. 3. Read S back from HBM for mask + softmax. 4. Write P = softmax(S) to HBM, read P and V back for PV. The NxN matrix crosses the HBM/SRAM boundary multiple times. S3 The Flash move: tile Q, K, V into SRAM-sized blocks. Keep each score tile on-chip. Online softmax: running max + running sum of exp per row, rescale as tiles stream. The full matrix never touches HBM. Memory O(N^2) to O(N). S4 Falsify the slide: the 2-4x speedup is NOT fewer FLOPs - same compute, no HBM bandwidth tax. The running-stats rescaling is exact, not approximate. Causal mask applied per-tile.