模型在昇腾910D上跑通后,性能问题常会集中到少数自定义算子。很多团队会立刻打开Triton或TBE,从Tiling、Block Size和数据搬运开始试。这样做控制力强,维护成本也高:模型shape、融合边界或CANN版本发生变化,原来的参数通常需要重新验证。
昇腾官方将算子开发定义为对神经网络基础单元的设计与实现,目标是让算法在昇腾AI处理器上高效运行。具体到工程现场,调优应该先找证据,再改代码。
先确认慢在哪里
据鲲鹏昇腾开发者社区专栏介绍,CANN 8的profiling模块可以记录算子运行耗时、通信时长、内存占用和下发延迟。这些信息能帮助开发者先回答几个问题:
- 慢的是计算,还是下发、通信与数据搬运?
- 相邻算子之间,是否存在中间结果反复写回HBM的情况?
- 热点是否会在业务常用的shape和batch下稳定出现?
- 单算子耗时下降后,端到端延迟是否真的会跟着变化?
以RMSNorm和Gated-SiLU为例,如果两者分开发射,且后一个算子立即消费前一个算子的输出,就值得检查融合机会。目标很明确:尽量减少中间结果写回HBM后再次读取。至于访存占比和融合收益,应由现场profile给出,不宜预设百分比。
Auto-Kernel负责搜索,边界仍要人来定
手写Triton或TBE时,开发者需要直接处理循环、Tiling和数据搬运。Auto-Kernel类工具换了一种分工:开发者描述算子的数学关系及约束,编译器转换为中间表示,再生成不同的Tiling、Block Size和流水方案进行测试。
放到RMSNorm与Gated-SiLU这个组合里,开发者要先讲清楚两段计算的数据依赖、输入输出、shape、数据类型和精度要求。编译器据此判断能否合并执行,并在可用的搜索空间里选择候选方案。
这里有一个容易忽略的前提:自动生成不保证天然更快。搜索空间是否完整、编译器是否支持相关语义、测试shape是否贴近业务,都会影响结果。搜索需要多久、性能能提高多少,也只能通过目标环境中的实测确定。
手写和自动生成,各自承担什么
| 对比项 | 手写Triton/TBE | Auto-Kernel自动生成 |
|---|---|---|
| 开发输入 | 循环、Tiling、数据搬运和调度细节 | 数学语义、shape、精度与融合约束 |
| 调优方式 | 人工修改参数并反复profile | 编译器生成候选方案并采样比较 |
| 版本变化 | 重新检查原有Tiling和内存假设 | 重新运行搜索并做性能回归 |
| 排障方式 | 底层控制直接,但维护内容较多 | 仍需结合生成结果和profile定位问题 |
| 更适合的场景 | 编译器暂不支持、访存模式不规则或需要精细控制的算子 | 语义可表达、融合边界清楚且需要重复适配的算子 |
两条路线并不冲突。已有的Triton/TBE实现可以作为基线,也可以在自动生成结果不理想时承担兜底。Auto-Kernel的价值,是先接走大量重复搜索工作,让工程师把时间留给编译器暂时处理不好的热点。
RMSNorm和Gated-SiLU怎么走一遍
这类融合算子的处理顺序可以压缩成一条线:
Profile → 划定融合边界 → 定义语义和约束 → 自动搜索 → 正确性回归 → 端到端复测
执行时还要补上几个容易漏掉的动作:
- 在同一模型、shape、batch和CANN版本下保存基线;
- 用CANN 8 profiling记录算子耗时、内存占用和下发延迟;
- 确认RMSNorm输出是否被Gated-SiLU直接消费,再决定是否融合;
- 先检查数值一致性,再比较单算子耗时;
- 回到完整模型,检查吞吐、端到端延迟和线上关心的P99指标;
- 对常用shape和动态shape分别回归,避免只优化一个测试点。
如果融合后的单算子更快,但完整模型没有明显变化,就要回头查看通信、下发或其他热点。继续围着同一个kernel打磨,未必能改善业务结果。
手调留给最后剩下的热点
Triton和TBE仍然有保留价值。遇到编译器无法表达的计算、搜索结果跨shape波动较大,或者项目确实需要精细控制底层调度时,手写实现依然更直接。
对于常规算子组合,可以先建立自动搜索基线。只有profile确认热点仍然存在,工程师再进入Tiling和数据搬运细节。这样一来,手工调优会从默认动作变成针对少数问题的精确投入。
国产算力迁移最怕反复做同一轮参数试验。更省人力的办法,是给手调设置门槛:先拿到CANN 8的profile证据,再试Auto-Kernel,最后才重写底层实现。
参考资料: