FlashAttention is not a new attention formula — it computes exactly the same softmax attention as the textbook version, to the last bit. What it changes is where the numbers live while you compute them. Standard attention writes a giant N-by-N score matrix out to slow GPU memory and reads it back; FlashAttention keeps the computation in fast on-chip memory by processing the matrix in tiles and never storing it whole. That single IO-aware idea — count the memory traffic, not just the FLOPs — makes attention several times faster and cuts its memory footprint from O(N^2) to O(N), which is what makes long context affordable. This piece builds the intuition from first principles; the companion ‘Flash Attention Math Deep Dive’ carries the full derivation and the backward-pass gradients.
The one idea: attention is memory-bound, not compute-bound
The usual way to reason about a kernel is to count floating-point operations. For attention that instinct misleads you. On a modern GPU the arithmetic units are so fast that the bottleneck is almost never the multiplies — it is moving data between the large, slow off-chip memory and the tiny, fast on-chip memory.
FlashAttention’s founding observation (Dao et al., 2022) is that standard attention spends most of its wall-clock time reading and writing a large intermediate matrix, not computing it. Fix the memory traffic and the kernel speeds up even though the FLOP count is unchanged. This is what IO-aware means: design around the memory hierarchy, minimizing bytes moved to and from slow memory, not around the operation count. Everything else follows from taking that one goal seriously.
The memory hierarchy: HBM versus SRAM
A GPU has two kinds of memory that matter here. HBM (high-bandwidth memory) is the large pool — tens of gigabytes — where your tensors normally live. It is fast in absolute terms (an H100 moves roughly 3 TB/s) but it is the slow tier relative to the chip. SRAM is the on-chip scratchpad attached to each streaming multiprocessor — only about 100–200 KB per SM, but an order of magnitude faster and far lower latency.
The gap is the whole story. Reading from SRAM is cheap; reading from HBM is the expensive event you want to avoid. A kernel is memory-bound when it stalls waiting on HBM while the compute units sit idle. FlashAttention’s design question is blunt: how few times can we touch HBM and still get the exact right answer?
Why standard attention wastes memory bandwidth
Standard attention, for one head, takes Q, K, V of shape [N, d] and computes S = QK^T (shape [N, N]), then P = softmax(S) row-wise, then O = PV (shape [N, d]). The trouble is S: it is quadratic in sequence length and it is written to HBM, read back for the softmax, and read again for the PV multiply.
For N = 8192 and fp16, a single N-by-N matrix is about 128 MB — per head, per layer. Materializing it means O(N^2) HBM reads and writes, and that traffic, not the matmuls, dominates the runtime. It also costs O(N^2) memory just to exist, which is why naive attention runs out of memory long before the model runs out of useful context.
The fix: tile the computation
FlashAttention refuses to ever hold S whole. It splits Q into row blocks and K, V into column blocks — say 128 each — sized so a Q block and a K/V block fit together in SRAM. The kernel loops: load one Q block, stream the K and V blocks past it, and for each pair compute the small block of scores on-chip.
Because each score block is formed, used, and discarded inside SRAM, the full N-by-N matrix is never written to HBM at all. HBM sees only the inputs Q, K, V and the final output O — all of size O(N·d). The tiling is the mechanical half of the idea. The subtle half is how you take a correct softmax over a whole row when you only ever see one block of that row at a time.
The problem tiling creates: softmax needs the whole row
Softmax is not a per-element function — it normalizes across an entire row: softmax(x_i) = exp(x_i) / Σ_j exp(x_j). To divide by that denominator you seemingly need every score in the row at once, which is precisely what tiling denies you. Worse, in practice softmax is computed in a numerically stable form by subtracting the row max first: exp(x_i - m) / Σ_j exp(x_j - m) with m = max_j x_j, and you do not know that max until you have seen the whole row either.
So tiling alone breaks the softmax. FlashAttention resolves this with online softmax: a way to compute the row’s normalized result incrementally, updating a running answer block by block. The exact algebra is the sibling article’s territory; here we build the intuition for why it works.
Online softmax: the running-max, running-sum idea
The trick is to carry two small running statistics per Q row as you stream the K/V blocks: the maximum score seen so far, m, and the running sum of exponentials l = Σ exp(score - m). When a new block arrives with a larger maximum, the old m is stale — every exponential you accumulated was scaled by the wrong offset.
The repair is a single rescale. If the max rises from m to m', you multiply the accumulated sum and output by exp(m - m'), which corrects every earlier term at once, then add the new block. Because softmax is invariant to subtracting a constant from all its inputs, this yields exactly the same numbers as a one-shot softmax — no approximation, only a reordering of arithmetic.
Rescale and accumulate: the forward pass in one loop
Put the pieces together and the forward pass is a tidy nested loop. For each Q block, initialize a running max m = -∞, a running denominator l = 0, and an output accumulator O = 0. Then, for each K/V block:
S_ij = Q_i · K_j^T # small block, in SRAM
m_new = max(m, rowmax(S_ij))
P_ij = exp(S_ij - m_new) # rescaled probabilities
l = exp(m - m_new)*l + rowsum(P_ij)
O = exp(m - m_new)*O + P_ij · V_j
m = m_newAfter the last K/V block, divide once: O_i = O / l, and write that row block out to HBM. Every quantity in the loop is a small tile living in SRAM; the only HBM writes are the final outputs.
O(N) memory: the footprint that unlocks long context
Because the N-by-N scores are never stored, the memory FlashAttention needs beyond its inputs and output is just the per-row running statistics — a handful of numbers per query row, i.e. O(N), not O(N^2). That is the difference between a context length that fits and one that does not.
Concretely, the quadratic 128 MB score matrix for N = 8192 evaporates; the working set is a few tiles of a few tens of KB. Quadrupling the context still quadruples the compute — attention is inherently O(N^2) in FLOPs and FlashAttention does not change that — but the memory stays linear, which is why 128k- and million-token contexts became practical on hardware that could never hold their score matrices.
The backward pass: recompute instead of store
Training needs gradients, and gradients normally need the attention probabilities P that the forward pass produced. Storing P would drag O(N^2) memory right back in — defeating the point. FlashAttention’s answer is recomputation: the forward pass saves only the cheap O(N) softmax statistics (the per-row m and l), and the backward pass regenerates each P tile on the fly from Q, K, and those saved statistics.
This trades extra compute for far less memory traffic — gradient checkpointing specialized to attention. Because the kernel was memory-bound, redoing arithmetic on-chip is nearly free relative to the HBM round trips it avoids. The exact gradient formulas are derived in the math companion; the architectural point is that the backward pass, too, never materializes the full matrix.
A tiny worked example
Take one query row with four keys, split into two blocks. Scores are block A = [1, 3] and block B = [2, 5]. Process A first: running max m = 3, and l = exp(1-3) + exp(3-3) = 1.135.
Block B arrives with a larger value, 5, so m' = 5. Rescale the old sum by exp(3-5) = 0.135, then add block B’s terms exp(2-5) + exp(5-5) = 1.050, giving l = 0.135×1.135 + 1.050 = 1.203. The plain one-shot denominator over all four scores (max 5) is also 1.203 — identical. The streaming answer is exact, not approximate.
Why the speedup is real
FlashAttention does more arithmetic than a naive kernel, yet runs two to four times faster on typical workloads. The paradox dissolves once you accept the kernel was memory-bound: the extra FLOPs run on idle compute units while the savings come from the HBM traffic that no longer happens.
The formal version is an IO-complexity result. Naive attention moves on the order of O(N^2) bytes through HBM; FlashAttention moves roughly O(N^2 d^2 / M) where M is the SRAM size — asymptotically fewer accesses whenever a meaningful tile fits on-chip. Fewer HBM accesses, same answer: that is the whole performance claim, and it is provable.
Practical implications and the CPU / small-model view
On GPUs, FlashAttention is now the default path: PyTorch’s scaled_dot_product_attention selects it automatically when shapes and dtypes qualify, and serving stacks such as vLLM and TensorRT-LLM build on it. You rarely call it by name; you benefit whenever attention runs.
The principle travels even where the exact CUDA kernel does not. On a CPU running a small language model the same hierarchy exists — DRAM is the slow tier, the caches are the fast tiers — and the same discipline pays off: block the computation so tiles stay resident in cache and avoid writing an N-by-N matrix out to DRAM. On any device the winning move is to fit the working set in the fast tier.
Common pitfalls and where it does not help
FlashAttention is exact, so it will not change your model’s outputs — if a swap seems to shift results, suspect a masking or dtype mismatch, not the algorithm. It does not reduce the O(N^2) compute either, so it is not a substitute for sparse or linear attention when you need to beat that asymptote; it makes dense attention cheaper, not sub-quadratic.
The wins also shrink at short sequence lengths, where the score matrix was never the bottleneck. And correctness lives in the details: causal masking must be applied per tile, and the running-statistics bookkeeping must be right or the numerics drift. These are the mechanics the math companion works through step by step; the takeaway here is the shape of the idea, not the last line of the kernel.