FlashAttention 完全拆解（三）：Warp 特化异步流水线与 FP8 低精度

Hopper 的三种新硬件能力#

第一，TMA（Tensor Memory Accelerator，张量内存加速器）。Hopper 为”HBM 与共享内存（SMEM）之间的数据搬运”单独造了一个硬件单元 Hopper 架构深入解析。之前 GPU 上搬数据要由线程亲自执行拷贝指令（cp.async 也要线程逐条发指令、按行计算地址），TMA 则不同：只需要一个线程向 TMA 描述”搬哪块、搬到哪”（含多维索引计算、越界处理、甚至内嵌 swizzle 布局变换），TMA 硬件就自动把整块数据从全局内存搬到共享内存。发起搬运的线程不阻塞，可以立刻去干别的。数据搬完由硬件置位一个屏障（barrier），其他 warp 可以”等这个屏障”。

第二，WGMMA（warpgroup MMA，warpgroup 级矩阵乘加指令）。Ampere 的 mma.sync 指令粒度是一个 warp（32 线程）；Hopper 的 WGMMA 把粒度扩大到 warpgroup（4 个 warp、128 线程），一次指令完成更大的矩阵乘，吞吐约为前代的两倍。更重要的是它是异步的：发出 WGMMA 后，张量核心（Tensor Core）在执行，执行它的 warp 不必等结果出来就能继续发射其他指令，之后再用”提交-等待”（commit-and-wait）原语同步结果。

第三，FP8 张量核心。WGMMA 支持 FP8（e4m3）操作数，吞吐是 FP16 的两倍。但 FP8 的使用有一堆布局约束（后面第三节细讲）。

此外还有一个容易被忽略但被 FA-3 用上的能力：setmaxnreg 指令允许 kernel 在运行时把寄存器从一组 warp 动态重新分配给另一组 warp——做数据搬运的 warp 用不了几个寄存器，做矩阵乘的 warp 越多越好。

论文把这些硬件能力总结成了下面这张内存层次表（论文表 1）：

注意 SMEM 的带宽（31 TB/s）比 HBM（3.35 TB/s）高一个数量级——分块算法的全部意义就是把数据尽量留在 SMEM 里反复用。

同步执行模型的代价：非矩阵乘操作吃掉一半周期#

FA-2 的 kernel 里，所有 warp 一视同仁：都要搬数据、都要算矩阵乘、都要算 softmax。在”同步”的执行模型下，这些操作串行发生：加载完 K、V 才能做 $QK^T$，$QK^T$ 完成才能做 softmax，softmax 完成才能做 $PV$。问题在于，矩阵乘有专用硬件（张量核心），softmax 里的指数运算只有吞吐低得多的 SFU（特殊函数单元）。

论文 §3.1 给出了两个关键数字（H100 SXM5，FP16）：

• FP16 矩阵乘峰值：989 TFLOPS；
• 指数（exponential）运算吞吐：3.9 TFLOPS——两者相差 256 倍。

再算前向注意力的操作构成：对 head dimension 128 的注意力，每行分数矩阵要做一次 $QK^T$（$\propto N d$ 次乘加）和一次行 softmax（$\propto N$ 次指数）。矩阵乘的 FLOPs 是指数操作 FLOPs 的 512 倍，但指数的吞吐只有矩阵乘的 $1/256$。于是：

指数计算花费的时间是矩阵乘的两倍——softmax 的指数运算可以占据 kernel 一半的执行周期。而且注意这里的 3.9 TFLOPS 是整个 GPU 的 SFU 吞吐，指数与矩阵乘虽在不同单元上可以部分重叠，但 FA-2 的同步模型让它们几乎没有机会重叠。FP8 下局面更糟：矩阵乘吞吐翻倍到 1978 TFLOPS，指数吞吐纹丝不动，比例失衡翻倍。

这就是 FA-2 在 H100 上只有 35% 利用率的直接原因：它把”吞吐差距 256 倍的两种操作”串行安排，让最慢的操作决定了整个 kernel 的速度。FA-3 的全部设计，一句话概括就是：让矩阵乘、指数、数据搬运这三种发生在不同硬件单元上的操作，尽可能同时进行。

角色划分：只搬数据的 warp 和只算数的 warp#

