返回文章 apply CMS 文章

FlashAttention-4：面向非对称硬件扩展的算法与内核流水线协同设计

FlashAttention-4 通过新的流水线、2-CTA MMA 模式和混合软硬件指数计算，解决了 Blackwell GPU 上注意力机制的性能瓶颈。

Together AI

FlashAttention-4Blackwell GPU注意力机制内核优化

成长分 / 100 70 综合收获、行动、留存与影响

为什么值得读了解如何应对 GPU 张量核心吞吐量远超其他资源（如共享内存带宽、SFU）的非对称扩展趋势。

学习 FlashAttention-4 中创新的前向/反向流水线、2-CTA MMA 和条件在线 softmax 重新缩放等关键技术。

关键洞察

Blackwell GPU 上注意力性能瓶颈在于前向的 softmax 指数运算和反向的共享内存流量，而非张量核心。
前向传递通过乒乓调度、软件仿真指数和条件重新缩放，将 softmax 与矩阵乘法重叠。
反向传递利用张量内存（TMEM）存储中间结果，并通过 2-CTA MMA 减少共享内存流量和全局原子归约。

转成行动

深入阅读

正文与原文对照

原文保真覆盖：全文原文字符：18129

抽象彩色背景，文字为 FlashAttention-4，小标签写着 Research。

像 Blackwell GPU 这样的现代加速器延续了非对称硬件扩展的趋势，其中张量核心吞吐量的增长速度远快于其他资源，如共享内存带宽、用于指数等超越运算的特殊函数单元（SFU）以及通用整数和浮点 ALU。例如，从 Hopper H100 到 Blackwell B200，BF16 张量核心吞吐量从 1 PFLOPs 增加到 2.25 PFLOPs，而 SFU 数量和共享内存带宽保持不变。

这种扩展不对称性对优化 Blackwell 架构上的注意力等复杂内核产生了深远影响。注意力机制的核心包括两个 GEMM（ $S=Q \cdot K^T$ 和 $O=P \cdot V$ ），中间有 softmax；实际上，它还涉及大量的管道和簿记工作：数据移动、同步、布局变换、逐元素操作、调度、掩码等。

一种天真的观点可能认为，GEMM 的速度完全控制着内核性能，并且可以有效地忽略这些其他注意力组件，至少在一阶近似下如此。然而，对 B200 进行“馈送与速度”分析实际上显示了相反的情况：主要性能瓶颈不在于张量核心执行 MMA 的速度，而在于（a）前向计算中用于 softmax 指数的 SFU 单元，以及（b）反向计算中的共享内存流量。

在这篇博客文章中，我们提出了 FlashAttention-4，一种算法和内核协同设计，最大限度地重叠矩阵乘法与其他资源瓶颈。在 B200 上使用 BF16，它达到了高达 1605 TFLOPs/s（71% 利用率），比 cuDNN 版本 9.13 快 1.3 倍，比 Triton 快 2.7 倍。

我们的主要算法和内核协同设计思路如下：

新的流水线以实现最大重叠：新的前向和反向软件流水线，利用 Blackwell 完全异步的 MMA 和更大的 tile 尺寸，重叠张量核心、softmax 指数和内存操作。前向（FWD）传递：通过 FMA 单元上的多项式近似实现指数函数的软件仿真，以缓解指数瓶颈，加上条件在线 softmax 重新缩放。*反向（BWD）传递*：将中间结果存储在张量内存中以减轻共享内存流量，结合 Blackwell 新的 2-CTA MMA 模式进一步减少共享内存流量并减少一半的原子归约，以及额外支持确定性执行模式以实现可重现训练。*调度*：新的 tile 调度器，以减轻因果掩码和可变序列长度带来的负载不平衡。

Blackwell 上的新硬件特性

张量内存（TMEM）：在 B200 上，每个 148 个 SM 拥有 256 KB 的 TMEM，这是一个连接到张量核心的片上暂存器，用于 warp 同步中间存储。完全异步的第五代张量核心：tcgen05.mma

是异步的，并累积在 TMEM 中。对于 BF16 和 FP16，最大的单 CTA UMMA 块为 128×256×16，大约是最大 Hopper WGMMA 原子的 2 倍。UMMA 由单个线程启动，减轻了寄存器压力，使得更大的块和更深的流水线成为可能，而不会出现 Hopper warpgroup MMA 的溢出痛点。这也使得 warp 专业化更加可行，一些 warp 移动块，而其他 warp 发出 MMA，以将矩阵乘累加与 softmax 和内存流量重叠。tcgen05.mma

