Counting bytes, not FLOPs
For one head, Q, K, V and O are each N×d and the work is about 4N²d FLOPs: 2N²d for QKᵀ, the same again for the value multiply. A kernel that materializes the score matrix makes several full passes over an N×N array — write S, read it for the softmax, write P, read P again. Call it c passes; in fp16 that is 2cN² bytes across HBM.
So the arithmetic intensity is about 2d/c FLOPs per byte — near 60 for a head dimension of 128 and four passes, while current datacenter parts want intensities in the low hundreds before dense tensor-core math saturates. Attention lands on the wrong side of the roofline by a wide margin, and the N×d input and output traffic is not what puts it there. The N² intermediate is. Delete it and the same FLOPs are fed by reading Q, K, V once and writing O once: traffic linear in N, arithmetic unchanged.
Why a compiler could not fuse this for you
The sharpest thing to say about FlashAttention is that it is not the fusion a compiler missed. Ordinary fusion — the sibling kernel-fusion article’s subject — chains operations so intermediates stay on chip, and it works because each output element depends on the matching input element. Softmax breaks that: its normalizer is a reduction along exactly the axis you want to tile away.
A pass that wants to merge QKᵀ with the softmax must produce a correct row denominator before the row exists. No dependence analysis over the operator graph can license that, because the justification is not in the graph — it is an algebraic identity about softmax being invariant to subtracting a constant from its inputs. That identity, derived in the Flash Attention Math companion, turns a whole-row reduction into a small running state a tiled loop can carry. FlashAttention is the rewrite that makes the fusion legal — hence hand-written CUDA rather than a compiler pattern.
The SRAM residency ledger
Take the hierarchy as given. The kernel author’s question is narrower: which tensors must be simultaneously resident to run one inner-loop iteration? A staged Q row-block Br×d, plus a K column-block and a V block each Bc×d, in shared memory; a score tile Br×Bc and an output accumulator Br×d in registers; and the running statistics, a couple of floats per query row.
Illustratively: at Br = Bc = 128, d = 64, fp16, each staged tile is 128×64×2 = 16 KB, so 48 KB before any double-buffering. Move to d = 128 and it doubles; add a second buffer for the incoming K/V tile and you are past what a kernel gets by default. Hence the opt-in dynamic shared-memory carveout above the static per-block limit, and hence Br and Bc being compile-time template parameters instantiated per head dimension and dtype rather than runtime knobs.
Tile shape is a register decision, and occupancy pays for it
Shared memory is only half the ledger, and usually not the binding half. The score tile and the output accumulator live in registers for the whole inner loop, in fp32 even when the operands are fp16, because the accumulation must stay stable across every K/V block. Each SM has a fixed budget of 32-bit registers shared by all resident threads, and the ISA caps any thread at 255. Widening Bc grows the score tile; widening Br grows both. Push too far and the compiler spills to local memory, which is HBM-backed — so a tile size chosen to save bandwidth spends it on the hottest loop.
The price is occupancy. Large tiles and heavy register use often leave one or two thread blocks resident per SM, which classic guidance calls a mistake: with few warps there is nothing to switch to when one stalls. It works anyway, because the stall has been designed out — latency hiding comes from instruction-level parallelism inside a deliberately fat loop and from asynchronous copy staging the next K/V tile while the current one feeds the tensor cores. Occupancy is a proxy for latency hiding, not the goal.
Inside the inner loop: two GEMMs and a softmax between them
Each iteration is two chained matrix multiplies. First S = QiKjᵀ, fp16 or bf16 in, fp32 accumulate, on tensor cores. Then, after the softmax step, P Vj into the output accumulator — tensor cores again. Between them sits work that does not touch them at all: row max, subtraction, exponential, row sum, and the rescale of everything accumulated so far.
The exponentials go to the special function units, whose throughput is far below the tensor cores’. Implementations shave what they can — folding log₂e into the softmax scale so the hardware’s base-2 exponential is used directly, and keeping the accumulator unnormalized so the division happens once at the end — but the structural problem survives: a low-throughput serial stage sits in the dependency chain between two high-throughput matrix stages. Overlapping the two classes of work, so one tile’s softmax runs while another’s GEMM does, is what recent kernel generations attack.
The backward pass: recompute, and where gradients collide
That the backward pass recomputes rather than stores is the companion article’s point: the forward saves O and one log-sum-exp scalar per query row, and the backward regenerates each S and P tile from Q, K and that scalar — extra arithmetic on operands already on chip in exchange for never touching an N² array.
The kernel-level complication is where the gradients accumulate. dQ for a query block takes a contribution from every K/V block; dK and dV for a key block take one from every query block. Those reductions run along opposite axes, so whichever axis goes on the outer loop, one of the three gradients is written by many thread blocks at once. Implementations resolve it with atomic accumulation into an fp32 dQ buffer, or by splitting the backward into separate passes with opposite loop orders, each keeping its gradient block-local. That choice, not the recomputation, is why backward kernels are the hard ones.