FlashAttention-4：面向 Blackwell 的算法-流水线协同设计

Blackwell 微架构的关键变化#

Blackwell B200 相对于 Hopper H100 的主要变化可以概括为：张量核心吞吐翻倍有余，其余几乎原地踏步。

这个表格揭示了核心矛盾：张量核心的吞吐翻了 2 倍以上，而执行 softmax 指数运算的 SFU（Special Function Unit）/MUFU（Multi-Function Unit）吞吐完全没有增加。共享内存带宽也几乎不变。

这意味着什么？在 H100 上，注意力计算的主要瓶颈是矩阵乘——张量核心一直在满负荷工作，softmax 和内存搬运可以隐藏在矩阵乘的延迟之后。而在 B200 上，矩阵乘不再是瓶颈，softmax 的指数计算和共享内存带宽分别成为了前向和反向传播的关键路径。

前向传播的瓶颈转移#

考虑一个典型的前向注意力计算（以 head dim = 128，序列长度 = 8K 为例）：

每个注意力 tile 需要完成的操作为：

• 两个矩阵乘：$S = QK^T$ 和 $O = PV$
• 一次 softmax：对 $S$ 的每一行计算指数、求和、归一化

在 H100 上，矩阵乘大约需要 512 个时钟周期，softmax 的指数运算大约需要 512 个周期（因为 SFU 一次只能处理 16 个元素，而一个 tile 有 128 行需要逐行计算指数）。但由于矩阵乘和 softmax 不能完全并行（softmax 必须等 $S$ 算完才能开始），且 H100 的张量核心速度刚好能让 softmax 几乎不成为瓶颈。

但在 B200 上，矩阵乘只需约 256 个周期（张量核心快了一倍），而 softmax 的指数运算仍然需要约 512 个周期（SFU 没变快）。结果是指数计算需要的时间是矩阵乘的 2 倍，成为前向传播的绝对瓶颈。

反向传播的瓶颈转移#

反向传播的计算图更加复杂——它涉及 5 次矩阵乘（计算 $dQ$、$dK$、$dV$ 需要 $dS = dO \cdot V^T$、$dP = dO \cdot V$ 的转置等操作）和一次对 $S$ 的指数运算（用于重算 softmax 梯度）。在 H100 上，这 5 次矩阵乘大约需要 2560 个周期，远大于共享内存的读写时间。

但在 B200 上，5 次矩阵乘只需约 1280 个周期（张量核心快了一倍），而中间结果的共享内存读写仍然需要约 3328 个周期（共享内存带宽没变）。结果是共享内存带宽成为了反向传播的主导瓶颈——与普遍认为”注意力计算受限于矩阵乘”的直觉完全相反。

FA4 的应对策略#

面对这个局面，FA4 的设计哲学可以概括为：既然矩阵乘不再是瓶颈，就把之前隐藏在其他计算后面的非矩阵乘操作重新设计，使其不再拖累整体吞吐。

具体来说：

• 前向传播：不再靠硬件 SFU 算指数，改用 FMA（Fused Multiply-Add）单元的软件近似，将有效指数吞吐提升数倍。
• 反向传播：利用 Blackwell 新引入的 TMEM（Tensor Memory）存放中间结果，避免经过共享内存，绕开共享内存带宽瓶颈。
• 整体流水线：利用 Blackwell 每 SM 支持 16 个 warp 的能力，构建更深的多级异步流水线，最大化各硬件单元的并发利用率。

Warp 特化架构#

FA4 的前向 kernel 使用了 16 个 warp（每 warp 32 线程，共 512 线程），划分为 6 种功能角色：

这种精细的角色划分在之前的 GPU 架构上是不可行的。H100 每 SM 只有 8-12 个 warp，无法支撑 6 种角色的流水线深度。Blackwell 将每 SM 的 warp 容量提升到 16 个，为这种细粒度特化提供了硬件基础。

每个 warp 角色的设计都围绕一个原则：让不同硬件单元同时处于忙碌状态。当加载 warp 正通过 TMA 从 HBM 搬运下一批 $K$、$V$ 数据时，softmax warp 正利用张量核心计算 $QK^T$，同时 MMA warp 正将上一个 tile 的 $P$ 和 $V$ 相乘。三种操作——内存搬运、矩阵乘、指数计算——在不同的硬件单元上并发执行。

流水线同步机制#

多级流水线的正确性依赖于精确的同步。FA4 使用 Blackwell 的命名屏障（named barriers）来协调各 warp 之间的数据依赖：

