flash attention

FlashAttention is an IO-aware attention algorithm that computes exact attention while never writing the full N×N score matrix to HBM. Instead of materializing all pairwise scores, it processes attention in tiles that stay in fast on-chip SRAM, folding each tile into a running softmax. The result is identical to standard attention, but with dramatically less memory traffic — which is what actually limits attention on modern accelerators.\n\n**Standard attention is memory-bound, not compute-bound.** The naive implementation forms S = QKᵀ, an N×N matrix for sequence length N, writes it to HBM, reads it back to apply softmax, writes the probabilities, then reads them again to multiply by V. The multiply-adds are cheap; the killer is shuttling those large intermediate matrices in and out of high-bandwidth memory. Traffic scales as O(N²), so as context length grows, attention spends nearly all its time moving data rather than computing.\n\n**Tiling plus an online softmax removes the round trip.** FlashAttention loads blocks of Q, K, and V into SRAM, computes each block's partial scores there, and combines them using an online softmax that keeps a running maximum and normalizer — so it never needs the whole row of scores at once. Because each tile is consumed on-chip and discarded, the O(N²) matrix is never written to HBM. The math is exact (not an approximation), but HBM traffic falls to roughly O(N), turning a memory-bound kernel into a far faster one.\n\n| Aspect | Standard attention | FlashAttention |\n|---|---|---|\n| N×N scores | written to HBM | kept in SRAM tiles |\n| HBM traffic | O(N²) | O(N) |\n| Softmax | full-row, materialized | online / running |\n| Result | exact | exact (identical) |\n| Bottleneck | memory bandwidth | much closer to compute |\n\n```svg\n\n```\n\n**It reshaped how long-context models are trained and served.** By cutting attention's memory traffic and footprint, FlashAttention made longer sequences practical without the quadratic memory blowup, and later versions (FlashAttention-2 and -3) pushed GPU utilization higher by improving work partitioning and exploiting newer hardware. It pairs naturally with the KV cache at inference: the cache limits how much context you can hold, while FlashAttention limits how expensive it is to attend over that context. Both are fundamentally about respecting the memory hierarchy.\n\nRead FlashAttention through a quant lens rather than a 'clever kernel' lens: attention's cost is bytes moved through HBM, and per the roofline a memory-bound kernel is capped by bandwidth long before FLOPs. FlashAttention lowers arithmetic intensity's denominator — HBM bytes — by keeping the N×N scores on-chip, so the same math lands higher on the roofline. The design question is how large a tile fits in SRAM and how few HBM bytes per token attention must touch, a measured traffic budget rather than a FLOP count.