warp specialization 的思想并不新鲜——论文引用的是 2011 年的 CudaDMA 工作，NVIDIA 的 CUTLASS 库也一直在用；新鲜的是 Hopper 硬件让专职化变得几乎免费：TMA 搬运只需要一个线程发一条指令，其余 31 个线程可以完全不用干活；而 WGMMA 的累加器放在寄存器里，消费者 warp 的寄存器越多越好。setmaxnreg 让这种分工在运行时动态调整。

FA-3 前向 kernel 把一个线程块内的 warp 分成两个角色（论文 Algorithm 1）：

• 生产者 warpgroup（producer）：专职发起 TMA 加载。执行 setmaxnreg 把自己的寄存器上限压得很低（TMA 搬运只需要一个线程发指令，其余线程几乎不用寄存器），把省下的寄存器让给消费者。它只做一件事：把 $Q_i$（一次）和 $\{K_j, V_j\}_{j=0}^{T_c-1}$（循环）加载进环形共享内存缓冲（circular SMEM buffer，$s$ 级流水），每次加载完成后通过命名屏障”提交”（commit）通知消费者；
• 消费者 warpgroup（consumer）：执行 setmaxnreg 把寄存器上限提到很高（WGMMA 的累加器全在寄存器里，具体数量按块大小调优），全部资源用来算：初始化 $O_i$、$\ell_i$、$m_i$，等 $Q_i$ 和 $K_0$ 就位后，循环执行两个 WGMMA（$S^{(j)}_i = Q_i K_j^T$ 和 $O_i = \tilde{P}^{(j)}_i V_j$）与 softmax 的状态更新。

两个角色之间用共享内存的 mbarrier（命名屏障）同步：生产者写满缓冲的第 $j$ 阶段后 commit，消费者等这个屏障、消费完数据后释放该阶段，生产者才能覆盖它。环形缓冲的深度 $s$ 决定了生产者可以领先消费者多少次迭代——生产者永远比消费者快几步，把 HBM 的延迟藏在后面。

这里有一个值得细想的点：为什么让生产者用 TMA、而不是让消费者自己加载？因为 TMA 是异步的——消费者发一条 TMA 指令后不需要等数据到位，可以先去算上一块数据，等真正需要新数据时再等屏障。FA-2 的同步模型里，“加载”是阻塞的：warp 发了拷贝指令就得等，等的时候张量核心空转。FA-3 把”等”从计算路径上拆走了：等数据的是生产者（它本来也没别的活干），消费者永远面对”数据已就位”的缓冲。

完整的前向主循环#

把角色划分落实到算法上，就是论文的 Algorithm 1（前向，无消费者内部重叠版本）。它的骨架与 FA-2 相同——外层遍历 $Q$ 行块（序列维并行）、内层遍历 $K$、$V$ 列块——但每个动作都换成了异步版本：

1
输入：HBM 中的 Q_i ∈ R^{B_r×d}，K, V ∈ R^{N×d}，键块大小 B_c，T_c = ⌈N/B_c⌉
2
 1: 初始化流水线对象，管理 s 级环形共享内存缓冲的屏障同步
3
 2: if 在生产者 warpgroup 中 then
4
 3:     释放预定数量的寄存器（setmaxnreg）
5
 4:     发起从 HBM 到共享内存的 Q_i 加载（TMA）
6
 5:     完成后 commit，通知消费者 Q_i 已加载
7
 6:     for j = 0 to T_c-1 do
8
 7:         等待缓冲第 (j % s) 阶段被消费完
9
 8:         发起 K_j、V_j 的 TMA 加载到第 (j % s) 阶段
10
 9:         完成后 commit，通知消费者
