FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning

Abstract

Scaling Transformers to longer sequence lengths has been a major problem in the last several years, promising to improve performance in language modeling and high-resolution image understanding, as well as to unlock new applications in code, audio, and video generation. The attention layer is the main bottleneck in scaling to longer sequences, as its runtime and memory increase quadratically in the sequence length. FlashAttention exploits the asymmetric GPU memory hierarchy to bring significant memory saving (linear instead of quadratic) and runtime speedup (2-4× compared to optimized baselines), with no approximation. However, FlashAttention is still not nearly as fast as optimized matrix-multiply (GEMM) operations, reaching only 25-40% of the theoretical maximum FLOPs/s. We observe that the inefficiency is due to suboptimal work partitioning between different thread blocks and warps on the GPU, causing either low-occupancy or unnecessary shared memory reads/writes. We propose FlashAttention-2, with better work partitioning to address these issues. In particular, we (1) tweak the algorithm to reduce the number of non-matmul FLOPs (2) parallelize the attention computation, even for a single head, across different thread blocks to increase occupancy, and (3) within each thread block, distribute the work between warps to reduce communication through shared memory. These yield around 2× speedup compared to FlashAttention, reaching 50-73% of the theoretical maximum FLOPs/s on A100 and getting close to the efficiency of GEMM operations. We empirically validate that when used end-to-end to train GPT-style models, FlashAttention-2 reaches training speed of up to 225 TFLOPs/s per A100 GPU (72% model FLOPs utilization).

DarcStar commentary

Our perspective on this work and how it informs our research — not part of the original paper.

FlashAttention-2 is the paper we point to when someone assumes a correct algorithm is a fast one. The first FlashAttention was already exact and memory-optimal — and still left more than half the GPU on the floor, reaching a fraction of what the matmul units can do. Not because the math was wrong, but because the work partitioning was: how the job splits across thread blocks and warps, how much traffic crosses shared memory, how full the machine stays. Fix that mapping and you get roughly another 2×, up near the efficiency of a raw GEMM — with the exact same output.

That gap between “provably correct” and “actually fast” is one we take seriously. It's why our order of operations is correctness first, kernel second — a right answer is table stakes, and the wall-clock win lives down in the hardware mapping, not in the flowchart. A FLOP-count cut isn't a speedup until it survives contact with occupancy and memory bandwidth, and plenty of clever-looking attention variants never do.

It also settled a build-versus-buy question for us. Work this good, this close to the metal, isn't something to reinvent — it's a mature, best-in-class layer to build on. The interesting ground isn't re-deriving the fast exact kernel; it's what you compose on top of it.