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

🌈 Abstract

The paper presents FlashAttention-3, a new algorithm for speeding up attention on Hopper GPUs. The key contributions are:

1. Exploiting asynchrony of the Tensor Cores and TMA to overlap computation and data movement via warp-specialization.
2. Overlapping block-wise matmul and softmax operations by pipelining the computations across iterations.
3. Leveraging hardware support for FP8 low-precision computation through block quantization and incoherent processing to further improve performance.

The authors demonstrate that FlashAttention-3 achieves 1.5-2.0x speedup over the previous FlashAttention-2 algorithm on H100 GPUs, reaching up to 740 TFLOPs/s (75% utilization) in FP16 and close to 1.2 PFLOPs/s in FP8. They also show that FP8 FlashAttention-3 achieves 2.6x lower numerical error compared to a baseline FP8 attention implementation.

🙋 Q&A

[01] Techniques to speed up attention on Hopper GPUs

1. What are the three main techniques used in FlashAttention-3 to speed up attention on Hopper GPUs? The three main techniques are:

• Exploiting asynchrony of the Tensor Cores and TMA to overlap computation and data movement via warp-specialization.
• Overlapping block-wise matmul and softmax operations by pipelining the computations across iterations.
• Leveraging hardware support for FP8 low-precision computation through block quantization and incoherent processing.

2. How does FlashAttention-3 achieve speedup compared to previous methods? FlashAttention-3 achieves 1.5-2.0x speedup over the previous FlashAttention-2 algorithm on H100 GPUs, reaching up to 740 TFLOPs/s (75% utilization) in FP16 and close to 1.2 PFLOPs/s in FP8.

3. How does FlashAttention-3 improve numerical accuracy in FP8 precision? FlashAttention-3 uses block quantization and incoherent processing techniques to reduce numerical error in FP8 precision, achieving 2.6x lower error compared to a baseline FP8 attention implementation.

[02] Comparison to prior work

1. How does FlashAttention-3 differ from the previous FlashAttention-2 algorithm? The key differences are:

• FlashAttention-3 exploits asynchrony and warp-specialization to overlap computation and data movement.
• It overlaps block-wise matmul and softmax operations through a 2-stage pipelining approach.
• It includes modifications to support FP8 precision, including layout transformations and techniques to improve numerical accuracy.

2. How does FlashAttention-3 compare to other attention optimization methods like sparse and low-rank approximations? The paper states that while sparse and low-rank attention approximation methods can scale to longer sequences, they typically do not offer the same model quality as standard attention. FlashAttention-3 focuses on optimizing the exact attention computation rather than approximating it.

3. How does FlashAttention-3 relate to alternative attention-free architectures like RWKV, H3, and Mamba? The paper notes that while these alternative architectures aim to address the limitations of attention, they still employ many layers of attention. The techniques developed in FlashAttention-3 to speed up attention computation can also benefit these alternative architectures.