11
10:     end for
12
11: else（消费者 warpgroup）
13
12:     按消费者 warp 数量重新分配寄存器（setmaxnreg）
14
13:     在片上初始化 O_i = 0、ℓ_i = 0、m_i = -∞
15
14:     等待 Q_i 加载完成
16
15:     for j = 0 to T_c-1 do
17
16:         等待 K_j 加载完成
18
17:         计算 S_i^(j) = Q_i K_j^T（SS-GEMM，提交并等待）
19
18:         保存 m_i_old，计算 m_i = max(m_i_old, rowmax(S_i^(j)))
20
19:         计算 P̃_i^(j) = exp(S_i^(j) - m_i)，ℓ_i = e^{m_i_old - m_i}·ℓ_i + rowsum(P̃_i^(j))
21
20:         等待 V_j 加载完成
22
21:         计算 O_i = diag(e^{m_i_old - m_i})^{-1}·O_i + P̃_i^(j)·V_j（RS-GEMM，提交并等待）
23
22:         释放缓冲第 (j % s) 阶段给生产者
24
23:     end for
25
24:     计算 O_i = diag(ℓ_i)^{-1}·O_i 和 L_i = m_i + log(ℓ_i)
26
25:     把 O_i、L_i 写回 HBM
27
26: end if

对比 FA-2 的算法，逐行的变化值得注意：

• 第 3、12 行：寄存器在运行时按角色重新分配——生产者用 setmaxnreg 把上限压低，消费者提高。WGMMA 的 FP32 累加器（$B_r \times B_c$ 的 $S$ 块）要放在寄存器里，这是消费者寄存器需求的大头；
• 第 4、8 行：TMA 加载只由生产者发起，单线程发指令，硬件完成搬运——消费者完全不需要感知数据搬运的发生；
• 第 5、9 行：commit 通知通过命名屏障实现，生产者与消费者的握手全部基于硬件屏障而非线程同步（__syncthreads）；
• 第 17、21 行：两个 GEMM 都标注了”SS”（第一个操作数从共享内存取）或”RS”（第一个操作数从寄存器取）前缀——$S = QK^T$ 的两个操作数都在共享内存，而 $O = \tilde{P} V$ 的概率块 $\tilde{P}$ 在寄存器里（它是上一步算出来的，没写回共享内存）。这是 FA-2 没有的寄存器-共享内存混合数据流；
• 第 7、22 行：生产者等”缓冲被消费”，消费者等”数据被加载”——两个方向的等待互相错开，这就是环形缓冲的意义。前 $s$ 次迭代生产者不会等待（缓冲从空开始填，没有东西可等）。

这套框架中，每个 warp 的指令流都极其简单：生产者是一连串 TMA 和屏障，消费者是一连串 WGMMA、逐元素操作和屏障。编译器面对这种简单的指令流更容易生成好的调度——这也是 warp specialization 的隐性收益。

反向传播：三种角色#

反向传播（论文附录 B.1，Algorithm 3）在前向的生产者-消费者之外增加了一个角色——dQ 写入 warp。原因是 FA-2 的反向里我们讲过：$\mathbf{dQ}_i$ 的梯度需要跨线程块累加（每个线程块给同一位置的 $\mathbf{dQ}$ 贡献一份），FA-2 用原子加处理。在 FA-3 里，多个线程块同时原子写全局内存的 $\mathbf{dQ}$ 会造成内存竞争——如果让消费者 warpgroup 自己去做这些原子加，它在等原子操作完成时就没法发射下一个矩阵乘。

FA-3 的解法是把原子加职责单独拆给一个 warp：消费者算完局部的 $\mathbf{dQ}^{(local)}_i$ 后写进共享内存、通知 dQ 写入 warp；dQ 写入 warp 用信号量（semaphore）串行地把这些局部结果原子累加到全局 $\mathbf{dQ}$。“谁负责跟全局内存打交道”从计算路径上被彻底拆走，消费者 warpgroup 只管在片上算。

反向前半段还有一处结构变化：FA-2 需要预计算 $D = \text{rowsum}(\mathbf{dO} \circ \mathbf{O})$ 并写回 HBM。FA-3 把它放进独立的预处理 kernel（Algorithm 3 第 1 行），避免在主 kernel 里多一趟同步。主 kernel 的消费者循环与 FA-2 的 Algorithm 2 结构一致（重算 $S$、$P$、五个矩阵乘），但全部换成异步 WGMMA。

伪依赖的本质：跨迭代流水#

