FlashAttention cut memory reads to speed up exact attention

An IO-aware rewrite, not a new approximation

In May 2022, researchers proposed computing standard, exact attention differently rather than approximating it. The FlashAttention paper argues that on modern GPUs the bottleneck for attention is not arithmetic but data movement between the GPU's large slow memory (HBM) and its small fast on-chip memory (SRAM). The method tiles the attention computation into blocks that fit in fast memory and recomputes intermediate values during the backward pass instead of storing them, so it avoids ever materialising the full attention matrix. The output is mathematically identical to standard attention; nothing about what the model computes changes, only how the computation is scheduled across memory.

What the paper measured

The paper reports a 15 percent end-to-end wall-clock speedup on BERT-large, roughly a threefold speedup on GPT-2 at a sequence length of 1,000 tokens, and a 2.4-fold speedup on the Long-Range Arena suite, all measured against then-standard implementations. It also reports that the longer context enabled by the lower memory footprint produced measurable quality changes on the tasks tested, including a 0.7 perplexity improvement on GPT-2 and a 6.4 point gain on a long-document classification task, and that models using the method were the first reported to do better than chance on the Path-X and Path-256 long-range tasks. A follow-up paper submitted in July 2023 found the original algorithm reached only 25 to 40 percent of a GPU's theoretical maximum FLOPs per second, and reworked the parallelism to reach 50 to 73 percent on an A100, roughly doubling throughput again.

What the benchmark conditions do not establish

These are speedups on specific model sizes, sequence lengths and GPU generations named in the papers; a different sequence length, batch size or accelerator can change the ratio, sometimes substantially. The project repository now documents further hardware-specific versions, including a build for Hopper GPUs and one for AMD accelerators, each with its own reported numbers rather than a single universal figure. A reader comparing runtimes should treat any single quoted multiplier as tied to the hardware and shapes it was measured on, not as a fixed property of the algorithm.

• What sequence length and batch size produced the quoted speedup, and do they match your workload?
• Which GPU generation was the comparison run on, and is a newer kernel version available for it?
• Is the baseline it was compared against still the best available alternative, or has that baseline itself moved on?

The mechanism is well specified and independently reproducible in the released code; the specific multipliers are properties of the hardware and shapes tested, not universal constants.