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

The problem: attention leaves much of the GPU idle

Attention is the operation that lets a transformer model mix information across positions in a sequence: each position compares itself with the others and combines their representations accordingly. Because it sits inside every layer, how fast the attention step runs strongly influences how fast such a model can be trained or served. Tri Dao submitted the FlashAttention-2 paper to arXiv on July 17, 2023 ^[1].

Hardware speed is usually described in FLOPs/s — floating-point operations per second — and every GPU has a theoretical maximum: the rate it could sustain if nothing but arithmetic were happening. The earlier FlashAttention reached only 25–40% of that maximum ^[1]. The gap was not caused by doing too much arithmetic but by how the arithmetic was arranged. The paper points to suboptimal work partitioning between thread blocks and warps on the GPU ^[1]. A thread block is a unit of work that the hardware schedules onto a processor; a warp is a group of threads inside a block that execute together in lockstep. If some blocks finish early while others are still running, or warps inside a block sit waiting on one another, the arithmetic units starve.

What FlashAttention-2 changes

The paper makes three changes ^[1]. The first reduces the number of non-matmul FLOPs — operations such as exponentials, maxima and sums that are not matrix multiplications and therefore cannot run on the specialized matrix units at full rate. The second parallelizes the attention computation across thread blocks even for a single head, so that one head's work can spread over more of the GPU and raise occupancy, the fraction of the hardware that is actually busy. The third distributes work between warps within each thread block, so that the warps need to exchange less through shared memory — a small on-chip store whose traffic can become a bottleneck of its own.

Results

Together these changes give roughly a 2x speedup over FlashAttention, and the kernel reaches 50–73% of the theoretical maximum FLOPs/s on an A100 GPU ^[1]. The paper breaks the figure down by pass: up to 73% of theoretical maximum throughput in the forward pass and up to 63% in the backward pass ^[1]. The backward pass is the part that computes gradients, and it is usually the harder of the two to keep efficient.

The number that matters most in practice is end-to-end training: when FlashAttention-2 is used to train GPT-style models, training speed reaches up to 225 TFLOPs/s per A100 GPU ^[1]. That figure covers a whole training step rather than one isolated attention kernel, so it describes the model as a whole rather than the attention operation alone.

Why it matters

Attention is called once per layer, so making it faster shortens every training step and every token generated. The reported percentages are also a useful kind of measurement: they say how close the implementation comes to the hardware ceiling, rather than only how much faster it is than the previous version. Moving from 25–40% to 50–73% means the same GPU delivers roughly twice the useful throughput on the same work ^[1].