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

This paper introduces FlashAttention, an exact attention algorithm designed to overcome the quadratic time and memory cost of self-attention that makes Transformers slow and memory-hungry on long sequences. Rather than approximating attention and trading away model quality, the authors argue the missing principle is IO-awareness: accounting for the reads and writes between levels of GPU memory. FlashAttention uses tiling to reduce the number of memory reads and writes between GPU high-bandwidth memory (HBM) and on-chip SRAM; the authors analyze its IO complexity, show it needs fewer HBM accesses than standard attention and is optimal for a range of SRAM sizes, and extend it to block-sparse attention that is faster than existing approximate methods.

FlashAttention trains Transformers faster than existing baselines, with a 15% end-to-end speedup on BERT-large, a 3x speedup on GPT-2, and a 2.4x speedup on long-range arena. By enabling longer context, it also yields higher-quality models, reporting 0.7 better perplexity on GPT-2 and 6.4 points of lift on long-document classification, and unlocks new capabilities, including the first Transformers to achieve better-than-chance performance on Path-X (16K length, 61.4% accuracy) and Path-256 (64K length, 63.1% accuracy).