> 本文记录我在 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 中的典型路径是： ```text 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 索引，并定义一个简单协议： ```text 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 到完整推理：我设的三道门 我没有让微基准结果直接决定是否合入。每一个候选都必须过三道门： 1. **算子正确性门**：与 PyTorch reference 比较，覆盖有效 shape、边界 shape、dtype 和非法输入； 2. **模型回归门**：检查完整 69 通道输出、通道顺序与任务指标； 3. **端到端门**：在同一 BW1000 节点、同一批样本下记录时间和显存，而不是只比较一个 kernel。 在同节点回归中，固定宽度 LayerNorm、residual-norm、token gather、norm-gather、bias-mask 和恢复路径组合后，稳定段平均推理时间由约 `0.15295～0.15320 s` 降至 `0.134999 s`。 这个结果并不意味着“自定义算子一定比框架快”。它只说明在这组稳定 shape、这版 DTK 和这条模型图上，减少通用调度和中间读写是有效的。 --- ## 8. 这次 Profile 给我的方法论 这次最值得复用的流程是： ```text 端到端计时异常 -> 采集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*](https://doi.org/10.1038/s41586-023-06185-3)。本文数字来自本地 BW1000 回归，不代表任何隐藏测试集成绩。