返回文章 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 隆重登场：[论文] [代码]

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 做一番“feeds and speeds”分析，结论恰恰相反：主要的性能瓶颈不在于张量核心做 MMA 的速度，而在于 (a) 前向计算中 softmax 指数运算所用的 SFU 单元，以及 (b) 反向计算中的共享内存流量。

在这篇博客中，我们介绍 FlashAttention-4，这是一种算法与内核协同设计，让矩阵乘法与这些其他资源瓶颈之间的重叠最大化。在 B200 上使用 BF16 时，它最高可达 1605 TFLOPs/s（利用率 71%），比 cuDNN 9.13 版本最多快 1.3×，比 Triton 快 2.7×。

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

全新流水线以实现最大重叠：全新的前向和反向软件流水线，利用 Blackwell 完全异步的 MMA 和更大的分块尺寸，让张量核心、softmax 指数运算和内存操作相互重叠。
前向（FWD）传播：用 FMA 单元上的多项式逼近实现指数函数的软件模拟，以缓解指数瓶颈，并配合条件式在线 softmax 重缩放。‍
反向（BWD）传播：将中间结果存放在张量内存中以缓解共享内存流量，结合 Blackwell 新的 2-CTA MMA 模式进一步降低共享内存流量，并将原子归约减半，同时额外支持确定性执行模式，以实现可复现的训练。‍
调度：全新的分块调度器，用于缓解因因果掩码和可变序列长度造成的负载不均衡。

Blackwell 上的新硬件特性

张量内存（TMEM）：在 B200 上，148 个 SM 各自拥有 256 KB 的 TMEM，这是一块接入张量核心的片上暂存存储，用于 warp 同步的中间结果存储。
完全异步的第 5 代张量核心：tcgen05.mma 是异步的，并在 TMEM 中累加。对于 BF16 和 FP16，最大的单 CTA UMMA 分块为 128×256×16，约为 Hopper 最大 WGMMA 原子的 2 倍。UMMA 由单个线程发起，减轻了寄存器压力，使更大的分块和更深的流水线变得可行，而不会遇到 Hopper warpgroup MMA 的寄存器溢出痛点。这也让 warp 特化更加可行：一部分 warp 搬运分块，另一部分发起 MMA，从而让矩阵乘累加与 softmax 和内存流量相重叠。tcgen05.mma 还可以从 TMEM 获取操作数 A。
2-CTA MMA：Blackwell 可以在同一集群内的一对 CTA 上执行一次 UMMA，横跨两个对等 CTA 的 TMEM。领导者 CTA 中的一个线程发起 MMA，但在其执行期间两个 CTA 都必须保持活跃。通过在这对 CTA 之间拆分 M 和 N，MMA 分块维度可以扩展到 256×256×16，既减少了冗余流量，又降低了每个 CTA 的占用。同一个内核中，跨 TMEM 和张量核心操作时，CTA 组的大小（1 或 2）必须保持不变。

Feeds and Speeds

对于 M=N=D=128

B200 上的 Feeds（每 SM）：

张量核心（BF16）：$\frac{8192 \text{ ops}}{cycle}$
指数单元：$\frac{16 \text{ ops}}{cycle}$
共享内存流量：$\frac{128 \text{ bytes}}{cycle}$

Speeds（每个分块的时钟周期数）：

前向（2 次 MMA + MN 次 exp）
1. 张量核心：$1024$
2. Exp：$1024$
3. SMEM：$768$
反向（5 次 MMA + MN 次 exp）：1-CTA
1. 张量核心：$2560$
2. Exp：$1024$
3. SMEM：$3328$

结论：前向受限于计算和指数运算，反向受限于共享内存带宽。因此我们在前向传播中让 softmax 与 MMA 重叠，在反向传播中减少共享内存流量。

前向传播：带条件重缩放的全新 softmax 流水线

前向传播有两次矩阵乘，QK^T 和 PV。在 Blackwell 上，张量核心快了很多，但指数单元（MUFU.EX2）没有变快。因此 softmax 不再只是“夹在两次矩阵乘之间的那点事”，而是必须精心流水线化的瓶颈。

FWD 传播概述：

每个 CTA 采用 $2x$ Q 和 $2x$ O 分块的乒乓调度：最大化 MMA 与 Softmax 之间的重叠
2x softmax warpgroup：按分块做 softmax，并通过同步避免在计算指数时相互重叠
- $2^x$ 的软件模拟：将 exp 计算分散到硬件的 MUFU 和在 FMA 上做软件模拟的两条路径
- 分阶段将 P 存入 TMEM：缓解寄存器压力
校正 warpgroup：指定专门的“校正”warpgroup 来执行重缩放，使其脱离关键路径
- 在线 softmax（条件式）重缩放：降低重缩放频率，以尽量减少非矩阵乘运算

流水线：乒乓式 Q 分块加专用的校正阶段

FlashAttention-4 为每个 CTA 计算两个查询分块，即 $Q^H$ 和 $Q^L$，每个分块覆盖 128 个查询词元，并以乒乓调度交替处理它们。

Blackwell 改变了 softmax 的映射方式。S = QK^T 的累加器分块大小为 128×128，存放在张量内存中；然而，读入寄存器后，按照硬件规定的分块划分方式，是每行一个线程。我们使用两个 128 线程的 warpgroup，每个 Q 分块对应一个，每个 softmax warpgroup 执行如下操作序列：

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

