FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (Dao, 2023)
The follow-up to flash-attention-paper — Tri Dao, 2023 (arXiv:2307.08691). The original flash-attention-paper proved attention is IO-bound and made it linear-memory and exact; this paper is the GPU-utilization sequel, the source of the hard FA-2 numbers the flash-attention page mentioned by name but without figures.
What changed from FlashAttention-1
The first version left a lot of the GPU idle — it hit only 25–40% of peak FLOPs/s on an A100. FA-2 closes that gap with three engineering moves:
- Fewer non-matmul FLOPs. Tensor cores are far faster at matmul than at other math, so the algorithm is retuned to minimize the non-matmul work that bottlenecks it.
- More parallelism. Attention for a single head is split across thread blocks (not just across batch and heads), raising GPU occupancy on long sequences with small batch sizes.
- Better work partitioning between warps inside a thread block, cutting shared-memory communication.
Numbers
- 50–73% of theoretical peak FLOPs/s on A100 — roughly 2× faster than FlashAttention-1.
- Up to 225 TFLOPs/s per A100 at 72% model-FLOPs utilization training GPT-style models end-to-end.
Why it matters for the spoke
The flash-attention page tracks FA-2/3/4 as “a moving target tracking new GPU generations”; this is the primary for the FA-2 rung of that ladder. It sharpens the spoke’s compute story: FlashAttention-1 answered can attention be exact and memory-cheap? (yes); FA-2 answers can it also saturate the hardware? — moving attention from ~30% to ~70% of peak. The descendant CUDA kernels live in the maintained flash-attention repo.
Tier
T1 — primary, single-author technical paper from the FlashAttention author, with reproducible A100 utilization benchmarks. Figures here are from the abstract; the arXiv PDF body was not fully extracted.
Related
flash-attention · flash-attention-paper · prefill-decode-kv-cache · kv-cache · llm-inference · vllm