FlashAttention
On 27 May 2022 Tri Dao, Daniel Fu, Stefano Ermon and Christopher Re of Stanford and Atri Rudra of the University at Buffalo posted FlashAttention, an exact attention algorithm that computes attention in tiles so as to reduce reads and writes between a GPU's high-bandwidth memory and its on-chip SRAM. It trained BERT-large 15 per cent faster than the MLPerf 1.1 record, GPT-2 3 times faster at 1K tokens, and the Long Range Arena 2.4 times faster at 1K to 4K.
Why it matters
Attention was made faster not by approximating it but by counting memory traffic, the cost approximate methods had left in place. Longer context became affordable: the first Transformers better than chance on Path-X (16K tokens, 61.4 per cent) and Path-256 (64K, 63.1 per cent). Falcon in 2023 names FlashAttention among its architectural choices.
The paper analyses the algorithm's IO complexity, shows that it needs fewer accesses to high-bandwidth memory than standard attention and is optimal for a range of SRAM sizes, and extends it to block-sparse attention. Longer context gave 0.7 better perplexity on GPT-2 and 6.4 points on long-document classification.