1
加载阶段 ──→ Softmax 阶段 ──→ MMA 阶段 ──→ 校正阶段 ──→ 收尾阶段
2
(TMA load)   (QK^T + softmax)  (P×V)       (rescale O)   (TMA store)
3

4
mbar_load_KV → mbar_P_full   → mbar_O_full → mbar_O_epi → TMA store
5
               mbar_S_full_P_full_O_rescaled

每一步数据传输都由显式的屏障保护。当加载 warp 完成数据搬运后，它通过 mbar_load_KV 通知 softmax warp 可以开始计算。Softmax warp 完成概率矩阵 $P$ 的计算后，通过 mbar_P_full 告知 MMA warp 可以开始 $P \times V$。MMA warp 完成后通过 mbar_O_full 触发校正 warp。校正 warp 完成重缩放后通过 mbar_O_epi 让收尾 warp 执行 TMA store。

关键细节：$P$ 矩阵被分为两个阶段写入 TMEM（阶段 0 占用偏移 64，阶段 1 占用偏移 192），这使得 MMA warp 可以在 Softmax 阶段 0 完成后立即开始消费 $P$ 矩阵的前半部分，无需等待 Softmax 阶段 1 完成。这是减少串行等待时间的关键设计。

混合指数计算：绕过 SFU 瓶颈#

FA4 最具突破性的创新之一是混合指数计算。传统的注意力实现依赖 GPU 的硬件指数单元（MUFU.EX2）来计算 softmax 中的 $\exp$ 操作。但在 Blackwell 上，MUFU.EX2 的吞吐（16 ops/cycle/SM）仅为张量核心吞吐（8192 ops/cycle/SM）的约 1/500，成为前向传播的绝对瓶颈。

FA4 的解决方案是：用 FMA 单元上的三次多项式近似来分担指数计算压力。具体来说，将一个 tile 中 128 个元素的指数计算按可调比例分配给硬件 MUFU 和软件 FMA 两个路径，两条路径并行执行。

软件路径的核心算法如下：

第一步：Cody-Waite 范围归约

将指数函数重新表述为整数部分和分数部分的乘积：

其中 $\lfloor x \rfloor$ 是 $x$ 的整数部分（向负无穷取整），$x - \lfloor x \rfloor \in [0, 1)$ 是分数部分。整数部分的 $2^{\lfloor x \rfloor}$ 可以通过 IEEE 754 浮点数的位操作高效实现——直接将 $\lfloor x \rfloor$ 加上 bias 后移入指数位即可。这完全不需要任何函数计算。

第二步：三次多项式逼近

对于分数部分 $f = x - \lfloor x \rfloor \in [0, 1)$，使用三次多项式逼近 $2^f$：

其中系数 $p_0 = 1.0$、$p_1 \approx 0.6951$、$p_2 \approx 0.2276$、$p_3 \approx 0.0771$，由数值软件包 Sollya 在 $[0, 1)$ 区间上最小化相对逼近误差得到。

多项式使用 Horner 方法 求值以减少乘法次数：

只需要 3 次 FMA 指令即可完成。FMA 指令的吞吐远高于 MUFU.EX2，且 FMA 单元在前向传播的 softmax 阶段本来就有空闲（因为矩阵乘已经由张量核心异步执行了），所以这个软件路径几乎不增加额外延迟。

第三步：组合结果

将整数部分和分数多项式的结果组合：将 $\lfloor x \rfloor$ 移入浮点数的指数位，乘以多项式逼近得到的 $2^f$ 尾数值，得到最终的 $2^x$。

精度论证：三次多项式的 FP32 相对误差约为 $10^{-5}$ 量级，高于硬件 MUFU 的精度。但注意力计算通常使用 BF16 精度，其有效尾数只有 7 位（约 $10^{-2}$ 量级的分辨率）。一旦结果被舍入到 BF16，软件逼近和硬件指数的舍入行为几乎无法区分。换句话说，在 BF16 精度下，三次多项式逼近等价于硬件指数。

算法还需要处理边界情况。$x < -127$ 时，$2^x$ 在 FP32 中下溢为零，因此将所有输入 clamp 到 $x \ge -127$。$x$ 过大时直接返回 $+\infty$（这在 softmax 中通常不会发生，因为 $QK^T$ 的值域有限）。

两个 softmax warp 组被显式同步，以避免同时使用 MUFU.EX2 造成争用——当一组使用硬件 MUFU 时，另一组使用 FMA 路径，交替进行。