关键细节在于 exp 是瓶颈所在的环节。我们显式地同步两个 softmax warpgroup，使它们不会同时计算 exp，从而减少 MUFU 的争用。

为了让重缩放脱离关键路径，内核把它分配给专用的 warpgroup。校正 warpgroup 负责计算：

仅当最大值跳变较大时才重缩放：
$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 粒度上做这一决策，以避免分支发散。

更快的指数运算：将 $2^x$ 分散到 MUFU.EX2 和 FMA（软件模拟）

Softmax 需要大量指数运算，而 MUFU 的吞吐量远低于张量核心。FlashAttention-4 让 exp2 的软件模拟与硬件 MUFU.EX2 路径并行运行，利用本来未被充分使用的 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，并运行 QK 和 PV 的梯度 MMA 以得到 dQ、dK、dP 和 dV，此外还有针对 P 和 dS 的逐元素运算。在 Blackwell 上，限制反向的不是 FLOPs，而是共享内存带宽

流水线：让 MMA 与 softmax 重叠

Hopper 时代的 FlashAttention-3 将 MMA 累加器保存在寄存器中，因此寄存器压力往往迫使调度更偏串行。在 Blackwell 上，累加器位于 TMEM 中，这使得在 CUDA 核心处理 P 和 dS 的逐元素运算时，可以让多个 MMA 同时在途。由于在我们的 roofline 分析中，指数运算的吞吐量与两次 MMA 相当，把它隐藏起来是值得的。

关键的重叠很简单：在为分块 j 计算 softmax 的同时，我们已经发起了分块 j−1 的 dK 和 dQ MMA

为减少共享内存流量，反向传播在重算 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 作为一个分块协同工作：每个 CTA 暂存操作数 B 的一半，并只保留自己的那份累加器切片。这使操作数 B 的共享内存流量大致减半。

归约轴冲突。为了削减 B 的流量，我们在五个反向 GEMM 中都采用 M=256、N=K=128 的 MMA 分块，但 dQ MMA 的性质引入了不匹配。在 FlashAttention 的反向传播中，每个 CTA 负责一个固定的 KV 分块（外层循环在 N 个 CTA 上并行），并在内层循环中遍历 M 个分块。dQ 的更新要在外层循环中对 KV 序列做归约。2-CTA MMA 拆分的是输出分块，而不是归约，而 dQ 的归约维度是 N，它已经被这对 CTA 拆分了。每个 CTA 对自己负责的那些行，仍然需要完整的归约。

解决方案：DSMEM 交换。我们利用集群内的分布式共享内存，在两个 CTA 之间交换一半的 dS 来解决这一问题。这会重新排布 dS，使其沿非归约轴划分：每个 CTA 负责 M/2 行，同时持有完整的 2N 归约。于是每个 CTA 的 dQ MMA 变为 (M/2, 2N)(2N, d)，在张量内存中累加出一个 (M/2, d) 的分块。在 2-CTA 模式下，S、dP、dV 和 dK 的 MMA 保持 M=256，而 dQ 使用 M=128，归约维度翻倍为 2N=256。随后我们重新排列流水线以隐藏 DSMEM 延迟：先计算当前分块的 dP，再计算上一个分块的 dQ。由于 dQ 分块可以与 P 一起放入 TMEM，它可以复用 S 所用的 TMEM 区域，因此 dP 和 dQ 不再像 1-CTA 模式下那样共用同一区域。按这种顺序，当前分块逐元素的 dS 与上一轮迭代的 dQ MMA 相重叠。

dQ 原子加法。一个附带的好处是，dQ 的这种分解使全局原子归约的次数减半。原子操作是非确定性的，而且开销很大，并且在每次内层循环迭代中都会发生。因此，在 2CTA 反向传播中，每个 CTA 只写入 dQ 分块的一半，执行的全局原子归约次数也只有 1-CTA 版本的一半。

确定性模式：可复现的 `dQ`，且不拖垮吞吐量

非确定性的来源是 dQ 的全局原子累加。FA4 提供了一种确定性模式，用信号量式的锁和内存栅栏将全局归约串行化，以强制固定的累加顺序。但确定性并不一定意味着“一切都停下来”。FA4 用 CTA swizzling 来减少锁争用，并对因果掩码采用最短处理时间优先（SPT）的排序来减少停顿。在我们的基准测试中，实际上确定性反向可达到非确定性吞吐量的约 85-90%。

调度

因果掩码和可变序列长度会使注意力的负载不均衡，因为不同的工作分块（worktile）的主循环长度不同，所以 FA4 改进了网格线性化，并采用最长处理时间优先（LPT）调度来缩短尾部。事实上，这些思路并非特定于 Blackwell 或任何特定的 GPU 架构，我们在 FA3 中也使用了它们。

对于因果掩码，标准的 (mblocks, heads, batches) 网格顺序会按从最短到最长的次序处理分块，这并不理想，因此 FA4 把 batch-head 以 L2 大小的区段做 swizzle，并按 batch-head 区段遍历网格：先倒序遍历 mblocks，再遍历每个区段内的 batch-head。

对于可变序列长度，由于不同批次的工作量不同，从 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 倍。