EP067: FlashAttention-2 Unlocks Massive Context Windows

Here is a short summary of the paper FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning:The Problem:Scaling Transformers to handle longer sequences is limited by the attention layer, which has quadratic runtime and memory costs. While the original FlashAttention algorithm significantly improved memory usage and speed, it remained inefficient compared to optimized matrix-multiply (GEMM) operations. It only reached 25-40% of a GPU's theoretical maximum throughput (FLOPs/s) because of suboptimal work partitioning and unnecessary shared memory reads/writes.The Solution:The author proposes FlashAttention-2, which introduces three key optimizations to maximize GPU efficiency:Algorithmic Tweaks: The algorithm is modified to reduce the number of non-matrix-multiply operations, allowing the GPU to spend more time executing much faster matrix-multiply operations.Enhanced Parallelism: In addition to parallelizing over batch size and attention heads, FlashAttention-2 parallelizes the computation along the sequence length dimension. This greatly increases GPU resource utilization (occupancy), especially when processing long sequences with small batch sizes.Improved Work Partitioning: The work is distributed more efficiently between "warps" (groups of threads) within a single thread block. By splitting the queries (Q) across warps while sharing keys (K) and values (V), it drastically reduces the need for communication and shared memory reads/writes.The Results:These updates yield roughly a 2× speedup over the original FlashAttention. FlashAttention-2 successfully reaches 50-73% of the theoretical maximum throughput on A100 GPUs and enables end-to-end training of GPT-style models at speeds up to 225 TFLOPs/s per A100 GPU.