FlashAttention is usually explained as an algorithm. This piece explains it as a kernel. The algorithmic core — tiling plus an online softmax that combines partial rows exactly — belongs to the Flash Attention architecture and math companions, where the rescaling identity is derived; taken as given here. What is left is what decides whether an implementation is actually fast: how many bytes cross HBM and why that number sets the runtime; what must be resident in shared memory and registers for one inner-loop iteration, and how that pins the tile shapes; where the tensor cores end up waiting on exponentials; how gradients accumulate along conflicting axes in the backward pass; and why decoding needs a differently shaped kernel.
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.
Masks, head dimension, and GQA at the tile level
Causal masking is nearly free in a tiled kernel and pure waste in a materializing one. Once the loop runs over K/V tiles, any tile lying entirely above the diagonal contributes nothing and is never visited — the loop bound shrinks, and only tiles straddling the diagonal need an element-wise mask. A kernel that builds S first has already paid to compute and store what it discards. The catch is load imbalance: late query blocks iterate over far more K/V tiles than early ones, so a naive one-block-per-query-tile launch leaves SMs idle at the tail, and schedulers compensate by splitting or reordering work.
Head dimension is a compile-time parameter: larger d shrinks the tiles that fit, and an unsupported d is the commonest reason a call silently falls off the fast path. Grouped-query attention helps the ledger directly — several query heads share one K/V head, so a staged K/V tile is reused across heads instead of reloaded.
Decode is a different kernel entirely
During autoregressive decoding there is no N×N matrix to avoid. The query is a single token, so the scores are one row against the whole KV cache. The FLOPs collapse, but the kernel is still memory-bound — more so, because it streams the entire cache out of HBM to produce one token with essentially no reuse. Tiling to keep an intermediate on chip buys nothing when the intermediate was already tiny.
What the forward kernel’s structure does buy is parallelism. With a single query row, batch × heads may not produce enough thread blocks to fill the machine, and long contexts make it worse: more work per block, no more blocks. FlashDecoding-style kernels split the key/value sequence across many blocks, have each compute a partial output with its own running statistics, and combine the partials in a second, tiny reduction kernel — the same rescaling identity the forward pass uses within a block, now applied across blocks.
Confirming the fast path in a profiler
The failure mode is silence. Frameworks dispatch attention through a backend selector, and an unsupported head dimension, dtype, additive bias, or non-contiguous layout quietly drops you onto a reference implementation that materializes the score matrix: same numbers, no warning, quadratic footprint again.
Three checks settle it. Read the kernel names in a timeline trace — the flash kernel and the fallback have different names and very different shapes. Compare measured DRAM traffic over the attention region against the floor implied by Q, K, V and O alone; traffic that scales quadratically with sequence length means the fast path is not running. Then check the fraction of kernel time with tensor cores active: a healthy forward kernel spends most of it inside the two GEMMs, and a low fraction points at the softmax stage or at spilled accumulators rather than at bandwidth.
N×N intermediate stopped travelling to HBM — and the algorithm had to be rewritten before that fusion was even legal, because a softmax reduction runs along the axis you want to tile away. Everything else follows: tile shapes bounded by a shared-memory and register residency budget, occupancy traded for on-chip reuse and pipelined loads, a low-throughput softmax wedged between two tensor-core GEMMs, and a backward pass whose difficulty is gradient accumulation along opposite axes rather than the recomputation. Two things transfer: moving less data usually beats doing less arithmetic — and check the profiler, because falling off the fast path is silent.