FlashAttention-2
lastmod 2026-03-08created 2023-07-212 edits
FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
Faster Attention through better parallelism and work partitioning
Features
- Language models with longer contexts have emerged, such as GPT-4 (32k), MPT (65k), and Calude (100k)
- Scaling the context length of Transformers is difficult because runtime & memory requirements grow to the fourth power
- FlashAttention, released last year, reduced memory usage and increased attention speed, and is used in many places
- It was already 2-4x faster at release, but there is still room for improvement. It is still not as fast as optimized matrix multiplication (GEMM), reaching only 25-40% of theoretical maximum FLOPs/s (up to 124 TFLOPs/s on an A100 GPU)
- FlashAttention-2 is 2x faster than the previous version, delivering up to 230 TFLOP/s on an A100 GPU
- Reached up to 225 TFLOPS in training GPT-style language models (72% model FLOP utilization)
- Tweaked the algorithm to reduce non-matmul FLOPs
- Better parallelism, changed how work is partitioned within each thread block
- Extended the number of Head Dimensions from 128 to 256