FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
Tri Dao2023arXiv 2023
Read it on arxiv.org(opens in a new tab)Why this one
Read this after FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. The first paper teaches the big lesson, avoid writing the attention matrix to slow memory. This one teaches the more painful lesson, a correct IO idea can still leave half the GPU idle if the work is sliced badly. People often talk about FlashAttention-2 as just a faster kernel, but the durable idea is lower level: parallelism is a resource you must budget across blocks, warps, shared memory, and non-matmul instructions. For an engineer serving or training transformers, this is a rare paper that connects algorithm shape to the actual hardware schedule. It is worth your evening because it changes how you read performance claims. You stop asking only whether an operation is fewer FLOPs or fewer bytes, and start asking whether the GPU can keep its expensive math units fed.
What to take away
- The speedup comes from repartitioning attention work, not from approximating attention or changing model outputs.
- Non-matmul FLOPs matter because GPU tensor cores are fast enough that scalar bookkeeping becomes visible.
- Parallelizing a single attention head across thread blocks improves occupancy for long sequences and small batch sizes.
Reads with
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
the original io aware algorithm that flashattention 2 parallelizes more carefully
- Ring Attention with Blockwise Transformers for Near-Infinite Context
extends blockwise exact attention when a single gpu schedule is no longer enough
- Efficient Memory Management for Large Language Model Serving with PagedAttention
shows the serving side bottleneck once attention kernels stop wasting hbm bandwidth