FlashAttention

What is FlashAttention?

FlashAttention reorganizes the attention computation to respect a fact GPU theory often ignores: moving data between GPU main memory (HBM) and on-chip SRAM costs far more time than doing math on it. Standard attention writes the full N×N score matrix to HBM — at 32k tokens that’s a billion entries per head, mostly memory traffic. FlashAttention tiles the computation so blocks of queries, keys, and values stay in fast SRAM, accumulating results with an online softmax and never materializing the big matrix. Same exact output, a fraction of the memory movement.

Why it changed what’s possible

Attention memory goes from quadratic to linear in sequence length, and wall-clock speed improves 2–4× at training-relevant lengths. The practical consequence: the 128k–1M context windows of modern models are affordable largely because FlashAttention (and successors FlashAttention-2 and -3, tuned for newer GPU generations) made long-sequence attention IO-efficient. It’s not an approximation like sliding-window or sparse attention — you give up nothing in accuracy.

FlashAttention vs PagedAttention — the pairing everyone confuses

They solve adjacent but different problems: FlashAttention makes computing attention efficient; PagedAttention makes storing the KV cache efficient. Production serving stacks use both — one per phase of the same bottleneck.

What people get wrong

• Calling it an approximation. Exact attention, computed smarter; benchmarks that “verify quality is preserved” are testing a tautology.
• Expecting gains everywhere. Short sequences and small models are barely attention-bound; the wins concentrate in long-context and training workloads.
• Ignoring hardware coupling: each FlashAttention generation targets specific GPU architectures — running FlashAttention-3 kernels expects Hopper-class hardware; older cards fall back to slower paths.

Primary source: FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Dao et al., 2022)

Why It Doesn’t Change the Math, Only the Speed

An important distinction from approximate attention methods is that FlashAttention computes exactly the same mathematical result as standard attention — it’s a pure systems-level optimization, not an approximation that trades accuracy for speed. Its speedup comes entirely from being “IO-aware”: standard attention implementations repeatedly write and read the full N×N attention matrix to and from GPU high-bandwidth memory, which is comparatively slow, while FlashAttention restructures the computation to keep intermediate results in the GPU’s much faster on-chip SRAM, processing the attention calculation in smaller blocks that never require materializing the full matrix in slower memory at all. This is why adopting FlashAttention requires no retraining or model changes — it’s a drop-in replacement for the attention computation that produces identical outputs, just faster and with less memory usage, which is why it was adopted essentially universally across training and serving frameworks within a year or two of its introduction.