Flash Attention 3 技术探讨圆桌
输入“/”快速插入内容
Flash Attention 3 技术探讨圆桌
Flash Attention 3 技术探讨圆桌
用户1672
用户8013
用户2055
用户1902用户32302025年11月29日修改
#腾讯会议：354-5795
分享主持人：某上市公司算法工程师、书生·浦语 SIG
FlashAttention-3:利用异步和低精度实现快速且准确的注意力机制
1.
论文翻译
作者信息：
Jay Shah¹, Ganesh Bikshandi¹, Ying Zhang², Vijay Thakkar³⁴, Pradeep Ramani³, 和 Tri Dao⁵⁶
¹Colfax Research ²Meta ³NVIDIA ⁴Georgia Tech ⁵Princeton University ⁶Together AI
摘要
注意力机制作为无处不在的Transformer架构的核心层，是大型语言模型和长上下文应用的瓶颈。FlashAttention提出了一种通过最小化内存读/写来加速GPU上的注意力计算的方法。它尚未利用最新硬件中的新功能，FlashAttention-2仅在H100 GPU上实现了35%的利用率。我们开发了三种主要技术来加速Hopper GPUs上的注意力计算：利用Tensor Cores和TMA的异步性来(1)通过warp专门化重叠整体计算和数据移动，(2)交错块级matmul和softmax操作，以及(3)块量化和不连贯处理，这些技术利用了FP8低精度的硬件支持。我们证明了我们的方法FlashAttention-3在H100 GPUs上实现了1.5-2.0倍的加速，FP16达到740 TFLOPs/s（75%利用率），FP8接近1.2 PFLOPs/s。我们验证了FP8 FlashAttention-3比基准FP8注意力实现的数值误差低2.6倍。
1 引言
对于Transformer架构[56]，注意力机制构成了主要的计算瓶颈，因为计算查询和键的自注意力分数具有序列长度的二次方缩放。扩展注意力以适应更长的上下文将解锁新的能力（对多个长文档进行建模和推理[24, 42, 49]以及大型代码库中的文件[30, 47]），新的模态（高分辨率图像[11]，音频[23]，视频[25]），以及新的应用（用户与长历史记录的交互[51]，具有长期视野的代理工作流[59]）。这在长上下文领域引发了对加速注意力计算的重大兴趣，包括通过近似[14, 27, 54]和软硬件优化([17, 29, 44])，甚至替代架构[22, 41, 53]。译者注：比如最近的TTT
在这项工作中，我们基于Dao等人[17]的工作，开发了将GPU执行模型和硬件特性知识整合到其高级设计中的精确注意力算法。在[17]中，Dao等人引入了FlashAttention，这是一种新的并行化注意力的分块策略，通过将所有注意力操作融合到单个GPU内核中，消除了中间读/写到慢速全局内存的需求。
Dao[15]将算法重构为FlashAttention-2，以在序列长度维度上并行化，并在键值矩阵的块上执行前向传播的内部循环，从而提高了GPU上的工作占用率和分布。然而，我们观察到FlashAttention-2在较新的GPU上相对于优化的矩阵乘法（GEMM）内核仍然存在较低的利用率，例如在Hopper H100 GPU上仅为35%。部分原因可能是实现级别的差异，比如在针对Tensor Cores时没有使用Hopper特定指令而是使用Ampere指令。
更根本地说，FlashAttention-2的算法遵循了一个简化的同步模型，在其设计中没有明确使用异步和低精度。异步是硬件专门化的结果，用于加速ML工作负载中最重要的操作：专用硬件单元执行矩阵乘法（Tensor Cores）或内存加载（Tensor Memory Accelerator - TMA），与执行逻辑、整数和浮点计算的其他CUDA核心分离。低精度如Hopper中的FP8和Blackwell中的FP4，延续了FP16（Pascal在2017年）和BF16（Ampere在2020年）的趋势，是一种proven技术，可以在相同功耗和芯片面积下获得更高的吞吐量。我们在§2.2中回顾了Hopper在这些方向上提供的功能。技术挑战在于重新设计FlashAttention-2以利用这些硬件特性：异步要求在matmul和softmax之间重叠计算，尽管其中一个依赖于另一个的输出，而低精度需要注意最小化量化误差，特别是在LLMs中的异常特征情况下[20, 52]。
为此，我们提出FlashAttention-3，它贡献并综合了三个新想法，以进一步提高新GPU架构上的性能：¹
1.
生产者-消费者异步：我们定义了一个warp专门化的软件流水线方案，利用数据移动和Tensor Cores的异步执行，将数据的生产者和消费者分割到不同的warps中，从而扩展了算法隐藏内存和指令发布延迟的能力。
2.
在异步块级GEMMs下隐藏softmax：我们将softmax中涉及的相对低吞吐量的非GEMM操作（如浮点乘加和指数）与GEMM的异步WGMMA指令重叠。作为这一过程的一部分，我们重新设计了FlashAttention-2算法，以规避softmax和GEMMs之间的某些顺序依赖。例如，在我们算法的2阶段版本中，当softmax在分数矩阵的一个块上执行时，WGMMA在异步代理中计算下一个块。
3.
硬件加速的低精度GEMM：我们调整前向传播算法以允许针对FP8 Tensor Cores进行GEMM，几乎将测量的TFLOPs/s翻倍。这需要弥合WGMMA在FP32累加器和FP8操作矩阵方面的不同布局一致性要求，假设这些矩阵在内存中已经布局好。我们使用块量化和不连贯处理技术来缓解转移到FP8精度导致的准确性损失。
为了实证验证我们的方法，我们在H100 SXM5 GPU上对FlashAttention-3进行了一系列参数的基准测试，并显示（1）FP16在前向传播中实现了1.5-2.0倍的加速（达到740 TFLOPs/s），在后向传播中实现了1.5-1.75倍的加速，（2）FP8达到接近1.2 PFLOPs/s，以及（3）对于长序列长度，
P32中。此外，使用块量化和不连贯处理的FP8 FlashAttention-3在存在异常特征的情况下比标准注意力精确2.6倍。
我们以许可证³开源FlashAttention-3，并计划将其集成到PyTorch和Hugging Face库中，以惠及最大数量的研究人员和开发人员。
注：
¹我们在 NVIDIA 的 Hopper 架构背景下描述我们的结果。然而，我们的算法适用于任何具有足够强大的异步执行和低精度能力的 GPU 架构。
²更准确地说，对于头维度 64，FlashAttention-3 FP8 领先，而对于头维度 128 和 256，在没有因果掩码的情况下与之相当，在有因果掩码的情况下落后。