还可以从 TMEM 获取操作数 A。2-CTA MMA：Blackwell 可以在同一集群中的一对 CTA 上执行一个 UMMA，跨越两个对等 CTA 的 TMEM。领导 CTA 中的一个线程启动 MMA，但两个 CTA 在其执行期间必须保持活动状态。这将 MMA 块维度扩展到 256×256×16，通过在对之间拆分 M 和 N，减少了冗余流量并降低了每个 CTA 的占用。CTA 组大小（1 或 2）在内核中的 TMEM 和张量核心操作中必须保持不变。

馈送与速度

对于 M=N=D=128

馈送（B200 上，每 SM）：

张量核心（BF16）： $\frac{8192 \text{ 操作}}{周期}$
指数单元： $\frac{16 \text{ 操作}}{周期}$
共享内存流量： $\frac{128 \text{ 字节}}{周期}$

速度（每块的时钟周期数）：

前向（2 个 MMA + MN 指数）
张量核心： $1024$
指数： $1024$
SMEM： $768$
后向（5 个 MMA + MN 指数）：1-CTA
张量核心： $2560$
指数： $1024$
SMEM： $3328$

要点：前向受计算和指数瓶颈限制，后向受共享内存带宽瓶颈限制。因此，我们在前向传递中将 softmax 与 MMA 重叠，并在后向传递中减少共享内存流量。

前向传递：带有条件重新缩放的新型 softmax 流水线

前向传递有两个矩阵乘法，QK^T

和 PV

。在 Blackwell 上，张量核心变得更快，但指数单元（MUFU.EX2）没有。因此，softmax 不再是“两个矩阵乘法之间的东西”，而是一个必须仔细流水线化的瓶颈。

前向传递简述：

乒乓调度：每个 CTA 有 $2x$ 个 Q 块和 $2x$ 个 O 块：最大化 MMA 和 Softmax 之间的重叠2x softmax warp 组：每个块执行 softmax，并同步以避免计算指数时重叠 $2^x$ 的软件仿真：将指数计算分布在硬件的 MUFU 和软件仿真的 FMA 上分阶段将 P 存储在 TMEM 中：缓解寄存器压力

校正 warp 组：指定“校正”warp 组执行重新缩放，以将其从关键路径中移除在线 softmax（条件）重新缩放：减少重新缩放频率以最小化非矩阵乘法操作

流水线：乒乓 Q 块加上专用校正阶段

FlashAttention-4 为每个 CTA 计算两个查询块—— $Q^H$ 和 $Q^L$ ——每个覆盖 128 个查询令牌，并以乒乓调度交替它们。

Blackwell 改变了 softmax 映射。S = QK^T

的累加器块为 128×128，位于张量内存中；然而，当读入寄存器时，根据硬件规定，块的划分是每行一个线程。我们使用两个 128 线程的 warp 组，每个对应一个 Q 块，每个 softmax warp 组执行以下操作序列：

每个线程从张量内存加载一行 128 个元素的 S 到寄存器中
规约 rowmax 和 rowsum
使用可调参数，决定 128 个元素中哪部分使用硬件的 MUFU 与软件模拟的 $e^x$
计算 P = softmax(S) 并转换为 BF16 精度
分阶段将 P 存回张量内存以缓解寄存器压力（而不是同时持有 S 的 128 个元素和 P 的 64 个（BF16）元素）
一旦存储了 P 的 $\frac{3}{4}$ 块，立即触发对应的 PV 矩阵乘法

关键细节在于 exp 是瓶颈部分。我们显式同步两个 softmax warp 组，使它们不同时计算 exp，从而减少 MUFU 争用。

为了将重新缩放移出关键路径，内核将其分配给一个专用的 warp 组。校正 warp 组计算：

仅在最大跳跃较大时重新缩放：
$O_j =\begin{cases}\exp(m_{j-1}-m_j)\,O_{j-1} + \exp(S_j-m_j)\,V_j, & \text{if } m_j - m_{j-1} > \tau,\\O_{j-1} + \exp(S_j-m_{j-1})\,V_j, & \text{otherwise.}\end{cases}$
在迭代结束时应用最终归一化 $O_{final} = \frac{O}{l_{final}}$
可选地计算并存储 LSE

最后我们仍然使用真实的最终统计量进行归一化，因此跳过小的重新缩放步骤可以在保持最终输出的同时，从关键路径中删除许多向量计算。我们在 warp 粒度上做出决策以避免分支发散。

