FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision #

§1 TL;DR #

针对 Hopper H100 GPU 的三项技术——producer-consumer warp-specialization、2-stage GEMM-softmax pipelining、FP8 block quantization + incoherent processing——将 FP16 attention 从 FA2 的 35% 利用率提升至 75%（740 TFLOPs/s），FP8 接近 1.2 PFLOPs/s。

§2 痛点 / 方法 / 结果 #

Q1 痛点 #

FlashAttention-2 的算法遵循同步执行模型，未利用 Hopper GPU 的硬件异步能力：

1. TMA 闲置：Hopper 的 Tensor Memory Accelerator 可异步搬运 GMEM↔SMEM，但 FA2 未分离数据搬运和计算。
2. WGMMA 串行化：Hopper 的 warpgroup-wide MMA (WGMMA) 可从 shared memory 异步发射，但 FA2 在 GEMM 和 softmax 之间存在强数据依赖，导致 Tensor Core 和 CUDA Core 无法重叠。
3. FP8 未利用：Hopper FP8 Tensor Core 可提供 2× matmul 吞吐，但 FP8 attention 面临布局不一致（FP32 accumulator vs FP8 operand A）和精度损失（outlier features）问题。

结果：FA2 在 H100 上仅 35% 利用率 vs GEMM 的 80–90%。

Q2 方法 #

技术 1——Producer-Consumer Warp-Specialization（§3.1）：

将 CTA 内的 warpgroup 分为 producer 和 consumer 两个角色：

• Producer：释放寄存器（setmaxnreg），专职通过 TMA 异步加载 Q, K, V 到 SMEM 的循环缓冲区。
• Consumer：获取额外寄存器，执行 WGMMA（Q·K^T 和 P·V）和 softmax。
• 通过 pipeline 对象 + barrier 同步实现 $s$-stage circular buffer。

Pingpong Scheduling：两个 consumer warpgroup 交替执行 GEMM——warpgroup 1 做 GEMM 时 warpgroup 2 做 softmax，反之。通过 bar.sync 强制 GEMM 调度顺序。

效果：从 570 → 620–640 TFLOPs/s（FP16, hdim 128, seq 8192）。

技术 2——2-Stage GEMM-Softmax Pipelining（§3.2）：

打破单次迭代内 GEMM1($\mathbf{Q}\mathbf{K}_j^\top$) → softmax → GEMM2($\tilde{\mathbf{P}}\mathbf{V}_j$) 的串行依赖：

• 迭代 $j$ 的 GEMM2 与迭代 $j+1$ 的 GEMM1 同时发射（commit but don't wait）
• 迭代 $j+1$ 的 softmax 在 GEMM1 完成后、GEMM2 开始前执行——与上一次 GEMM2 重叠

需要额外寄存器存储 $\mathbf{S}_{\text{next}}$（$B_r \times B_c \times 4$ bytes），可能与更大 block size 竞争。

SASS 分析（Appendix B.2）确认编译器确实将第一个 WGMMA 与 softmax 交错。

技术 3——FP8 Block Quantization + Incoherent Processing（§3.3）：

• 布局问题：FP8 WGMMA 仅支持 k-major 操作数，但 V 在内存中是 head-dim-contiguous（mn-major）。解决：in-kernel transpose via LDSM/STSM 指令（在 producer warpgroup 中，与 WGMMA 重叠）。FP32 accumulator 到 FP8 operand A 的布局不匹配通过 byte permute 指令解决。
• Block quantization：对 Q/K/V 的每个 $B_r \times d$ 或 $B_c \times d$ block 独立量化（独立 scaling factor），可与 RoPE 融合无额外开销。
• Incoherent processing：用随机正交矩阵 $M$（Hadamard × random $\pm 1$ diagonal）预乘 Q 和 K——因 $MM^\top = I$，注意力输出不变，但 outlier 被"扩散"到所有分量，减少量化误差。可用 Fast Walsh-Hadamard Transform $O(d \log d)$ 计算，融合进 RoPE。

核心技术壁垒：2-stage GEMM-softmax pipelining 的关键在于精确理解 WGMMA 的异步语义——commit 和 wait 的分离使得下一次 GEMM 的发射可以与当前 softmax 重叠，但需要额外寄存器 buffer 来打破数据依赖。3-stage 尝试失败（编译器拒绝同时 overlap 两个 WGMMA），说明硬件调度器行为的可预测性是实现此技术的关键约束。

