本文记录我在 BW1000gfx936上优化盘古气象模型推理的分析过程。重点不是给出一份“万能算子清单”而是展示如何从一张真实 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-Weather1. 先看原始证据而不是先写 kernel我对完整推理路径做了 PyTorch device profile按self_device_time_total排序。下表摘自原始operator_table_device.txt算子或区段Device 时间调用次数单次平均aten::addmm35.244 ms112314.680 μsaten::_softmax17.774 ms28634.789 μsaten::native_layer_norm14.521 ms34427.087 μsaten::bmm12.470 ms56222.677 μs输出传输阶段10.153 ms110.153 ms输入准备阶段7.902 ms17.902 msaten::gelu_6.634 ms16414.600 μsaten::index_select5.740 ms32179.360 μs高空恢复阶段4.970 ms14.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 时直接回退或报错避免“看似通用、实际没测过”的路径进入推理。微基准结果如下ShapePyTorch LayerNormgfx936 kernel[1,131040,192]0.7791 ms0.1139 ms[1,32760,384]0.2910 ms0.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 ms0.1687 ms[1,32760,384]0.3511 ms0.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。ShapePyTorch gathergfx936 gather[1,131040,192] → 135360 tokens0.2920 ms0.1219 ms[1,32760,384] → 34560 tokens0.1536 ms0.0473 ms进一步观察发现norm1的输出只被窗口 gather 消费。于是我把 LayerNorm 和 gather 合并省掉完整 normalized-token tensor 的一次写入和再次读取Shape分离 norm gather融合 norm-gather[1,131040,192]0.2358 ms0.1656 ms[1,32760,384]0.0904 ms0.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.152950.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 回归不代表任何隐藏测试集成绩。