关键观察：$j$ 次迭代的 softmax 只依赖第 $j$ 次迭代的 $S$ 块，$j$ 次迭代的第二个 GEMM 只依赖第 $j$ 次迭代的 $\tilde{P}$ 块。迭代之间并没有数据依赖——所以可以把”GEMM1(j+1) → softmax(j) → GEMM2(j)“这三个来自不同迭代的操作叠在一起：前者的输出不是后者的输入，流水线就成立了。FA-3 用了两级手段实现这个重叠。

第一级：乒乓调度（pingpong scheduling）——两个消费者 warpgroup 交替#

FA-3 把消费者拆成两个 warpgroup，让它们一个做 GEMM、一个做 softmax，交替轮换（乒乓球式）：

图里同一种颜色代表同一次迭代 $j$：可以看到 warpgroup A 在处理迭代 $j$ 的 GEMM 时，warpgroup B 正在处理迭代 $j-1$ 的 softmax——两者在不同的硬件单元上同时执行（GEMM 用张量核心，softmax 用 CUDA 核心/SFU）。论文承认”实际中乒乓调度并不如图示这么干净”（迭代长度不同、单元负载不均衡），但实测收益明确：FP16 前向、head dim 128、序列长度 8192 时，单靠这一项就把吞吐从 570 TFLOPS 提到 620–640 TFLOPS。

第二级：warpgroup 内 2 级流水（2-stage pipelining）——提交但不等待#

乒乓调度只在”两个 warpgroup”层面重叠，每个 warpgroup 内部仍是串行。FA-3 进一步在同一 warpgroup 内部做流水（论文 Algorithm 2）。核心技巧是 WGMMA 的提交-不等待（commit but do not wait）语义：

1
 7: for j = 1 to T_c-1 do
2
 8:     等待 K_j 加载完成
3
 9:     计算 S_next = Q_i K_j^T（WGMMA，提交但不等待）
4
10:     等待 V_{j-1} 加载完成
5
11:     计算 O_i = O_i + P̃_cur·V_{j-1}（WGMMA，提交但不等待）
6
12:     等待 S_next 的 WGMMA 完成
7
13:     基于 S_next 计算 m_i、P̃_next、ℓ_i   ← softmax（迭代 j）
8
14:     等待 O 的 WGMMA 完成，然后重缩放 O_i
9
15:     释放缓冲第 (j % s) 和第 ((j-1) % s) 阶段
10
16:     把 S_next 复制到 S_cur
11
17: end for

对照图理解（下面这张是论文的 Figure 2，2-stage 流水线的示意）：

逐行解读这套流水线的时序：

• 第 9 行：发出迭代 $j$ 的第一个 WGMMA（$Q_i K_j^T$），不等待结果——张量核心开始算 $S_{j}$，warp 立即继续；
• 第 11 行：发出迭代 $j-1$ 的第二个 WGMMA（$\tilde{P}_{j-1} V_{j-1}$），也不等待——此时张量核心上有两个矩阵乘排队执行；
• 第 12 行：等到 $S_j$ 算完——注意此时第 11 行的 WGMMA 应该也快完成了（它更早发出）；
• 第 13 行：对 $S_j$ 做 softmax（最大值、指数、和）——这是纯 CUDA core 操作，与此同时张量核心正在执行第 11 行那个矩阵乘。这就是”把 softmax 藏在 GEMM 下面”的字面含义；
• 第 14 行：等第 11 行的矩阵乘完成，重缩放 $O$。

妙处在于：第 13 行的 softmax 与第 11 行的矩阵乘在时间上重叠——一个用 CUDA core，一个用张量核心，互不抢占。论文正文的原话是：“主循环中，迭代 $j$ 的第二个 WGMMA（第 11 行）与迭代 $j+1$ 的 softmax 操作（第 13 行）重叠执行。”

这个设计有两个工程上的坑，论文都点到了：

坑一：编译器重排。 NVCC 经常重新排列指令以优化调度，可能打乱精心编排的 WGMMA/非 WGMMA 交错序列，导致重叠消失或行为异常。论文在附录 B.2 里对 SASS（汇编）代码做了分析，确认编译器生成的代码确实如预期重叠——但这份确认是针对当时 CUDA 12.3 编译器的，换个编译器版本可能就得重新验证。这也是手写 CUDA 异步 kernel 的普遍痛点。