Q3 结果 #

• FP16 forward：1.5–2.0× over FA2，最高 740 TFLOPs/s（75% of H100 peak 989 TFLOPs/s）。
• FP16 backward：1.5–1.75× over FA2。
• FP8 forward：接近 1.2 PFLOPs/s（hdim 256, non-causal）。
• FP8 精度：2.6× lower RMSE vs baseline FP8 per-tensor quantization（9.1e-3 vs 2.4e-2）。
• FP16 精度：与 FA2 相同（1.9e-4 RMSE），均比标准 attention 好 1.7×。

§3 架构 / 方法图 #

sequenceDiagram participant TMA as Producer WG
(TMA loads) participant WG1 as Consumer WG 1
(WGMMA + softmax) participant WG2 as Consumer WG 2
(WGMMA + softmax) Note over TMA: setmaxnreg (release regs) TMA->>TMA: TMA load Q_i → SMEM loop j = 0 to T_c-1 TMA->>TMA: TMA load K_j, V_j → SMEM[j mod s] Note over TMA,WG2: Barrier sync (producer-consumer) WG1->>WG1: WGMMA: S = Q·K_j^T WG2->>WG2: softmax(prev block) Note over WG1,WG2: Pingpong: roles swap WG2->>WG2: WGMMA: S = Q·K_j^T WG1->>WG1: softmax(prev block) WG1->>WG1: WGMMA: O += P̃·V_j WG2->>WG2: WGMMA: O += P̃·V_j end WG1-->>TMA: O_i → HBM

Target operation #

• 输入：$\mathbf{Q}, \mathbf{K}, \mathbf{V} \in \mathbb{R}^{N \times d}$，head dim $d \in \{64, 128, 256\}$
• 输出：$\mathbf{O} = \text{softmax}(\alpha \mathbf{Q}\mathbf{K}^\top)\mathbf{V}$，$\alpha = 1/\sqrt{d}$
• 精度：FP16/BF16 或 FP8 (e4m3) 输入，FP32 accumulator
• Mask：causal（可选）

Hardware model (H100 SXM5) #

级别	容量	带宽
GMEM (HBM3)	80 GiB	3.35 TB/s
L2	50 MiB	12 TB/s
SMEM (per SM)	228 KiB	31 TB/s (全 GPU)
RMEM (per SM)	256 KiB	—

参数	值
SMs	132
FP16 matmul peak	989 TFLOPs/s
FP8 matmul peak	1979 TFLOPs/s
Special function peak	3.9 TFLOPs/s
Matmul / special function 比	253×

关键硬件原语：TMA（async GMEM↔SMEM）、WGMMA（async Tensor Core from SMEM）、setmaxnreg（动态寄存器重分配）、LDSM/STSM（SMEM transpose）。

§4 作者证明 #

符号表 #

符号	含义
$B_r, B_c$	Q 行 block / K,V 列 block 大小
$T_c$	K/V block 数 $\lceil N/B_c \rceil$
$s$	Circular buffer stages
$\alpha$	Scaling factor $1/\sqrt{d}$
$L$	logsumexp $= m + \log(\ell)$
$M$	随机正交矩阵（incoherent processing）

方程物理意义 #

Attention 前向：与 FA2 数学等价——在线 softmax 逐 block 累积 $(m, \ell, \tilde{\mathbf{O}})$，最终 $\mathbf{O} = \text{diag}(\ell)^{-1} \tilde{\mathbf{O}}$。

Incoherent processing 不变性：$(\mathbf{Q}M)(\mathbf{K}M)^\top = \mathbf{Q}MM^\top\mathbf{K}^\top = \mathbf{Q}\mathbf{K}^\top$（$M$ 正交）。

FP8 outlier 分布模型：

$$X \sim \mathcal{N}(0, 1) + \mathcal{N}(0, 100) \cdot \text{Bernoulli}(0.001)$$

0.1% 的 entries 有 $\sigma = 10$ 的 outlier——模拟 LLM 中观察到的 massive activations。

无形式化性能模型——所有效率论证通过 benchmark 和 ablation 实证。

6 项检查 #

