[Paper Review] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

Tri Dao, Daniel Y. Fu+3 more|arXiv (Cornell University)|May 27, 2022

Advanced Neural Network ApplicationsComputer Science457 citations

TL;DR

FlashAttention computes exact attention with drastically reduced memory IO by tiling and recomputation, achieving faster training and enabling longer context, plus a block-sparse variant for further speedups.

ABSTRACT

Transformers are slow and memory-hungry on long sequences, since the time and memory complexity of self-attention are quadratic in sequence length. Approximate attention methods have attempted to address this problem by trading off model quality to reduce the compute complexity, but often do not achieve wall-clock speedup. We argue that a missing principle is making attention algorithms IO-aware -- accounting for reads and writes between levels of GPU memory. We propose FlashAttention, an IO-aware exact attention algorithm that uses tiling to reduce the number of memory reads/writes between GPU high bandwidth memory (HBM) and GPU on-chip SRAM. We analyze the IO complexity of FlashAttention, showing that it requires fewer HBM accesses than standard attention, and is optimal for a range of SRAM sizes. We also extend FlashAttention to block-sparse attention, yielding an approximate attention algorithm that is faster than any existing approximate attention method. FlashAttention trains Transformers faster than existing baselines: 15% end-to-end wall-clock speedup on BERT-large (seq. length 512) compared to the MLPerf 1.1 training speed record, 3$ imes$ speedup on GPT-2 (seq. length 1K), and 2.4$ imes$ speedup on long-range arena (seq. length 1K-4K). FlashAttention and block-sparse FlashAttention enable longer context in Transformers, yielding higher quality models (0.7 better perplexity on GPT-2 and 6.4 points of lift on long-document classification) and entirely new capabilities: the first Transformers to achieve better-than-chance performance on the Path-X challenge (seq. length 16K, 61.4% accuracy) and Path-256 (seq. length 64K, 63.1% accuracy).

Motivation & Objective

• Motivate memory IO as a bottleneck in self-attention on GPUs and propose an IO-aware exact attention method.
• Reduce reads/writes of the large N×N attention matrix by tiling inputs and performing softmax incrementally.
• Avoid storing the full attention matrix for backward by recomputing from on-chip statistics and output.
• Extend to block-sparse attention to enable faster approximate attention at long sequence lengths.
• Provide open-source implementation and empirical validation across baselines and long-context tasks.

Proposed method

• Reformulate attention with tiling to load K and V in blocks into SRAM and accumulate O across Q blocks.
• Use algebraic aggregation to compute softmax in blocks with maintained statistics m and ell for numerical stability.
• Apply recomputation during backward by storing only O and softmax statistics to reconstruct S and P as needed.
• Fuse all steps in a single CUDA kernel to minimize memory traffic and avoid materializing the full N×N matrices.
• Provide IO-complexity analysis showing O(N²d²/M) HBM accesses for FlashAttention vs. Ω(Nd+N²) for standard attention.
• Extend to block-sparse FlashAttention with a fixed sparsity mask to reduce IO proportional to sparsity.

Experimental results

Research questions

• RQ1 How can attention be computed exactly while minimizing GPU HBM accesses?
• RQ2 Can tiling and recomputation yield wall-clock speedups over standard attention without sacrificing exactness?
• RQ3 How does block-sparse FlashAttention trade accuracy for IO efficiency and speed?
• RQ4 What are the IO lower bounds for exact attention across SRAM sizes, and can they be approached by practical algorithms?
• RQ5 Do IO-aware implementations enable longer context and higher-quality Transformer models in practice?

Key findings

• FlashAttention achieves up to 7.6× speedup over GPT-2 baseline on attention computation and reduces HBM reads/writes substantially.
• For typical head dimensions and SRAM sizes, FlashAttention requires far fewer HBM accesses than standard attention and is memory-footprint efficient (O(N) beyond inputs/outputs).
• Training speed improves: BERT-large 15% faster than MLPerf 1.1 record; GPT-2 up to 3× faster than HuggingFace baseline; LRA 2.4× faster.
• Long-context benefits include 0.7 perplexity improvement on GPT-2 and 6.4-point lift on long-document classification; Path-X and Path-256 achieving better-than-random performance with long sequences.
• Block-sparse FlashAttention yields 2–4× faster speeds than FlashAttention, scaling to 64K sequences while maintaining comparable quality.