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

FLASH ATTENTION sped up attention on GPUs but achieved only 35% utilization on H100 GPUs. To address this, FLASH ATTENTION-3 introduces three techniques for Hopper GPUs: warp-specialization to overlap computation and data movement by assigning warps to producer and consumer roles; interleaving operations to combine block-wise matrix multiplication (matmul) and softmax operations for asynchronous execution; and support for FP8 precision using block quantization and incoherent processing to leverage hardware capabilities. These techniques result in a 1.5-2.0× speedup with FP16, up to 1.2 PFLOPs/s with FP8, and 2.6× lower numerical error than baseline FP8 attention, improving computational efficiency and hardware utilization on modern GPUs.

The paper builds on the existing work on FlashAttention, and introduces an advanced method to accelerate the attention mechanism in Transformer model. It specifically targets the newer GPU architectures, like NVIDIA H100, by exploiting asynchrony in Tensor Cores and Tensor Memory Accelerators. There are three techniques that are proposed: 1) producer-consumer asynchrony across warps; 2) compute asynchrony within consumer warps; 3) hardware accelerated low precision GEMM. The extensive empirical results show impressive performance, with significant speedups compared to state-of-the-art methods, and with reduced numerical errors.

This is a system paper that introduces FlashAttention-3, an optimized version of FlashAttention for NVIDIA's SM90+ GPUs. The key contributions (summarized by the authors) are:

• Taking advantage of warp specialization in NVIDIA SM90+ and designing a producer-consumer computation paradigm to allow better intra-CTA overlapping between producers and consumers;
• Designing 2-stage software pipeline to allow better overlapping within consumers;
• Designing on-the-fly layout transformation to support FP8 WGMMA and using incoherent processing to reduce quantization error by 3x.

This paper presents FlashAttention-3, which speedup the commonly-used attention operator on Hopper GPUs. The paper proposes to leverage the asynchronous execution of the Tensor Cores and Tensor Memory Acceleator to better utilize the GPU hardware. Specifically, the paper proposes three techniques: (1) It overlaps the data movement and GEMM computation by warp specialization; (2) It proposes to overlap the GEMM and softmax operations by interleaving these operations; and (3) It leverages FP8 tensor core by block quantization. The experimental results show that FlashAttention-3 can achieve superior speedup by 1.5-2.0x with FP16 and can outperform vendor libraries.

We thank the reviewers for their enthusiastic support, their careful read of the paper, and their thoughtful questions and suggestions. We are very happy that the reviewers find the paper “has immediate impact to the community”, the ideas “novel and insightful”, and the writing “clear and easy to understand”.

Since the submission in May, we have made FA3 even faster, more accurate, easier to use, and more general:

• New optimizations: (1) A persistent scheduler to start loading the Q, K, V for the next subproblem while writing out the output O of the current subproblem, thus speeding up the cases of small seqlen (2) Optimized for the case where sequences in the batch have variable lengths (e.g. in Llama 3 training) (3) Reduced the numerical error for GQA/MQA by accumulating the gradient dK and dV of different heads in FP32, before converting the sum to BF16. This brings 2-3x smaller numerical error compared to the standard method of computing dK and dV for different heads in BF16, then summing the gradients.
• Speed: FA3 version in August is 6-25% faster, reaching up to 700 TFLOPS for headdim 128 FP16 (vs 650 in May), 800 TFLOPS for headdim 256 FP16 (vs 640 in May), and 1300 TFLOPS for headdim 256 FP8 (vs 1230 in May). This is thanks to a technique called pingpong scheduling that better overlaps the GEMM of one warpgroup with the softmax of another warpgroup. For FP16/BF16, attention speed from FA3 is now on par with matmul speed from cuBLAS (arguably the most optimized operation) for similar sizes (700-770 TFLOPS), suggesting that we are using the hardware as efficiently as possible. This is another validation that asynchrony can play a major role in optimizing for modern hardware.
• Integration with other libraries: we are working with PyTorch developers to integrate FA3 to PyTorch to benefit the largest number of researchers and engineers. We are also working with inference and training libraries (Hugging Face, vLLM, Megatron-LM, TransformerEngine, DeepSpeed) to help with their integration effort to speed up transformer training and inference with FA3.
• Generality: We are working with Triton and Jax/XLA developers to implement some of the techniques in FA3 in the Triton & XLA compilers. We are also collaborating with other researchers and engineers on hardware other than H100 (elaborated below).