条件 Softmax 重缩放#

Online softmax 算法在遍历 $KV$ 序列块时维护两个运行统计量：

• $m_j$：当前已看到的最大 logit 值
• $\ell_j$：当前已累加的指数和

当遇到更大的 $m_j$ 时，需要将之前累加的部分输出 $O_{j-1}$ 按比例 $\exp(m_{j-1} - m_j)$ 重缩放。在标准的 online softmax 实现中（包括 FA3），每次 $m_j$ 变化都要重缩放。

FA4 的关键洞察：不需要每次 $m_j$ 变化都重缩放。如果 $m_j - m_{j-1}$ 很小，跳过重缩放带来的数值误差可以忽略。FA4 引入了一个阈值 $\tau$：

当 $m_j - m_{j-1} \le \tau$ 时，直接跳过重缩放，用旧的 $m_{j-1}$ 代替 $m_j$ 进行局部的指数计算。在所有 $KV$ 块遍历完成后，最终做一次全局归一化：$O_{\text{final}} = O / \ell_{\text{final}}$，保证数学上的正确性。

这个简单但高效的技巧消除了约 90% 的重缩放操作。在 FA3 中，重缩放位于关键路径上——校正 warp 必须等 MMA warp 完成上一轮、且 softmax warp 完成新 tile 的 softmax 之后才能执行，造成流水线停顿。FA4 的重缩放操作大幅减少后，校正 warp 很少成为瓶颈，流水线得以保持满载。

TMEM：直接连接张量核心的本地存储#

Blackwell 架构引入了一种新的片上存储——TMEM（Tensor Memory），每 SM 256 KB。TMEM 的特殊之处在于它直接连接到张量核心，是张量核心的专用工作内存，不像共享内存需要通过交叉开关（crossbar）路由。

FA4 在反向传播中利用 TMEM 存放中间结果，从而避免了这些结果经过共享内存。具体来说，$S = QK^T$ 的计算结果直接写入 TMEM，随后被 softmax 梯度重算和 $dS$ 的计算直接消费——整个过程不经过共享内存。这大幅减少了共享内存的读写压力，使得共享内存带宽不再成为瓶颈。

在前向传播中，TMEM 也被用于存放 $S$ 和 $P$ 矩阵，使 softmax warp 和 MMA warp 之间的数据交换在 TMEM 内部完成：

要注意 TMEM 只在数据中心级 Blackwell GPU（B200、B300）上可用。消费级 Blackwell GPU（RTX 5090、RTX PRO 6000）没有 TMEM，因此 FA4 在这些 GPU 上不可用。

2-CTA MMA 模式#

2-CTA MMA 是 Blackwell 张量核心的一项新功能：两个 CTA（Cooperative Thread Array，即线程块）协作完成一次矩阵乘法。两个 CTA 共享各自的 TMEM，操作数在两者之间分配。

FA4 在反向传播的 $dQ$ 计算步骤中利用了 2-CTA MMA 模式。传统的做法是每个 CTA 独立加载完整的 $K$（用于计算 $dQ = dS \cdot K$ 时所需的操作数），导致每个 CTA 都要从共享内存读取一遍 $K$，造成重复的共享内存访问。在 2-CTA MMA 模式下，两个配对的 CTA 各自只加载 $K$ 的一半，然后共享操作数，将共享内存访问量减半。

此外，传统方法中每个 CTA 需要对自己的部分 $dQ$ 结果执行原子加法（atomic add）将其累加到全局内存中的 $dQ$ 梯度张量。2-CTA MMA 模式将两个 CTA 的输出在 TMEM 中合并后再执行一次原子加法，将原子操作数减半。

这个优化虽然看起来简单，但对反向传播的整体吞吐有显著影响——共享内存带宽是反向传播的瓶颈，任何减少共享内存流量的措施都能直接转化为加速。论文在 B200 上测得的反向 TFLOPS 印证了这一点：借助 TMEM 与 2-CTA MMA 这两项针对性优化，FA4 的反向 kernel 在长序列下同样达到了远超 Triton、领先 cuDNN 的吞吐水平。

确定性模式#

FA4 的反向传播还支持确定性执行模式。在默认的非确定性模式下，由于浮点加法的非结合性，原子加法的执行顺序会影响最终 $dQ$ 的值，导致每次运行产生微小差异。确定性模式通过固定归约顺序来保证每次运行得到完全一致的结果，代价是约 10-15% 的吞吐损失。