1. 数学等价性：FP16 FA3 输出与 FA2 数学相同（仅执行顺序和 overlap 不同）；FP8 引入量化误差但在受控范围内。✔
2. Warp-specialization 正确性：producer-consumer 通过 pipeline barrier 同步，circular buffer 保证 K/V 不被覆盖。✔
3. 2-stage pipelining 依赖分析：$\mathbf{S}_{\text{next}}$ 的 WGMMA 可以在 $\mathbf{S}_{\text{cur}}$ 的 softmax 之前 commit（不 wait），因为二者写不同寄存器；$\tilde{\mathbf{P}} \mathbf{V}$ 的 wait 延迟到 rescaling 前。✔
4. FP8 布局一致性：byte permute 解决 FP32 acc → FP8 operand A 布局差异；LDSM/STSM 解决 V transpose。✔
5. Incoherent processing 无偏性：Hadamard 矩阵 × random diagonal 产生正交矩阵，$MM^\top = I$ 保证注意力输出不变。✔
6. 3-stage 失败分析：编译器仅 overlap 第一个 WGMMA 和 softmax，第二个 WGMMA 不被 overlap——加上更高寄存器压力，3-stage 劣于 2-stage（Appendix B.3）。✔

Design space #

优化轴	选择	拒绝的替代方案	约束
数据搬运	TMA warp-specialization	FA2 同步 load	Hopper TMA 异步能力
GEMM-softmax overlap	2-stage pipelining	3-stage / 无 overlap	3-stage 编译器不配合 + 寄存器爆炸
Warpgroup 调度	Pingpong	自由调度	bar.sync 强制 GEMM 顺序避免 softmax 串行
FP8 V 布局	In-kernel transpose (LDSM/STSM)	预处理 transpose kernel	内存 bound 推理场景预处理浪费带宽
FP8 精度	Block quant + incoherent	Per-tensor scaling	Per-tensor 对 outlier 损失太大

Optimization techniques inventory #

技术	目标瓶颈	硬件原语	实测贡献
Warp-specialization	数据搬运延迟	TMA, setmaxnreg	570 → 582 TFLOPs/s (+2.1%)
2-stage GEMM-softmax pipelining	Tensor Core 闲置	WGMMA async commit/wait	582 → 661 TFLOPs/s (+13.6%)
Pingpong scheduling	同上	bar.sync	570 → 620–640 TFLOPs/s
FP8 in-kernel transpose	V 布局不一致	LDSM/STSM	解锁 FP8 路径
Block quantization	FP8 量化粒度	—	9.1e-3 → 9.3e-3 RMSE（微小）
Incoherent processing	Outlier 导致 FP8 误差	Hadamard FFT	2.4e-2 → 9.1e-3 RMSE（主要）

§5 实验与数据 #

FP16 Forward Speed（H100 SXM5） #

Head dim	Causal	FA3 (TFLOPs/s)	FA2 (TFLOPs/s)	cuDNN (TFLOPs/s)	FA3/FA2
64	No	497	324	413	1.53×
128	No	648	370	595	1.75×
128	Yes	616	335	539	1.84×
256	No	756	326	581	2.32×
256	Yes	642	298	509	2.15×

FA3 在中长序列（≥ 1K）全面超越 cuDNN（闭源 H100 优化库）。

FP8 Forward Speed #

Head dim	Causal	FA3 FP8	cuDNN FP8	FA3/cuDNN
64	No	613	438	1.40×
128	Yes	881	922	0.96×
256	No	1171	1139	1.03×
256	Yes	1024	1099	0.93×

FP8 causal masking 场景下 cuDNN 领先——FA3 FP8 缺少 persistent kernel + load balancing。

Ablation（2-stage pipelining） #

配置	Time (ms)	TFLOPs/s
FA3 (both enabled)	3.538	661
No pipelining, with warp-spec	4.021	582
Pipelining, no warp-spec	4.105	570

Pipelining 贡献 +13.6%，warp-specialization 贡献 +2.1%，叠加 +16%。

Numerical Accuracy #

方法	RMSE
Standard FP16	3.2e-4
FA2 FP16	1.9e-4
FA3 FP16	1.9e-4
Baseline FP8 (per-tensor)	2.4e-2
FA3 FP8 (block quant + incoherent)	9.1e-3
FA3 FP8 (no block quant)	9.3e-3
FA3 FP8 (no incoherent)	2.4e-2

关键发现：incoherent processing 是 FP8 精度提升的主要来源（去掉后 RMSE 回到 baseline），block quantization 贡献微小。

