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

Why attention is a memory problem

Attention lets every position in a sequence look at every other position, which is what makes Transformers effective — and what makes them expensive as sequences grow. Most efforts to speed it up focus on reducing arithmetic. The authors of FlashAttention — Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra and Christopher Ré — argue that a principle has been missing: making attention algorithms IO-aware, that is, accounting for the reads and writes between levels of GPU memory rather than counting only floating-point operations. The paper was submitted to arXiv on May 27, 2022 ^[1].

The hardware gap that motivates this is stark. An A100 GPU has 40-80GB of high bandwidth memory (HBM), but each of its 108 streaming multiprocessors has only 192KB of on-chip SRAM ^[1]. Moving data between the two is far slower than computing on data already in place. A conventional attention implementation writes the large matrix of attention scores out to HBM and then reads it back for later steps, so a substantial share of time goes to memory traffic rather than to arithmetic.

Tiling, without giving up exactness

FlashAttention is an IO-aware exact attention algorithm. It uses tiling: the inputs are broken into blocks small enough to fit in on-chip SRAM, and the attention computation proceeds block by block, with intermediate results kept on chip and combined incrementally instead of being written out to HBM and read back again. The paper reports that this requires fewer HBM accesses than standard attention and is optimal for a range of SRAM sizes ^[1].

The word "exact" is central. Many earlier ways of making attention cheaper were approximations that skipped or compressed parts of the score matrix, trading model quality for speed. FlashAttention computes the same attention result as the standard formulation; the savings come from how data is moved, not from computing something different.

Measured speedups

The paper reports faster training than existing baselines on several workloads. Against the MLPerf 1.1 training speed record, it achieves a 15% end-to-end wall-clock speedup on BERT-large; it reaches 3x speedup on GPT-2, and 2.4x on the Long Range Arena benchmarks ^[1]. Because these are wall-clock comparisons on training runs rather than counts of operations, they reflect the memory-traffic argument directly: the same mathematical work finishes sooner.

Longer context, and a hard benchmark

Speed is not the only consequence. FlashAttention and its block-sparse extension make longer contexts practical, which the authors link to higher-quality models. They report the first Transformers to perform better than chance on the Path-X challenge, where sequences are 16K tokens long, reaching 61.4% accuracy ^[1]. Path-X is notable because attention over such long sequences had previously been out of reach for standard implementations, so a model could do no better than guessing.

What the work changed in how the problem is framed

The lasting contribution described in the paper is a way of thinking as much as a kernel: treat the memory hierarchy as a first-class part of the algorithm design, and attention becomes considerably cheaper without any approximation. That framing connects an abstract operation — scaled dot-product attention — to the concrete constraints of the hardware it runs on, and it explains why gains of this size were available at all. It also shifts the practical boundary of what sequence lengths a Transformer can be trained on, since both time and memory are governed by the same traffic between HBM and SRAM ^[1].