Standard attention has a step that nobody designs on purpose. It computes an N by N matrix of scores, writes it to high-bandwidth memory, reads it back to apply softmax, writes the result, and reads it again to multiply by V. At a sequence length of 4096 that matrix is 32 MB per head in bf16, and it crosses the memory bus four times. The arithmetic is trivial. The traffic is not.
FlashAttention removes the matrix. It loads a block of queries and a block of keys and values into on-chip SRAM, computes just that tile of scores there, folds it into a running result, and moves on. The tile never leaves the chip. What used to be a term proportional to N squared becomes a term proportional to N squared times the head dimension divided by the block size, and the constant is large enough to matter more than the asymptotics suggest.
The part that makes this legal is the running softmax. An ordinary softmax needs the maximum of the whole row before it can exponentiate anything, which is precisely what forces you to have the whole row. Instead keep a running maximum and a running sum, and when a later block contains a larger value, rescale what you have accumulated by the exponential of the difference. The result is not an approximation of the full softmax. It is equal to it.
The trade is that keys and values get re-read once per query block. That is a real cost, and it is why block size is a tuning parameter rather than a constant: bigger blocks mean fewer re-reads, until the working set stops fitting in SRAM and the kernel spills, at which point you have lost everything you came for. Push the block sliders past the budget to see that boundary.
One honest limit of the picture: the head dimension here is fixed at 64 and the block sizes are chosen for legibility rather than for a real GPU's shared memory, so the traffic ratio is faithful while the absolute numbers are illustrative.
The maths
- Online softmax update
The running maximum and running sum after block k. When a later block contains a larger score, the accumulated sum is rescaled by the exponential of the difference. This is what makes tiling legal, and the result equals the full-row softmax exactly rather than approximately.
- Standard attention HBM traffic
Q, K, V and O account for the 4Nd term. The rest is the score matrix crossing the bus four times: written, read for softmax, written again, read for the value product. The causal mask halves it. This term dominates the moment N exceeds the head dimension.
- FlashAttention HBM traffic
No N by N term. Q and O are touched once; K and V are re-read once per query block, which is the price of never writing the score matrix. Note this is still quadratic in N: the win is the constant, roughly Br divided by d, not the asymptotics.
- SRAM working set per tile
The Q block, the output accumulator, the K and V blocks, the score tile and the running statistics. The output accumulator must stay resident across the whole inner sweep, which is exactly why the query block sits on the outer loop. Exceed the budget and the kernel spills.
Related terms
- AttentionAttention is the mechanism a model uses to weigh how strongly each token in its context relates to the others when predicting the next one. It is the basis of how a model actually uses context.
- InferenceInference is the act of running a trained model to get an answer: text goes in, a prediction comes out. Every message you send to a coding agent is an inference. It is the opposite end of the lifecycle from training.
- Context windowThe context window is the maximum amount of text, measured in tokens, that a model can consider for a single request. It is a hard ceiling, and it is the main resource you manage when working with an agent.