第 5 课：IAW 模型 FA2 vs FA3 实测

理论说 FA-3 快 1.5–2×，实测里它真的总是更快吗？

本课基于一份真实模型的 kernel-trace 级性能分析（飞书文档原文）。它会把第 4 课的"Hopper 异步理论"拉回现实：在非 Hopper 平台上，FA-3 走的是另一条路。

1. 背景：IAW 模型的 attention 在算什么

IAW 是一个多模态模型（vision + LLM）。在 sequence packing 训练场景下，attention 的完整路径是：

输入: (B=1, total_len, H, D) 稠密张量（含 padding 位置）
  │
  ├─ unpad: 去掉 padding token，压成紧凑张量 (total_valid, H, D)
  │         ↑ 按 indices 取行（gather）
  │
  ├─ flash_attn_varlen_func: 在紧凑张量上做 attention（核心 kernel）
  │
  └─ pad: 把紧凑结果散回 (B, total_len, H, D)
          ↑ 按 indices 填回（scatter）

attention 分两块，开销占比悬殊：

• vit 部分（dense attention，12 次/层）：72ms / 326ms ≈ 22%。
• llm 部分（变长 attention，24 次/层）：14ms / 151ms ≈ 9.4%。

2. FA2 vs FA3：unpad/pad 其实差别不大

一个常见误解：FA2 用 C++ 的 bert_padding 做搬运，FA3 用本地 index_select，所以 FA3 更快。kernel trace 打破了这个说法：

步骤	FA2 (bert_padding)	FA3 (本地)
unpad	at::native::_scatter_gather_elementwise_kernel	at::native::indexSelectLargeIndex
pad	vectorized_elementwise + index_elementwise	相同
本质	都是 PyTorch native op 的薄封装，差异很小

也就是说，FA2 的 bert_padding 并不是定制 CUDA kernel。性能差异的根源不在搬运，而在核心 flash kernel 的架构。

3. 核心差异：手写 CUDA vs CUTLASS

维度	FA2	FA3
框架	手写 CUDA（flash::flash_fwd_kernel）	CUTLASS 模板（cutlass::FlashAttnFwdSm80）
架构适配	通用 kernel	按 GPU 架构特化
tiling	标准	CollectiveMainloop 优化 tiling + pipeline

4. 平台对照：PPU810E 与"可移植性边界"

这份实测跑在 阿里自研 PPU810E 上——一个 CUDA 兼容的非 NVIDIA 平台（torch.cuda 可用，kernel 以 cutlass:: / at::native:: 名字出现）。CUTLASS 在它上面选中了 Sm80 路径。

5. 实测数据：eager 模式 FA3 确实更快

LLM 变长 attention（packed），eager 模式：

配置	FA2 eager	FA3 eager	FA3 加速
32×64（多短序列）	0.254 ms	0.128 ms	1.98×
32×128（中等 packing）	0.274 ms	0.175 ms	1.56×
1×2048（单长序列）	0.378 ms	0.233 ms	1.63×
1×8192（单长序列）	2.065 ms	1.674 ms	1.23×

短序列 packed 场景优势最大，长序列优势收窄。FA3 eager 全面领先，来源是 CUTLASS flash kernel。

6. 反转：compile 之后 FA3 反而变慢

把 torch.compile 加进来，故事变了：

配置	FA2 compile 加速	FA3 compile 加速
32×64	1.49×（变快）	0.51×（变慢近一半）
32×128	1.36×	0.74×（变慢）
1×2048	1.46×	0.86×（变慢）
1×8192	1.01×	1.00×（持平）

定量拆解 FA3 compile 为什么亏（典型短序列）：

FA3 eager:   3×native index_select(9μs) + FillFunctor(3μs) + cutlass(100μs) + 2×native pad(6μs) = 118μs
FA3 compile: 2×triton index_select(30μs) + FillFunctor(3μs) + cutlass(100μs) + 2×triton pad(30μs) = 163μs
                                                                                  ↑ 多 45μs

• triton dispatch 比 native 慢 5×（15μs vs 3μs）：Inductor 固有开销，无法消除。
• FillFunctor 无法 fuse：它在 C++ extension 内部（graph break 之后），Dynamo 够不着。
• fusion 只省掉 1 次 kernel + 1 次冗余写，约 −15μs。

7. Vision dense：FA3 反而最慢

vit 走的是 dense（非 packed）attention，而且实际跑的是 SDPA（PyTorch 内置 flash_sdp，不是 flash_attn 包）。三方对比（每层，12 heads MHA，seq=1024）：

配置	SDPA	FA2	FA3
B=1	0.068	0.124	0.095
B=149（trace 匹配）	6.544	6.007	7.194

12 层总开销（B=149）：SDPA 78.5ms / FA2 72.1ms（最快但不可用）/ FA3 86.3ms。

8. 选型建议

场景	推荐路径	是否 compile	原因
LLM attention（packed）	FA3 eager	整体 compile	CUTLASS kernel 最快；MLP（87% 计算量）的 compile 收益覆盖 attention 的劣化
Vision attention（dense）	SDPA	可选	PyTorch 内置、免装包、比 FA3 快
评测 / 推理	FA3 eager	不 compile	无需编译，FA3 eager 已最快

典型 packed 32×256 路径延迟：FA2 eager 0.353 → FA2 comp 0.255 → FA3 eager 0.257 → FA3 comp 0.253（compile 对 FA3 无额外收益）。

9. 一句话总结

10. 快速测验

11. 延伸阅读

本课的全部数据与 kernel trace 来自 飞书文档：FA2 vs FA3 性能对比分析（IAW 模型 / PPU810E），含完整 benchmark 脚本与 Asight Systems 报告。

对照回看 第 4 课：FA-3 Hopper 异步执行与 FP8，体会"理论上的 Hopper 异步"与"实测中的 CUTLASS 通用路径"的裂缝。下节课预告：FlashAttention-4：Blackwell 重架构与 CuTe-DSL。