FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Attention compares every token with every other token, so a sequence of N tokens produces an N × N matrix of scores. Standard implementations compute that matrix, write it to the GPU's main memory (HBM), read it back to apply softmax, write the result, and read it again to multiply by the values. Many papers tried to make attention cheaper by approximating it with fewer operations, but they rarely ran faster in wall-clock time.

Tri Dao, Daniel Fu, Stefano Ermon, Atri Rudra and Christopher Ré at Stanford and Buffalo identified why: on modern GPUs, attention is limited by memory traffic, not arithmetic. An A100 has about 20 MB of on-chip SRAM running at roughly 19 TB/s, and 40 to 80 GB of HBM at 1.5 to 2 TB/s. Writing and reading an N × N matrix through HBM is what takes the time.

FlashAttention computes exact attention without ever storing that matrix. It splits queries, keys and values into blocks that fit in SRAM, and for each block computes scores, applies softmax and multiplies by values on-chip, writing only the output. The trick that makes this possible is an online softmax: keep a running maximum and running sum per row, and rescale earlier partial results when a later block changes them, so the final answer is exactly the same. For training, the backward pass recomputes the scores from the stored inputs and these small statistics rather than reading a stored N × N matrix. Everything is fused into one CUDA kernel.

It does slightly more arithmetic and much less memory traffic. On GPT-2 medium attention, traffic fell from 40.3 GB to 4.4 GB and runtime from 41.7 ms to 7.3 ms. End to end, GPT-2 trained 3× faster than the Hugging Face implementation, BERT-large 15% faster than the MLPerf 1.1 record, and attention memory grew linearly with sequence length. A block-sparse version went further, to 64K tokens.

FlashAttention and its successors are now inside nearly every serving and training stack, from PyTorch's built-in attention to vLLM and SGLang. You rarely call it directly, but it shapes what is possible and what to check.

It is why long context became affordable. Attention memory linear in sequence length is what allows contexts of tens or hundreds of thousands of tokens on one GPU. When a customer asks why a model supports long inputs, the answer includes this kernel.

Check it is actually on. Older model code, custom attention variants, unusual head sizes or older GPUs can silently fall back to standard attention. If long-context latency or memory looks wrong, confirm the fast kernel is being used.

Think in memory traffic, not FLOPs. The paper's lesson generalizes: on GPUs, many operations are bound by moving data. When profiling a slow model, measure memory bandwidth use, not just FLOPs; fusion and keeping data on-chip often matter more than reducing math.

Exact, not approximate. FlashAttention changes nothing about model outputs (up to floating-point rounding), which makes it a safe optimization to enable, unlike approximate attention methods that change model quality.

Hardware specifics matter. The paper notes kernels may not transfer across GPU architectures; later versions were tuned for newer GPUs. Performance claims depend on the GPU generation.