Flash Attention V3 发布，大模型进化再次加速

最新推荐文章于 2025-04-06 22:21:09 发布

强化学习曾小健

最新推荐文章于 2025-04-06 22:21:09 发布

阅读量1.7k

点赞数 21

文章标签： windows

本文链接：https://blog.csdn.net/sinat_37574187/article/details/140567629

版权

Flash Attention V3 发布，大模型进化再次加速

翻译：某上市公司算法工程师

原文链接：https://tridao.me/publications/flash3/flash3.pdf

欢迎对FlashAttention-3感兴趣的同学

全文约 8000 字，预计阅读时间 12 分钟

作者信息：

Jay Shah¹, Ganesh Bikshandi¹, Ying Zhang², Vijay Thakkar³⁴, Pradeep Ramani³, 和 Tri Dao⁵⁶

¹Colfax Research ²Meta ³NVIDIA ⁴Georgia Tech ⁵Princeton University ⁶Together AI {jayshah,ganesh}@colfax-intl.com, yingz@meta.com, {vithakkar,prraman}@nvidia.com, tri@tridao.me

摘要

注意力机制作为无处不在的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倍。

引言

对于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架构上的性能：¹

生产者-消费者异步：我们定义了一个warp专门化的软件流水线方案，利用数据移动和Tensor Cores的异步执行，将数据的生产者和消费者分割到不同的warps中，从而扩展了算法隐藏内存和指令发布延迟的能力。
在异步块级GEMMs下隐藏softmax：我们将softmax中涉及的相对低吞吐量的非GEMM操作（如浮点乘加和指数）与GEMM的异步WGMMA指令重叠。作为这一过程的一部分，我们重新设计了FlashAttention-2算法，以规避softmax和GEMMs之间的某些顺序依赖。例如，在我们算法的2阶段版本中，当softmax在分数矩阵的一个块上执行时，WGMMA在异步代理中计算下一个块。
硬件加速的低精度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）对于长序列长度，FP16优于FP8，与NVIDIA的cuDNN库中最先进的注意力实现相比具有竞争力²。我们还验证了FP16 FlashAttention-3产生的数值误差与FlashAttention-2相同，并且优于标准注意力实现，因为中间结果（例如，softmax重缩放）保持在FP32中。此外，使用块量化和不连贯处理的FP8 FlashAttention-3在存在异常特征的情况下比标准注意力精确2.6倍。

我们以许可证³开源FlashAttentio