FlashAttention makes exact attention IO-aware

Tri Dao, Christopher Ré and colleagues at Stanford reformulated the attention operation at the heart of the Transformer to account for how data moves through a GPU’s memory hierarchy. Where much prior work had tried to make attention cheaper by approximating it — accepting some loss of accuracy for lower cost — FlashAttention computed the same result exactly, but ordered the computation so that it read from and wrote to the GPU’s slow high-bandwidth memory far less often.

The bottleneck the authors identified was not arithmetic but memory access: standard attention materialised a large intermediate matrix in slow memory, and the time spent shuttling that matrix back and forth dominated. By tiling the computation and keeping intermediate values in the GPU’s fast on-chip memory — a technique the paper called IO-awareness — FlashAttention avoided ever writing the full matrix out. The authors reported a 15% end-to-end speedup on BERT-large and substantial gains on GPT-style training, while using memory that grew linearly rather than quadratically with sequence length, which made much longer context windows practical.

Because the method changed nothing about a model’s outputs — it was a faster route to an identical result — it carried no accuracy trade-off to weigh, and adoption was rapid. FlashAttention and its successors became standard components of the training and inference stacks at major labs, one of the pieces of unglamorous systems work that made the long-context models of the following years affordable to run.