FlashAttention & IO-awareness

1

Concept

The formula for attention is compact, but a straightforward implementation creates large intermediate tensors. For sequence length $N$, the product $QK^\top$ has $N^2$ scores per head. Writing that matrix to accelerator high-bandwidth memory (HBM), reading it for softmax, writing probabilities, and reading them again to multiply $V$ can move far more data than the formula suggests.

Modern accelerators have a hierarchy. Small on-chip SRAM is fast but limited; HBM is larger but slower to access. Arithmetic units can wait while tensors travel between them. IO-awareness means designing an algorithm around this movement, not merely counting floating-point operations. FlashAttention’s key observation is that attention can be computed in tiles that fit fast memory without storing the complete score matrix.

Split queries into row blocks and keys/values into column blocks. Load one query tile and a key/value tile into SRAM, compute their local scores, and update a partial output. Then proceed to the next key/value tile. The obstacle is softmax: normalizing a row seems to require seeing all its scores simultaneously.

Online softmax solves that obstacle. For each query row, maintain a running maximum $m$ and a running exponential sum $\ell$. When a new score block has maximum $m'$, define the new maximum $m_{new}=\max(m,m')$. Previously accumulated sums and outputs are rescaled by $e^{m-m_{new}}$; the new block is scaled by $e^{s-m_{new}}$. After all blocks, the accumulated weighted value is divided by the accumulated normalization.

This is the same stable softmax identity applied incrementally. The algorithm changes operation order and therefore may differ at the last floating-point bits, but it does not intentionally approximate, sparsify, or truncate attention. Causal masking, dropout during training, variable lengths, and gradients require additional bookkeeping, all handled within tiled kernels.

Recomputation can be cheaper than storage. In the backward pass, FlashAttention can recompute local scores rather than reading a saved $N^2$ matrix. That increases some arithmetic while reducing memory traffic and peak memory. The trade makes sense because matrix arithmetic is fast relative to repeated HBM transfers on supported hardware.

The benefit depends on shape and implementation. Very short sequences may be dominated by launch overhead. Head dimension, dtype, mask, dropout, hardware generation, and library version determine available kernels. A framework can silently fall back to a slower path when a feature is unsupported. Measure the actual selected kernel rather than inferring it from an API flag.

During inference, prefill benefits strongly because it processes many positions. Single-token decode has a different shape and is often dominated by reading model weights and the KV cache; specialized decode attention kernels are related but the same headline speedup should not be assumed. Paged cache layouts add another addressing layer that kernels must support.

Correctness tests compare outputs and gradients against a reference across masks, odd lengths, dtypes, and extreme logits. Performance tests record end-to-end phase latency and memory, not only a favorable kernel microbenchmark. The durable insight is broader than one implementation: when compute becomes cheap, algorithm design must count bytes moved through the memory hierarchy. FlashAttention is exact attention rearranged around that physical reality.