更快的指数：在 MUFU.EX2 和 FMA（软件模拟）之间分配 $2^x$

Softmax 需要大量指数运算，而 MUFU 吞吐量远低于张量核心吞吐量。FlashAttention-4 通过同时运行硬件 MUFU.EX2 路径和软件模拟的 exp2（使用原本可能未被充分利用的 FMA 单元）来提高有效 exp 吞吐量。

范围缩减（Cody-Waite）： 我们使用经典的 Cody-Waite 范围缩减技术将指数计算分解为整数部分和小数部分： $2^x = 2^{n} \cdot 2^{f}$ 。在 IEEE 754 float32 中，乘以 $2^{n}$ 只是指数更新。

$2^{x_{frac}}$ 的多项式近似（Horner 方法）： 为了近似 $2^{f}$ ，我们将其重写为 Horner 形式以便高效求值。

2^{x_{\mathrm{frac}}} \approx p_0 + p_1 x_{\mathrm{frac}} + p_2 x_{\mathrm{frac}}^{2} + p_3 x_{\mathrm{frac}}^{3}

系数 p0 = 1.0、p1 ≈ 0.6951、p2 ≈ 0.2276、p3 ≈ 0.0771 使用 Sollya 软件包选择，以最小化 $[0, 1)$ 上的相对近似误差。

指数位移位与加法： 最后一步是将整数部分 n 和小数近似 $2^{f}$ 组合成 $2^{x} \approx 2^{n}\cdot 2^{f}$ 。由于 $2^f \in[1,2)$ 的 float32 指数为 127，乘以 $2^{n}$ 只需将整数 n 移位到指数域，然后加上 $2^{f}$ 的尾数位。

反向传播：共享内存流量占主导

优化 FlashAttention 反向传播就像把一张过大的地毯塞进房间：压平一个角落，另一个角落又翘起来。反向传播的计算量约为前向传播的张量核心工作的 2.5 倍，它串联了五个 MMA 操作以重新计算 S，并运行 dQ、dK、dP 和 dV 的 QK 和 PV 梯度 MMA，以及 P 和 dS 的逐元素工作。在 Blackwell 上，FLOPs 不是反向传播的限制因素；共享内存带宽才是。

流水线：将 MMA 与 softmax 重叠

softmax

Hopper 时代的 FlashAttention-3 将 MMA 累加器保留在寄存器中，因此寄存器压力通常迫使采用更串行的调度。在 Blackwell 上，累加器位于 TMEM 中，这使得在 CUDA 核心处理 P 和 dS 的逐元素工作时，保持多个 MMA 同时进行变得可行。由于在我们的屋顶线模型中，指数运算的吞吐量与两个 MMA 相当，因此隐藏它是值得的。

关键的重叠很简单：当我们计算 tile j 的 softmax 时，我们已经为 tile j−1 发出了 dK 和 dQ MMA。

为了减少共享内存流量，反向传播相对于前向传播以转置的 tile 重新计算 S 和 P，因此中间结果已经是 $S^T$ 和 $P^T$ 。然后我们可以将 $P^T$ （以及后来的 $dS^T$ ）直接存储在 TMEM 中，其布局分别与 dV 和 dK MMA 消耗的操作数 A 完全一致。

TMEM 无法同时容纳五个完整的累加器和中间结果，因此 FA4 跨阶段重用 TMEM 列：S 和 P 共享一组列，dP、dS 和 dQ 共享另一组。

2-CTA 反向传播：减少共享内存流量和全局原子加操作

共享内存流量。 即使采用了改进的流水线，并且十个 GEMM 操作数中有两个保存在张量内存中，反向传播仍然受限于共享内存带宽。我们通过 Blackwell 2-CTA MMA 模式缓解了这一问题，该模式将输出累加器在 CTA 对之间进行分区。对于 M=256 且 N=K=128，两个 CTA 作为一个 tile 协作：每个 CTA 暂存操作数 B 的一半，并仅保留自己的累加器切片。这大致将操作数 B 的共享内存流量减半。

归约轴冲突。 我们在五个反向 GEMM 中使用 M=256 且 N=K=128 的 MMA tile 来减少 B 的流量，但 dQ MMA 的性质引入了不匹配。在 FlashAttention 反向传播中，每个 CTA 拥有一个固定的 KV tile（外层循环在 N 个 CTA 上并行化），并在内层循环中遍历 M 个 tile。dQ 更新在外层循环中对 KV 序列进行归约。2-CTA MMA 分割的是输出 tile，而不是归约，而 dQ 的归约维度是 N，该维度已经在 CTA 对之间分割。每个 CTA 仍然需要为其拥有的行进行完整的归约。

