FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision

Why attention is the bottleneck

Attention is the core layer of the Transformer architecture, and the paper takes as its starting point the claim that attention — rather than the layers around it — is the bottleneck for large language models and long-context applications ^[1]. The reason is structural: attention makes every token in a sequence exchange information with other tokens, so the amount of work grows with sequence length, and how well this layer is implemented decides how much context a model can afford.

GPUs such as NVIDIA's H100 can multiply matrices at enormous rates, but an attention layer is not one clean matrix multiplication. It is a chain of matrix products interleaved with a softmax normalization, and intermediate results must be moved between different levels of memory. A kernel that leaves the arithmetic units waiting on memory wastes a large part of the machine. The paper measures this gap directly: FlashAttention-2, the previous version of the method, achieves only 35% utilization on the H100 ^[1].

Three techniques for Hopper GPUs

The paper develops three techniques aimed specifically at Hopper-generation GPUs ^[1].

The first is overlapping computation and data movement through warp-specialization ^[1]. A GPU kernel runs many threads organized into groups called warps; warp-specialization means giving different warps different jobs, so that some warps are dedicated to fetching data while others do arithmetic on data that has already arrived. The transfer then hides behind the computation instead of stopping it.

The second is interleaving block-wise matrix multiplication with softmax operations ^[1]. Attention is computed in blocks rather than all at once, because no single processor can hold a whole long sequence. Softmax needs running statistics over the scores, which cannot be produced by matrix multiplication alone and would otherwise leave the fast matrix units idle. Interleaving the two kinds of work keeps those units busy while the normalization is being worked out.

The third is block quantization with incoherent processing, which takes advantage of the hardware's FP8 support ^[1]. FP8 stores each number in 8 bits instead of 16, which lets the hardware move more values per unit of time, at the cost of coarser representation. Quantization is applied block by block, and incoherent processing is a transformation applied before quantization so that the reduced precision does less damage to the result.

Results

On H100 GPUs the method is 1.5–2.0x faster than FlashAttention-2, with FP16 reaching up to 740 TFLOPs/s — 75% utilization — and FP8 approaching 1.2 PFLOPs/s ^[1]. On the H100 SXM5 GPU specifically, the FP16 version gives a 1.5–2.0x speedup in the forward pass and 1.5–1.75x in the backward pass ^[1]. Speed is not bought with accuracy: the FP8 version achieves 2.6x lower numerical error than a baseline FP8 attention ^[1].

Why it matters

If attention is what limits large language models and long-context applications ^[1], then making that single layer faster on the same hardware is one of the most direct ways to make longer contexts practical. The utilization figure matters as much as the speedup: at 75% in FP16, most of the machine's arithmetic capacity is doing useful work, where the previous version left roughly two thirds of it idle ^[1]. The numerical-error result addresses the usual objection to low precision, namely that it costs accuracy ^[1].

Authors and scope

The paper was submitted to arXiv on July 11, 2024, by Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani and Tri Dao ^[1]. Its techniques are designed for Hopper-generation GPUs, and the reported results are measured on H100 hardware ^[1].