在诸如强化学习（RL）等对可复现性要求严格的场景中，确定性模式非常关键——训练过程中任何非确定性都可能导致策略梯度估计的偏差累积。

论文对确定性反向传播做了细致的消融：在因果注意力下，SPT（短路径优先）、LPT 及其反向 mblock 顺序等调度策略的相对收益被逐一拆解；在非因果注意力下，则对比了带 batch/head swizzle 与朴素调度。这些消融说明确定性模式带来的 10-15% 吞吐损失并非固定代价，调度策略的选择可以在保证可复现性的同时尽量缩小这一差距。

LPT 调度#

在实际推理中，不同请求的 prompt 长度可能差异很大——从几个 token 的短对话到数十万 token 的长文档。此外，因果掩码（causal mask）使得不同 query 位置对应的 $KV$ 长度不同（位置 $i$ 只需关注前 $i$ 个 token）。这两种情况都会导致各 query 的计算量不均衡，而 GPU 的并行执行模型（SIMT）要求同一 warp 内的线程同时完成才能进入下一阶段，负载不均意味着快的线程在等慢的线程。

FA4 采用 LPT（Longest Processing Time First）调度策略处理变长序列和因果掩码：将计算量最大的 query 块优先分配给处理单元，减少尾部等待时间。在多 query 并行处理时，LPT 调度遵循”先处理最重的任务”原则，使得整体完成时间逼近理论下界。

论文报告的实验数据显示，LPT 调度为多头注意力（MHA）带来 4-8% 的额外 FLOPs 增益，对于 MQA/GQA（Multi-Query / Grouped-Query Attention）则可达到 14%。这是因为 MQA/GQA 中 $KV$ 头数远少于 $Q$ 头数，$QK^T$ 在 $K$ 维度上的负载不均更加显著。

FlexAttention 集成#

PyTorch 的 FlexAttention API 允许用户在 Python 层面通过 score_mod 函数定义自定义的注意力变体（ALiBi、滑动窗口、文档掩码、soft-capping 等），而不需要编写 CUDA 代码。FA4 实现了作为 FlexAttention 后端的集成，在 Hopper 和 Blackwell GPU 上可用。

当用户定义一个 score_mod 函数后，FlexAttention 将其 JIT 编译为 FA4 kernel 的一个变体。相比 Triton 后端，FA4 后端的加速在 1.2× 到 3.2× 之间，具体取决于注意力模式的复杂度。这意味着研究者可以快速实验新的注意力变体，同时享受到手写 CUDA kernel 级别的性能——在此之前，这两者是不可兼得的。

Kernel 级别性能#

在 B200 GPU 上的 BF16 精度下，FA4 的前向 kernel 达到了最高 1613 TFLOPs/s，约为 B200 理论峰值（约 2250 TFLOPS BF16）的 71%。对于注意力计算而言，这是一个里程碑式的数字——以往的注意力 kernel（包括 FA3 在 H100 上）很难突破 50-60% 的硬件利用率，因为非矩阵乘操作（softmax、内存搬运）长期占据大量时间。

对比基线：

• 1.3× 于 cuDNN 9.13（NVIDIA 的官方注意力库，同样针对 Blackwell 优化）
• 2.7× 于 Triton 实现的注意力 kernel
• FA4 在序列长度 4K 以上时优势最为显著，因为此时 softmax 和共享内存带宽的瓶颈效应最强

FA4 的性能优势并不局限于标准的 head dim = 128 配置。在 DeepSeek V3 架构常用的 head dim (192, 128) 配置下，FA4 同样保持了对 cuDNN 的稳定领先——这说明它对非标准头维度的兼容并非”只对默认配置调优”，而是算法层面的通用设计。

端到端推理性能#

在实际 LLM 推理中（BF16，batch size = 32，8K 上下文）：

在更长序列（32K+）上，FA4 的优势进一步扩大，因为注意力计算在总延迟中的占比随序列长度平方增长。对于需要处理超长上下文的 RAG、代码库问答和 Agent 工作负载，FA4 在 B200 上的 2× 级别加速直接转化为用户体验的显著改善。

成本效率#

虽然 B200 的按需价格高于 H100，但由于单卡吞吐约为 H100 的两倍（尤其是在长序列下），部署相同吞吐量所需的 GPU 数量减少，从而降低了跨 GPU 通信开销。在 spot 实例定价下，B200 + FA4 的每百万 token 成本与 H100 + FA3 几乎持平，但延迟显著更低。