解决方案：DSMEM 交换。 我们通过使用集群内的分布式共享内存在两个 CTA 之间交换一半的 dS 来解决这个问题。这重新打包了 dS，使其沿非归约轴进行分区：每个 CTA 拥有 M/2 行，同时持有完整的 2N 归约。每个 CTA 的 dQ MMA 变为 (M/2, 2N)(2N, d)，在张量内存中累积一个 (M/2, d) 的 tile。在 2-CTA 模式下，S、dP、dV 和 dK 的 MMA 保持 M=256，而 dQ 使用 M=128，归约加倍为 2N=256。然后我们重新排序流水线以隐藏 DSMEM 延迟：在当前 tile 的 dP 计算之前，先计算前一个 tile 的 dQ。由于 dQ tile 与 P 一起适合 TMEM，它可以重用 S 使用的 TMEM 区域，因此 dP 和 dQ 不再像 1-CTA 模式那样共享一个区域。通过这种排序，当前 tile 的逐元素 dS 与上一次迭代的 dQ MMA 重叠。

dQ 原子累加。 作为附带好处，dQ 分解将全局原子归约的次数减半。原子操作是非确定性的且代价高昂，它们出现在每个内循环迭代中。因此，在 2CTA 反向传播中，每个 CTA 只写入一半的 dQ 块，并且执行的全局原子归约次数是 1CTA 对应版本的一半。

确定性模式：可重现的 `dQ`

不牺牲吞吐量

非确定性的来源是 dQ 的全局原子累加。FA4 提供了一种确定性模式，通过信号量风格的锁和内存屏障来序列化全局归约，以强制执行固定的累加顺序。然而，确定性并不意味着“一切停止”。FA4 通过 CTA 交错减少锁竞争，并使用最短处理时间优先（SPT）排序来处理因果掩码，以减少停顿。在实践中，确定性反向传播在我们的基准测试中能达到非确定性吞吐量的约 85-90%。

调度

因果掩码和可变序列长度导致注意力负载不均衡，因为不同的工作块具有不同的主循环长度，因此 FA4 改进了网格线性化，并应用最长处理时间优先（LPT）调度来减少尾部延迟。实际上，这些想法并非 Blackwell 或任何特定 GPU 架构所特有，我们在 FA3 中也使用了它们。

对于因果掩码，标准的 (mblocks, heads, batches) 网格顺序会从最短到最长处理块，这是次优的，因此 FA4 将 batch-heads 交错成 L2 大小的部分，并按 batch-head 部分遍历网格，以相反顺序迭代 mblocks，然后遍历每个部分内的 batch-heads。

对于可变序列长度，由于不同批次涉及不同的工作量，从 LPT 调度启发式的角度来看，给定的批次处理顺序通常是次优的。为了解决这个问题，我们可以启动一个预处理内核，按每个工作块的最大执行时间对批次进行排序，并写入一个虚拟到实际批次索引的映射，注意力内核使用该映射按排序顺序遍历批次；此外，元数据可以被缓存，因此排序不会带来性能损失。在撰写本文时，我们已经验证了这一想法，并在 FA3 中实现了它，我们预计在不久的将来将排序和其他元数据准备更普遍地整合到 F4 中。

语言与框架：CuTe-DSL

FA4 完全使用 CuTe-DSL（CUTLASS 的 Python 内核 DSL）实现。内核用 Python 编写；DSL 降级为 PTX，然后 CUDA 工具包编译为 GPU 机器码。编程模型镜像了 CuTe/CUTLASS 的抽象，并带有 PTX 逃生口，同时将编译时间比 C++ 模板减少了约 20-30 倍。

注意力基准测试

我们展示了 FlashAttention-4 在 B200（BF16）上的结果，并将其与 FlashAttention-2 以及 Triton、Gluon 和 cuDNN 的实现进行比较。对于 cuDNN，我们比较了 cuDNN 9.13 和最新版本 9.19.1.2。从版本 9.13 和 9.14 开始，我们与 cuDNN 团队合作，将 FlashAttention-4 的一些技术整合到 cuDNN 中，以便我们的工作能惠及尽可能多的从业者。在反向传播中，FlashAttention-4 在长序列长度上始终优于其他基线。在前向传播中，FlashAttention-4 比 cuDNN 9.13 快 1.1-1.3 倍，比 Triton 快 2.1-2.7 倍。