We now respond to some common questions from the reviewers.

1. Generality: The techniques developed in FA3 are not limited to Nvidia Hopper GPUs. As mentioned in the intro, asynchrony and low-precision are the general trend of AI accelerators, due to Amdahl's law. Amdahl’s law states that the overall performance improvement gained by optimizing a single part of a system is limited by the fraction of time that the improved part is actually used. As tensor cores in hardware accelerate matmuls exponentially with each new generation of hardware, the serial sections of the code (softmax and causal masking in attention) can dominate the runtime. The only way to overcome these latencies, other than adding expensive hardware units that compute them, is to overlap their execution with other concurrent work. We chose to validate our algorithms and ideas on Hopper GPUs due to good hardware and software support (async barriers, CUTLASS), but these ideas also apply to other modern accelerators. In fact, we are collaborating with AMD engineers to implement the overlapping of GEMM and Softmax on AMD GPUs such as MI300X, as well as Google researchers on a version of FA3 on TPUs. We are also exploring FP8 support for Ada GPUs (e.g. 4090) to potentially double the attention throughput on these consumer cards.
2. Impact: With FA3, we have significantly sped up attention, one of the two main layers in Transformers. Attention speedup can have a major impact when the sequences are long, which is increasingly common (e.g. Llama 3.1 has 128k context length). In a LLM inference production setting (with an already very optimized inference engine), we have measured that switching from FA2 to FA3 yields up to 30-40% speedup in time to first token for Llama-3 405B. This is a major improvement for inference services running on hundreds or thousands of GPUs. For 405B at 128K sequence length, attention takes about 50% of the time and MLP takes 50% of the time, and making attention 2x faster will yield about 1.3x speedup. One can combine FA3 with other techniques to speed up the MLP (e.g. sparsity, FP6/FP4) to get even higher speedup.
3. Interoperability: FA3 is a drop-in replacement for FA2, save for the FP8 case where the version of FA3 in May requires V to have layout BHDS (sequence length dimension must be contiguous), due to the constraint of the FP8 tensor cores. We have since developed a variant of the FP8 kernel that supports V input in the standard BSHD layout (with the head dimension as the contiguous mode). This is important to integrate FP8 attention into standard libraries, such as PyTorch and vLLM, and also circumvents a technical challenge with variable sequence length in terms of loading memory addresses not aligned to 16 bytes. Recall from the paper that the FP8 WGMMA instruction only supports k-major operand B, so this variant necessitates doing an "in-kernel" transpose on tiles of V after they are loaded from GMEM to SMEM. Moreover, in order to overlap this transpose with other operations to minimize its impact on speed, we change the kernel scheduling. In short, we take advantage of special LDSM-STSM instructions that both minimize register usage and support transposition for SMEM <=> RMEM copy in order to place the in-kernel transpose of V in the producer, thereby further leveraging warp specialization with warpgroup register reallocation. Our implementation achieves comparable performance to cuDNN, exceeding it for headdim 64 and matching for headdim 128/256.

Overall we are excited about generalizing our techniques to other hardwares to unlock new use cases for long context models.

The work introduces FlashAttention 3, an optimized version of FlashAttention for modern GPUs, which is a key building block for efficient transformer models. While a main focus is on NVIDIA's Hopper GPUs which currently lack such efficient kernels, the innovations of this paper are also promising and valuable for GPUs from other vendors and various other compute backends. Compatibility with new lower-precision floating point formats such as FP8 is another strength of this work, which is very valuable for the wider community.