一次 Profile 如何决定 DCU 算子优化顺序：盘古气象模型 BW1000 实录

weixin_47594597

542人浏览 · 2026-07-18 00:05:07

weixin_47594597 · 2026-07-18 00:05:07 发布

本文记录我在 BW1000（gfx936）上优化盘古气象模型推理的分析过程。重点不是给出一份“万能算子清单”，而是展示如何从一张真实 Profile 表，决定哪些工作值得做、哪些原型应当及时停止。

摘要

在 DCU 上优化模型，最容易犯的错误是看到 addmm 排第一，就立刻重写 GEMM。我的实际 Profile 显示，112 次 aten::addmm 共占 35.244 ms，确实是单项第一；但同一张表还显示 34 次 native_layer_norm 占 14.521 ms、32 次 index_select 占 5.740 ms，以及大量短逐元素 kernel。最终最稳定的端到端收益并不来自替换 GEMM，而来自围绕固定 shape 的 LayerNorm、窗口 gather 和相邻逐元素路径做专用化。

本文给出实际 Profile、选型依据、HIP 扩展组织方式、正确性门禁和被拒绝的原型。所有数字来自同一 BW1000 节点的本地回归；它们用于说明分析方法，不代表隐藏集成绩。

关键词：海光 BW1000、gfx936、PyTorch Profiler、HIP、LayerNorm、算子融合、Pangu-Weather

1. 先看原始证据，而不是先写 kernel

我对完整推理路径做了 PyTorch device profile，按 self_device_time_total 排序。下表摘自原始 operator_table_device.txt：

算子或区段	Device 时间	调用次数	单次平均
`aten::addmm`	35.244 ms	112	314.680 μs
`aten::_softmax`	17.774 ms	28	634.789 μs
`aten::native_layer_norm`	14.521 ms	34	427.087 μs
`aten::bmm`	12.470 ms	56	222.677 μs
输出传输阶段	10.153 ms	1	10.153 ms
输入准备阶段	7.902 ms	1	7.902 ms
`aten::gelu_`	6.634 ms	16	414.600 μs
`aten::index_select`	5.740 ms	32	179.360 μs
高空恢复阶段	4.970 ms	1	4.970 ms

另外两个数字让我警惕：Profile 中有 173 次 aten::copy_，277 次 aten::empty，而且 hipDeviceSynchronize 的 CPU 等待明显。这说明问题不只是“某个算子算得慢”，还包含临时张量、数据搬运和同步带来的端到端成本。

2. 为什么没有直接接入一套通用 GEMM 工程

addmm 占用最高并不自动等于“先优化 addmm”。我先检查了它的形状和上下游关系，结论是：这些 Linear 已经处于 PyTorch/DTK 的高效矩阵乘路径中，而模型的隐藏维度和权重格式并不适合直接套用外部通用 GEMM 项目。

更重要的是，替换 GEMM 会带来新的问题：

权重格式转换是否额外占显存；
小矩阵、多调用场景是否被 launch 开销淹没；
新实现是否改变 FP16 累加和舍入；
外部工程能否被干净地放入提交包；
新依赖是否在目标 DTK 环境稳定编译。

我把优先级转向两个更适合专用化的热点：固定宽度 LayerNorm 和 shifted-window 的 token gather。它们的共同点是 shape 固定、调用频繁、前后依赖稳定，而且可以减少中间 tensor 的读写。

这个选择不是“避开大问题”，而是先处理端到端收益更确定、正确性边界更清楚的问题。

3. 固定宽度 LayerNorm：从通用实现切到模型已知的形状

盘古学生模型运行时的主要隐藏宽度只有 192 和 384。模型加载阶段将 LayerNorm affine 参数等价折叠到后续 Linear，因此设备侧 LayerNorm 只需完成标准化。

gfx936 的 wavefront 为 64 线程。我的 kernel 设计为一个 wave 处理一行，256-thread block 并行处理四行，FP16 输入输出、FP32 进行均值和方差归约。这里没有试图支持任意宽度；输入不是 192 或 384 时直接回退或报错，避免“看似通用、实际没测过”的路径进入推理。

微基准结果如下：

Shape	PyTorch LayerNorm	gfx936 kernel
`[1,131040,192]`	0.7791 ms	0.1139 ms
`[1,32760,384]`	0.2910 ms	0.0410 ms

数值检查中，单层最大绝对差约为 0.001953125，平均绝对差约为 2.6e-8。原因不是公式变化，而是 reduction 顺序不同。这个结果必须结合完整模型输出继续验证，不能只看单层误差就认为“可用”。

4. 融合 residual add + norm：先确定 FP16 舍入点

Transformer block 中的典型路径是：

residual = FP16(residual + branch)
normalized = LayerNorm(residual)