论文承认的弱项 #

1. FP8 causal masking hdim 128/256：cuDNN 领先（缺 persistent kernel）。
2. 未针对 LLM 推理优化（decode 场景）。
3. 3-stage pipelining 不如 2-stage（编译器行为不可控）。
4. FP8 attention 在大规模训练中的效果未验证。

§6 论证链 #

步骤	论点	证据
1	FA2 在 H100 上仅 35% 利用率 vs GEMM 80–90%，根因是同步执行模型	§1 para 2；§3.1 分析 H100 异步硬件能力
2	Hopper 特有异步原语（TMA, WGMMA, setmaxnreg）允许 producer-consumer 分离	§2.2 Table 1（硬件层级）；§3.1 Algorithm 1
3	Softmax 占 attention ~50% cycles（特殊函数 253× 慢于 matmul），必须与 GEMM overlap	§3.1 para 2 计算（989 TFLOPs/s matmul vs 3.9 TFLOPs/s exp）
4	2-stage pipelining 打破 GEMM-softmax 串行依赖	§3.2 Algorithm 2；Appendix B.2 SASS 分析确认 overlap
5	FP8 精度通过 incoherent processing 解决 outlier 问题	§4.3 Table 3（2.6× RMSE 改善，去掉 incoherent 后完全丧失）
6	三项技术叠加达 75% 利用率（FP16）和 ~1.2 PFLOPs/s（FP8）	§4.1 Fig. 5/6/7 全部 benchmark

§7 实现 cross-reference #

代码仓库：flash-attention

• 实现基础：CUTLASS [57] 的 WGMMA 和 TMA 抽象
• Warp-specialization：setmaxnreg 动态寄存器分配，producer 释放 registers → consumer 获得更多
• 2-stage pipeline：通过 WGMMA commit/wait 分离实现；$\mathbf{S}_{\text{next}}$ 在额外寄存器中缓存
• FP8 transpose：LDSM (ldmatrix) + STSM (stmatrix) 指令的 transpose 模式
• FP8 byte permute：将 FP32 accumulator 布局转换为 FP8 operand A 布局
• Backward pass：三角色 warp-specialization（producer + consumer + dQ-writer），dQ-writer 通过 semaphore 异步累加

关键实现细节 #

1. Pingpong 不完美：实测中 bar.sync 不能保证完美交替——SASS 显示有部分 stall，但整体仍比无 pingpong 好 10%+。
2. 3-stage 失败根因：NVCC 编译器将第二个 WGMMA 的所有 HGMMA 指令打包在一起（不与 softmax 交错），原因未知。这说明编译器对 WGMMA 的调度策略是当前实际能力的硬约束。

可移植性 #

• 强绑定 Hopper 架构（TMA, WGMMA, setmaxnreg, LDSM/STSM）
• 论文声称算法"一般性适用于任何具有 robust async execution 和 low-precision 能力的 GPU"，但实际代码深度依赖 H100 ISA
• 对 AMD CDNA4 / Intel Ponte Vecchio 的移植需要等价原语

部署上下文 #

• Serving stack：可替代 FA2 成为 vLLM / SGLang 的 H100 attention kernel
• Regime：prefill（compute-bound）最大受益；论文承认未优化 decode
• Fusion scope：softmax + matmul；RoPE 和 block quantization 可融合进前置 kernel
• Persistent kernel：FP16 有 persistent kernel + load balancing，FP8 暂无（解释了 FP8 causal 落后 cuDNN 的原因）

Software → Hardware 启示 #

• 更快的特殊函数单元：H100 上 exp 吞吐（3.9 TFLOPs/s）与 matmul（989 TFLOPs/s）差 253×。若 ISA 提供 4× 更快的 exp/log，softmax 占比从 ~50% 降至 ~12%，attention 可接近 GEMM 效率。
• 编译器透明的 WGMMA 调度：3-stage pipelining 失败因编译器黑盒行为。若 ISA 提供显式 WGMMA 调度控制（类似 CUDA graph），可恢复 3-stage 的理论优势。
• FP8 k-major 约束放宽：若 FP8 WGMMA 支持 mn-major 操作数，in-kernel V transpose 可完全消除。
• 更大寄存器文件：2-stage pipelining 的寄存器压力与 block size 竞争；更大 RMEM 可同时支持大 tile 和深 pipeline。