flashattention

FlashAttention does the same arithmetic as naive attention and is several times faster. The answer is a three-line recurrence and one IO count, and interviewers want both.

The one-query recurrence is a warm-up. A real kernel carries a vector of maxima and a vector of sums for a tile of rows, and the rescale becomes a diagonal matrix. The answer writes it, runs it, and shows the edge case that turns it into NaN.

The textbook IO count says FlashAttention re-reads K and V once per query block, which at 8K is only 2x better than naive. The answer explains why the real HBM number is 65x lower, and which cache makes the difference.

Storing the attention matrix for backward costs gigabytes per layer; recomputing it costs extra FLOPs. The answer derives the gradient, shows the one scalar per row that makes recomputation exact, and the identity that removes the last N-squared term.

FlashAttention-2 reaches well under half of an H100's BF16 peak. The answer explains why in one ratio, exponentials against matmuls, and maps each FA3 technique to the stall it hides.

A 4K window over 32K tokens keeps one eighth of the keys, and the speedup over causal attention is not 8x. The answer counts tiles exactly, classifies which need masking, and explains where the rest of the gap goes.

Packing removes padding but puts different requests' tokens side by side in one tensor, so the kernel's tile boundaries no longer match sequence boundaries. The answer shows the indexing and a bug that a causal test can never catch.

An error that grows with context length is a running sum losing increments, and BF16 stops adding 1 to a total at 256. The answer turns the symptom into a bisection, finds the state held in the wrong dtype, and shows the error curve that proves it.