《FlashAttention-3：以异步与低精度实现快速准确注意力》

问题：注意力是瓶颈

注意力（attention）是 Transformer 架构的核心层，这篇论文的出发点是一个判断：拖住大语言模型（large language model）和长上下文应用的瓶颈在注意力本身，而不是它周围的其他层^[1]。原因出在结构上——注意力让序列里每个词元（token）与其他词元交换信息，计算量随序列长度增长；这一层实现得好不好，直接决定模型能负担多长的上下文。

英伟达 H100 这类 GPU 做矩阵乘法极快，但一个注意力层并不是一次干净的矩阵乘法。它是一串矩阵乘法与 softmax 归一化的交替，中间结果还要在不同层级的内存之间搬运。如果实现让算力单元停下来等数据，机器的一大半就空转掉了。论文把这个差距量了出来：上一代方法 FlashAttention-2 在 H100 上的利用率只有 35%^[1]。

三项面向 Hopper 的技术

论文提出了三项专门针对 Hopper 架构 GPU 的技术^[1]。

第一项是通过「线程束特化」（warp-specialization）让计算与数据搬运重叠进行^[1]。GPU 程序里的线程按组运行，一组称为一个线程束（warp）；线程束特化就是把不同的线程束分派去做不同的事，一部分专门负责取数据，另一部分专门对已经到位的数据做运算。数据搬运于是藏在计算背后，而不是把计算卡住。

第二项是把分块的矩阵乘法与 softmax 运算交错执行^[1]。注意力是一块一块算出来的，因为没有任何一个处理器能一次装下整条长序列。softmax 需要维护分数的动态统计量，这一步没法只靠矩阵乘法完成，否则会让高速的矩阵单元闲着。把两类工作交错开，归一化在算的时候，矩阵单元仍然有活干。

第三项是块量化（block quantization）配合不相干处理（incoherent processing），利用硬件的 FP8 支持^[1]。FP8 用 8 位而不是 16 位存储每个数，同样时间能吞吐更多数值，代价是表示更粗糙。量化按块进行，不相干处理则是在量化之前对数据做一次变换，让降低精度对最终结果的影响更小。

结果

在 H100 GPU 上，这一方法比 FlashAttention-2 快 1.5–2.0 倍，其中 FP16 最高达到 740 TFLOPs/s，即 75% 的利用率，FP8 则接近 1.2 PFLOPs/s^[1]。具体到 H100 SXM5 这块卡，FP16 版本在前向传播中带来 1.5–2.0 倍加速，反向传播中带来 1.5–1.75 倍加速^[1]。速度并不是拿精度换来的：FP8 版本的数值误差比作为基线的 FP8 注意力低 2.6 倍^[1]。

为什么重要

如果限制大语言模型和长上下文应用的就是注意力^[1]，那么在同样的硬件上把这一层做快，就是让长上下文真正可用的一条最直接的路径。利用率这个数字和加速比同样重要：FP16 下达到 75%，意味着机器的算力大部分在做有用的工作，而上一代方法有大约三分之二被闲置^[1]。误差那一项结果回应的是人们对低精度的常见顾虑，即精度下降会带来准确度损失^[1]。

作者与范围

论文于 2024 年 7 月 11 日提交到 arXiv，作者为 Jay Shah、Ganesh Bikshandi、Ying Zhang、Vijay Thakkar、Pradeep Ramani 和 Tri Dao^[1]。三项技术是为 Hopper 架构 GPU 设计的，报告的成绩也都在 H100 硬件上测得^[1]。