一个常见但危险的写法是把两个步骤放进同一 kernel 后，全程用 FP32 做加法再归一化。公式看起来等价，结果却未必一致：原路径在 residual add 后已经发生一次 FP16 舍入。

我的实现明确保留这个舍入点：先产生与原路径一致的 FP16 residual destination，再在同一次读取中做 FP32 LayerNorm reduction。

Shape	原 add + norm	融合 kernel
`[1,131040,192]`	0.8990 ms	0.1687 ms
`[1,32760,384]`	0.3511 ms	0.0922 ms

这件事让我意识到，融合算子的验收不能只问“少了几个 kernel launch”，还要问“原始路径在哪一步降精度、在哪一步写回、是否允许改变这些语义”。

5. Window gather：把 padding 也纳入索引协议

shifted-window 路径原来触发了 32 次通用 index_select。这类操作本身并不复杂，但它要处理 padding、窗口重排和索引张量，频繁调用时会放大调度和读写开销。

我的实现为每种窗口形状预计算 INT32 索引，并定义一个简单协议：

index >= 0：读取源 token
index == -1：输出零

这样 padding 不需要先构造完整填零张量，再进入 gather。kernel 使用 __half2 搬运连续 FP16 token。

Shape	PyTorch gather	gfx936 gather
`[1,131040,192] → 135360 tokens`	0.2920 ms	0.1219 ms
`[1,32760,384] → 34560 tokens`	0.1536 ms	0.0473 ms

进一步观察发现，norm1 的输出只被窗口 gather 消费。于是我把 LayerNorm 和 gather 合并，省掉完整 normalized-token tensor 的一次写入和再次读取：

Shape	分离 norm + gather	融合 norm-gather
`[1,131040,192]`	0.2358 ms	0.1656 ms
`[1,32760,384]`	0.0904 ms	0.0554 ms

这里的收益主要来自减少大 tensor 的往返，而不是“融合”这个词本身。

6. 哪些原型我主动放弃了

为了避免文章只展示成功案例，下面列出我实际测试后没有保留的路线：

原型	观察到的问题	决定
bias + mask + softmax 全融合	大 shape 略快，其他 shape 变慢，数值路径更复杂	不保留
自定义 tanh GELU	浅层和深层微基准都慢于 PyTorch	不保留
SDPA 替换	当前窗口 shape 下更慢，显存更高	不保留
CUDA Graph	增加常驻显存，端到端综合收益为负	不保留
attention 多分块	没有降低最终峰值，额外增加调度	不保留
DeepGEMM 接入	当前矩阵形状、权重格式和提交体积不匹配	不接入

“没有进入最终版本”的实验同样重要。它们帮我确认：这条模型路径的瓶颈不是所有 Transformer 都共有的那几个名词，而是具体 shape、具体运行时和具体内存生命周期共同决定的。

7. 从 microbenchmark 到完整推理：我设的三道门

我没有让微基准结果直接决定是否合入。每一个候选都必须过三道门：

算子正确性门：与 PyTorch reference 比较，覆盖有效 shape、边界 shape、dtype 和非法输入；
模型回归门：检查完整 69 通道输出、通道顺序与任务指标；
端到端门：在同一 BW1000 节点、同一批样本下记录时间和显存，而不是只比较一个 kernel。

在同节点回归中，固定宽度 LayerNorm、residual-norm、token gather、norm-gather、bias-mask 和恢复路径组合后，稳定段平均推理时间由约 0.15295～0.15320 s 降至 0.134999 s。

这个结果并不意味着“自定义算子一定比框架快”。它只说明在这组稳定 shape、这版 DTK 和这条模型图上，减少通用调度和中间读写是有效的。

8. 这次 Profile 给我的方法论

这次最值得复用的流程是：

端到端计时异常
  -> 采集device profile
  -> 先看调用次数和shape，再看总时间
  -> 选稳定、可解释的局部路径
  -> 写reference与边界测试
  -> 微基准
  -> 完整模型回归
  -> 再决定保留或删除

如果少了“调用次数和 shape”这一步，很容易把最高耗时算子当作唯一目标；如果少了“完整模型回归”，又会把局部最快实现带进端到端退化。Profile 的价值不是给出答案，而是把优化范围缩小到可以实验、可以证伪的几个假设。

第三方依赖与来源说明

本文实验调用 PyTorch Profiler、主办方提供的 OneScience、ROCm/DTK 和 MIOpen。HIP 扩展源码由本人针对 gfx936 编写；Pangu-Weather 模型设计参考 Bi 等人的公开论文 Accurate medium-range global weather forecasting with 3D neural networks。本文数字来自本地 BW1000 回归，不代表任何隐藏测试集成绩。