坑二：寄存器压力。 2 级流水要求同时持有 $S_{cur}$ 和 $S_{next}$ 两个分数块（迭代 $j-1$ 的还没消费完，迭代 $j$ 的已经算出来），每个线程块多占 $B_r \times B_c \times 4$ 字节的寄存器（FP32）。这会与”用更大块大小提升效率”竞争——块越大寄存器越紧张，论文的建议是”根据 profiling 结果做权衡”。

3 级流水：为什么反而更差#

论文还尝试了 3 级流水（附录 B.3，Algorithm 4），理想目标是同时重叠三个操作：迭代 $j+2$ 的第一个 WGMMA、迭代 $j+1$ 的 softmax、迭代 $j$ 的第二个 WGMMA：

结果：3 级流水表现反而更差。论文给了两个原因：

1. 编译器不配合：SASS 分析显示，编译器只把”第一个 WGMMA”与 softmax 重叠了，“第二个 WGMMA”没有被重叠——编译器认为的调度顺序与设计意图不一致，且论文作者也说不清编译器为什么这样重排；
2. 寄存器压力更大：3 级流水需要在寄存器里额外保存 $\tilde{P}_i$ 和缩放因子 $\mathit{scale\_o}$（$B_r \times B_c \times$ 数据类型大小 + $B_r \times 4$ 字节），为容纳这些寄存器只能把块调小，块变小又损失了 GEMM 效率。

这个负面结果很有价值：它说明了异步流水线的深度不是越大越好，受制于寄存器文件容量和编译器的实际行为。2 级流水是 FA-3 在 Hopper 上找到的甜点。

消融实验：两项技术各贡献多少#

论文表 2 给了三组消融数据（FP16 前向，head dim 128）：

三个数字告诉我们：warp 特化（生产者-消费者异步）单独贡献约 570→582（约 2%），而 GEMM-softmax 流水线在与 warp 特化合体后贡献了 582→661（约 13%）。两者合力约 16%。流水线（把 softmax 藏起来）是更大的单项贡献——这与前文”指数运算占一半周期”的分析一致：解决了最大的瓶颈，收益最大。注意消融里的”570”是基线，与乒乓调度一节提到的 570 一致。

布局约束一：FP8 WGMMA 只接受 k-major 操作数#

WGMMA 要求矩阵操作数以特定方式在内存中连续。对一个 $M \times K$ 矩阵 $A$ 来说，若 $A$ 在 $M$ 维连续则称为 mn-major，在 $K$ 维连续则称为 k-major。FP16 WGMMA 对两种布局都接受；FP8 WGMMA 只接受 k-major。

注意力里出问题的是第二个 GEMM：$O = \tilde{P} V$，其中 $V$ 是序列维连续的（按惯例 Q、K、V 都以 head 维为连续维，即每行是一个头的 $d$ 个元素）。第二个 GEMM 需要 $V$ 的 tile 在序列长度维连续——TMA 加载本身不能改变连续维（它只是搬运）。论文列出了三个候选方案：

1. 1a：把转置融合到前一步的 epilogue（如 RoPE 的输出直接写成转置后的布局）——但这对标准库集成太困难；
2. 1b：独立的转置预处理 kernel——在内存受限的推理场景太浪费（多一趟 HBM 往返）；
3. 内核内转置：数据加载进 SMEM 后用指令转置——FA-3 选了这个。

内核内转置用的是 LDSM（ldmatrix）/STSM（stmatrix） 指令：一个 warp 的线程集体从共享内存加载到寄存器、或从寄存器存回共享内存，以 128 字节为粒度，拷贝的同时可以改变布局（转置）。关键工程细节：这个转置放在生产者 warpgroup 里执行（生产者反正只发 TMA，有大量闲置的指令发射能力），且第一次迭代之后，对下一个 $V$ tile 的转置可以隐藏在涉及当前 $V$ 和下一个 $K$ 的两个 WGMMA 的阴影里——转置不需要额外时间。

布局约束二：FP8 WGMMA 的累加器与操作数布局不一致#

前向的两个 GEMM 是背靠背的：第一个 GEMM（$S = QK^T$）的输出是 FP32 累加器，直接作为第二个 GEMM（$\tilde{P}V$）的 FP8 操作数。问题在于：FP8 WGMMA 的 FP32 累加器在寄存器里的布局，与它作为操作数 A 时要求的布局不一样（论文图 3、图 4 画了具体的线程-寄存器映射）。如果直接喂，第二个 WGMMA 会算出错误的结果。

