FlashAttention Algorithm Solves Transformer Memory Bottlenecks
A new IO-aware exact attention algorithm called FlashAttention speeds up Transformer training by optimizing memory reads and writes between GPU memory levels.
Transformer models struggle with long sequences due to the quadratic time and memory complexity of standard self-attention. While approximate attention methods attempt to solve this issue by sacrificing model quality, they frequently fail to deliver actual wall-clock speedups. A new paper introduces FlashAttention, an exact attention algorithm that addresses this gap by making the computation IO-aware and accounting for data transfers between GPU high bandwidth memory and on-chip SRAM.
FlashAttention uses a technique called tiling to significantly reduce the number of memory reads and writes during the attention calculation. This approach proves mathematically optimal for a range of SRAM sizes and requires fewer HBM accesses than standard attention implementations. The researchers also extend this method to create block-sparse FlashAttention, which provides an approximate attention variant that outperforms all existing approximate methods in speed.
This breakthrough yields impressive real-world training results, including a 15% end-to-end speedup on BERT-large and a 3x speedup on GPT-2. By enabling longer context windows, FlashAttention improves model quality and unlocks entirely new capabilities, allowing Transformers to achieve better-than-chance performance on extremely long sequences like Path-256 with a sequence length of 64K.