FlashAttention

Attention is expensive at long sequence lengths because every token can attend to every other token. The naive implementation materializes a large sequence-by-sequence score matrix, writes it to GPU high-bandwidth memory, reads it back for softmax, then writes and reads again. FlashAttention keeps the math exact but reorganizes the computation so tiles stream through fast on-chip SRAM.

IO, not FLOPs#

Modern GPUs have enormous compute throughput, but moving bytes to and from HBM is often the bottleneck. Naive attention performs many HBM round trips for an intermediate matrix that does not need to exist all at once. FlashAttention is IO-aware: it treats memory movement as the thing to minimize.

Online softmax#

The trick is that softmax can be computed block by block if the kernel tracks running maxima and normalization terms. FlashAttention loads a tile of queries, keys, and values into SRAM, updates the partial softmax result, and moves on. It never stores the full attention matrix. The output is the same as standard attention up to normal floating-point differences.

Versions#

FlashAttention-1 proved the IO-aware algorithm. FlashAttention-2 improved work partitioning and GPU occupancy. FlashAttention-3 targets Hopper-era hardware more directly, using newer instructions and scheduling to increase throughput. The practical result is simple: longer context windows, larger batches, and less memory pressure during training and serving.

Practical takeaways#

FlashAttention is not an approximation like sparse attention. It is an implementation win. If a serving or training stack supports it, you usually want it enabled. It pairs with PagedAttention but solves a different problem: FlashAttention reduces attention compute memory traffic, while PagedAttention manages the KV cache across requests.