FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
Author: Tri Dao (Princeton University & Stanford University)
Published: July 18, 2023
"By optimizing work partitioning across GPU warps and parallelizing along the sequence length dimension, FlashAttention-2 achieves up to 2x speedup over its predecessor, reaching 50-73% of theoretical maximum device throughput."
Why This Matters
Transformers have revolutionized AI, but they suffer from a fundamental bottleneck: the attention mechanismThe core layer of Transformers that computes pairwise relationships between all tokens in a sequence. scales quadratically ($O(N^2)$) with sequence length. This makes processing long documents, high-resolution images, or long-form videos incredibly expensive.
While the original FlashAttention solved the memory footprint issue by avoiding the materialization of the massive $N \times N$ attention matrix in slow High Bandwidth Memory (HBM), it still left significant hardware performance on the table. It only achieved 25-40% of a GPU's theoretical maximum compute throughput. FlashAttention-2 addresses this inefficiency directly, unlocking the ability to train models with context lengths of 16k, 32k, or even 100k+ much faster and cheaper.
The Big Idea
The core bottleneck on modern GPUs isn't just compute; it's the asymmetry between memory speeds. Matrix multiplications (GEMM) run on blazing-fast specialized Tensor CoresSpecialized hardware units on modern GPUs designed specifically for rapid matrix multiplication., while element-wise operations (like softmax, masking, and dropout) run on standard CUDA cores and are heavily memory-bound.
The FlashAttention-2 formula for speed:
1. Reduce non-matmul FLOPs: Tweak the online softmax math so we perform fewer scaling operations.
2. Parallelize along the sequence length: Utilize more GPU Stream Multiprocessors (SMs) even when batch sizes are small.
3. Optimize warp-level layout: Eliminate the need for warps to communicate via slow shared memory during the forward pass.
How It Works
1. Warp Partitioning: Eliminating the "Split-K" Bottleneck
In FlashAttention-1, the keys ($K$) and values ($V$) were split across different warps (groups of 32 threads) while keeping the query ($Q$) accessible to all. This forced warps to write intermediate attention outputs to shared memory, synchronize, and then sum them up.
FlashAttention-2 reverses this: It splits the query $Q$ across different warps while keeping $K$ and $V$ accessible. Each warp computes its slice of the output directly in registers without needing to communicate or write intermediate states back to shared memory.
Interactive Warp Partitioning Comparison
2. Sequence Length Parallelism
FlashAttention-1 parallelized computation primarily over batch size and attention heads. However, when sequence lengths are very long, the batch size is typically small to fit in memory, leaving many GPU Streaming Multiprocessors (SMs) idle.
FlashAttention-2 adds an additional dimension of parallelism: the sequence length dimension. It schedules different row blocks of the attention computation on separate thread blocks, maximizing GPU occupancy and ensuring all SMs stay fully utilized.
3. Mathematical Tweaks to Reduce Non-Matmul FLOPs
GPUs have specialized Tensor Cores that run matrix multiplications extremely fast (up to 16x faster than non-matrix operations). FlashAttention-2 tweaks the online softmax algorithm to minimize these expensive non-matrix operations.
Interactive Playground
Adjust the parameters below to see how sequence length and head dimension scale the theoretical FLOPs, and compare the execution time across different attention implementations.
Tiling & SRAM Memory Simulator
Step through the tiling loop to see how blocks of $Q$, $K$, and $V$ are loaded into fast SRAM cache to compute attention incrementally.
Load block $Q_i$ from slow GPU High Bandwidth Memory (HBM) into fast on-chip SRAM.
Results & Benchmarks
FlashAttention-2 was benchmarked on NVIDIA A100 GPUs and demonstrated massive throughput gains over standard PyTorch attention and the original FlashAttention.
Attention Forward Pass Speed (A100 GPU, TFLOPs/s)
Note: Higher is better. FlashAttention-2 reaches up to 73% of the theoretical maximum throughput (312 TFLOPs/s on A100) for FP16/BF16 matrix multiplications.
Limitations & Open Questions
Despite its incredible performance, FlashAttention-2 has a few notable limitations:
- Hardware Specificity: The low-level optimizations are highly tuned to NVIDIA GPU architectures (specifically Ampere and Hopper). Adapting these to other architectures like AMD or Apple Silicon requires significant custom engineering.
- Manual Block-Size Tuning: Different head dimensions require manually tuned block sizes to maximize register usage and avoid memory spilling.
- Compilation Complexity: Writing and debugging highly optimized CUDA kernels is notoriously difficult, making custom modifications to the attention mechanism hard for typical ML practitioners.
Glossary
High Bandwidth Memory (HBM)
The main off-chip memory on modern GPUs. It has high capacity (e.g., 40-80GB on A100) but is relatively slow compared to on-chip SRAM cache.
SRAM (Static Random-Access Memory)
Super-fast, on-chip memory located right next to the GPU's processing cores. It has very limited capacity (e.g., 192KB per SM on A100) but incredible bandwidth.
Warp
The basic unit of execution on NVIDIA GPUs, consisting of a group of 32 threads executing the same instruction in parallel (SIMT).
Cite this Paper
@article{dao2023flashattention2,
title={FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning},
author={Dao, Tri},
journal={arXiv preprint arXiv:2307.08691},
year={2023}
}