FA-3 的解法是用 byte-permute（字节置换）指令在寄存器里就地变换布局：FP8 WGMMA 的 FP32 累加器在寄存器里的列分组顺序是 $\{d_0, d_1, d_2, d_3, d_4, d_5, d_6, d_7\}$，而作为 FP8 操作数 A 时要求 $\{d_0, d_1, d_4, d_5, d_2, d_3, d_6, d_7\}$ 的顺序（论文图 3、图 4 画了两种布局的线程-寄存器映射）——用字节置换指令把 $d_2, d_3$ 与 $d_4, d_5$ 交换即可，该置换每 8 字节重复一次。逻辑效果是置换 $P$ 块的列顺序；为了让第二个 GEMM 仍然算对，内核内转置在写 $V$ tile 时写出匹配的行置换——两个置换互相抵消，数学结果不变。

这两个布局问题的解决让 FP8 WGMMA 可以连续地执行两个背靠背 GEMM——论文明确感谢了 cuDNN 团队贡献的内核内转置思路。

精度问题一：块量化（block quantization）#

FP8（e4m3）只有 3 位尾数，动态范围极窄。简单地对整个 Q、K、V 用一个全局缩放因子（per-tensor scaling）量化，精度损失严重。FA-3 采用块量化：对每个 $Q_i$、$K_j$、$V_j$ 块分别维护一个缩放因子：

其中 $\tilde{Q}_i$、$\tilde{K}_j$ 是量化后的 FP8 块，$s_{Q_i}$、$s_{K_j}$ 是各自的缩放标量。块级缩放与 FA 的块迭代结构天然契合：每块一个缩放因子不需要任何额外全局操作，而且分数矩阵 $S$ 的块缩放 $s_{Q_i} s_{K_j}$ 可以直接融合进在线 softmax 的重缩放步骤——几乎零成本。

精度问题二：非相干处理（incoherent processing）#

块量化解决”动态范围”问题，但 LLM 的激活值里还有离群值（outlier）维度：极少数维度上的值比其他维度大一个量级（这在很多 LLM 里被反复观察到）。量化时这些离群值吞掉大量缩放预算，导致其他维度精度崩溃。

FA-3 采用 QuIP/QuIP# 等量化工作的”非相干处理”思路：量化前把 Q 和 K 乘以同一个随机正交矩阵 $M$。因为 $M$ 正交（$MM^T = I$）：

注意力输出数学上完全不变。而 $QM$ 的每个元素都是 $Q$ 各元素的随机加权和——离群值的能量被”摊平”到所有维度，每个维度的动态范围都变得均匀，量化误差随之下降。实践中 $M$ 取”随机 $\pm 1$ 对角矩阵 × Hadamard 矩阵”的乘积，与一个 $d$ 维向量相乘只需 $O(d \log d)$（Hadamard 的快速变换），且可以与 RoPE 旋转融合，零额外计算成本。论文引用了 Chee et al. 的 QuIP 和 Tseng et al. 的 QuIP# 作为这一技术的来源。

数值误差验证#

论文用带离群值的分布验证精度（§4.3）：$Q$、$K$、$V$ 的元素按 $N(0,1) + N(0,100) \cdot \text{Bernoulli}(0.001)$ 生成——正常分布上叠加 0.1% 的”大离群项”，模拟真实 LLM 激活。以 FP64 计算为参考，比较 RMSE：

读这张表有三个发现。第一，FP16 下 FA-2 和 FA-3 完全一致（1.9e-4），且都比标准实现好 1.7 倍——因为它们的 softmax 中间结果（指数和、累加器）保持 FP32，标准实现把中间结果存成 FP16。第二，FA-3 的 FP8 相比 per-tensor 基线误差低 2.6 倍（9.1e-3 vs 2.4e-2）。第三（表格里最有信息量的一点）：去掉非相干处理后误差直接回到 2.4e-2——非相干处理是 FP8 精度的主要贡献者，块量化单独只有微弱影响（9.3e-3 vs 9.1e-3）。论文正文的表述是”块量化与非相干处理共同把 FP8 的数值误差降低 2.6 倍”，但消融揭示了主次。

