FlashAttention-2: complete rewrite, about 2x faster, with MQA/GQA and head-dim up to 256
In one sentence On July 17, 2023 Tri Dao published FlashAttention-2: about 2x faster than FlashAttention, 50-73% of theoretical FLOPs on A100 (vs 25-40%), parallelism along sequence length too, MQA/GQA support and head dimensions up to 256.
One year after FlashAttention, Tri Dao published a complete rewrite. The original version did not fully use the GPU: work was split only by batch and attention heads, so with long sequences and small batches many compute units sat idle.
FlashAttention-2 also splits the work along the sequence, reduces the operations that are not matrix multiplications (slow on modern GPUs) and organizes work better between thread groups. It supports head dimensions up to 256 and the MQA and GQA attention variants, used to reduce KV cache memory.
The result is roughly twice the speed of FlashAttention.
Companies
Princeton University, Stanford University
Tools
FlashAttention-2, CUDA, PyTorch
Tags
Sources