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

注意力机制是无处不在的Transformer架构中的核心层，是大型语言模型和长上下文应用的瓶颈。FlashAttention提出了一种通过最小化内存读/写来加速GPU上的注意力机制的方法。然而，它尚未利用最新硬件中存在的新功能，FlashAttention-2仅在H100 GPU上实现了35%的利用率。我们开发了三种主要技术来加速Hopper GPU上的注意力机制：利用张量核心和TMA的异步性，通过warp-specialization和交错块状matmul和softmax操作来重叠整体计算和数据移动，以及利用硬件支持FP8低精度的块量化和不连贯处理。我们证明了我们的方法FlashAttention-3在H100 GPU上实现了1.5-2.0倍的加速，FP16可达到740 TFLOPs/s（75%利用率），FP8接近1.2 PFLOPs/s。我们验证，FP8 FlashAttention-3的数值误差比基线FP8注意力低2.6倍。

FlashAttention-3论文的主要目的是解决在大型语言模型和长上下文应用中，注意力机制是GPU性能瓶颈的问题，试图提高GPU上的注意力机制性能。

FlashAttention-3的关键思路是利用Tensor Cores和TMA的异步性来最小化内存读/写，通过warp-specialization和交错块状matmul和softmax操作来重叠整体计算和数据移动，并使用硬件支持的FP8低精度实现块量化和不连贯处理。

FlashAttention-3方法在H100 GPU上实现了1.5-2.0倍的加速，FP16最高可达到740 TFLOPs/s（75%利用率），FP8最高可达到近1.2 PFLOPs/s。与基线FP8注意力相比，FP8 FlashAttention-3的数值误差降低了2.6倍。

近期在该领域的相关研究包括：《Attention is not all you need: Pure attention loses rank doubly exponentially with depth》、《Reformer: The Efficient Transformer》、《Linformer: Self-Attention with Linear Complexity》等。