另外一个容易忽略的要点：FP8 只用于前向。论文明确写”我们改造前向算法以利用 FP8 张量核心”——反向的五个矩阵乘仍用 FP16。附录 C.2 的 FP8 基准也只有前向。原因是反向的梯度累加对精度更敏感，且反向的布局约束更复杂（$dQ$ 的跨块累加在 FP8 下几乎不可行）。所以”FA-3 是 FP8 注意力”这个说法要精确为”FA-3 的前向支持 FP8”。

FP16 前向：最高 740 TFLOPS#

下面这张是论文的 Figure 5（FP16/BF16 前向，6 个子图覆盖三种头维度 × 有无因果掩码）：

挑最常用的 head dim 128 报具体数字（无掩码，图 5c；单位 TFLOPS）：

四个观察：

1. 短序列（512）FA-3 略低于 cuDNN（467 vs 497），1K 起反超——论文正文的原话是”中长序列（1K 及以上）FA-3 甚至超过 NVIDIA 闭源、专门为 H100 优化的 cuDNN”。短序列下流水线还没填满，cuDNN 的深度调优占优；
2. 相对 FA-2，FA-3 在 16K 时 648/370 = 1.75 倍；论文摘要给的整体区间是 1.5–2.0 倍（head dim 256 时差距更大）；
3. 序列越长相对优势越大——长序列下 $K$、$V$ 块的遍历次数多，流水线和 TMA 的收益被放大；
4. 加因果掩码后（图 5d）趋势不变：16K 时 FA-3 616 vs FA-2 335，约 1.8 倍；掩码跳过的块让两个实现都变快，但 FA-3 的异步流水线在变长的跳过块循环里同样保持优势。

head dim 256（图 5e，无掩码）时 FA-3 达到 756 TFLOPS——这就是论文摘要”FP16 最高 740 TFLOPS、约 75% 利用率”的来源。为什么 head dim 越大效率越高？因为单个 WGMMA 的矩阵更大，张量核心利用更充分，流水线更不容易被小矩阵的启动开销打断。反过来 head dim 64 时最高只有 497 TFLOPS——注意力 kernel 的效率对 head dim 极其敏感，这也是后来 DeepSeek 选 (192, 128) 这类非标准 head dim 时各家都要重新调 kernel 的原因。

反向：1.5–1.75 倍#

反向（FP16/BF16，图 6）相对 FA-2 快 1.5–1.75 倍。head dim 128、无掩码：FA-3 在 16K 时 474 vs FA-2 的 291 TFLOPS。反向的收益小于前向，因为反向有 5 个矩阵乘，softmax 重算只占其中一小部分，“藏 softmax”的空间更小；反向的主要收益来自 warp 特化（dQ 写入 warp 消除了原子加阻塞）和 TMA。

FP8 前向：接近 1.2 PFLOPS#

论文图 7 是 FP8 前向（head dim 256，无掩码/有掩码）：

FP8 前向在长序列下达到接近 1.2 PFLOPS（论文摘要原话）。相对 FP16 的 740 TFLOPS 又提了约 60%——没有完全翻倍，因为 softmax 仍在 FP16/FP32 下计算、块量化和转置有开销，但方向明确：FP8 把注意力的上限推到了 PFLOPS 级别。注意图 7 里 FP8 的序列长度用的是 512、1024、2048、4224、8448、16896——≥4K 时取 132 的倍数（H100 SXM5 的 SM 数），避免最后一波线程块不满载导致的性能波动（wave quantization 效应，附录 C.2 解释了这点）。

与标准实现对比：3–16 倍#

论文摘要还报了一个区间：FA-3 相对 PyTorch 标准注意力实现快 3–16 倍（标准实现短序列 16K 时直接 OOM——32 头 × 16K² 的分数矩阵已经超出显存）。这个对比说明 IO 感知 + 异步 + 低精度三代的收益是乘法叠加的：IO 分块省掉了 $\Theta(N^2)$ 的 HBM 流量，异步把剩余时间里的利用率从 35% 提到 75%，FP8 又把每个周期的计算量翻倍。