ops-elementwise:小算子的融合艺术
本文探讨了ElementWise算子融合在AI推理中的优化价值。通过将Add、Mul、Sub等小算子合并为复合Kernel,可显著减少调度开销(原开销是计算时间的10-20倍)。典型融合场景包括残差连接、激活函数与Dropout等,昇腾NPU采用连续算子融合、GEMM尾处理融合等策略。实测显示,LLaMA-7B模型在融合后Kernel数量减半,Launch开销降低50%,整体延迟下降7%。这种优化
Add、Mul、Sub、Div——这些逐元素运算的计算量几乎为零,但在推理中出现的频率最高。一个 Transformer Block 里几十次 Add(残差连接、偏置加),几十次 Mul(Attention 的 scale、Dropout 的 mask 乘)。
每个小算子独立 Launch 一次就有 5-15μs 的 Runtime 调度开销。把这些小算子融合到一个 Kernel 里——省掉的 Launch 时间比计算时间还多。
为什么简单算子也会拖慢推理
A + B 的 Vector Unit 执行时间约 0.5μs(4096 个元素的加法)。但独立 Launch 这个算子需要:
- Runtime Task 创建:3-5μs
- Driver 提交:2-5μs
- Vector Unit 执行:0.5μs
调度开销是计算时间的 10-20 倍。
ElementWise 为什么适合融合
ElementWise 算子之间没有数据依赖关系——连续的 Add、Mul、Scale 可以合并成一个复合算子。融合后的 Kernel 只做一次 Launch,调度开销分摊到多个操作上。
// 不融合:3 次 Launch
z = x + y; // Launch Add
w = z * scale; // Launch Mul
u = w + bias; // Launch Add
// 融合:1 次 Launch
u = (x + y) * scale + bias; // Launch Fusion
融合后的 Launch 开销减少了 66%。Vector Unit 在 L1 上流水线执行加法和乘法——x + y 的结果在 Vector 寄存器中直接传给 * scale,不需要写 DDR 再读。
昇腾NPU如何减少 Kernel 启动
ops-elementwise 的融合策略:
- 连续 ElementWise 融合。 检测连续出现的 Add、Mul、Scale、Bias 等操作,合并为一个 Composite Kernel
- 激活函数融合。 GELU 这种需要多个 Vector 指令的激活函数也可以跟前面的 ElementWise 合并——
x * scale + bias → GELU - 与 GEMM 的 Epilogue 融合。 ElementWise 算子作为 GEMM 的 Epilogue——GEMM 算完立即做 Add/Mul/Scale,中间结果不落 DDR
Transformer 中的典型融合场景
残差连接的 Add 融合。 Attention 子层的输出 + 输入残差——这个 Add 不独立 Launch,而是作为 LayerNorm 的输入的一部分,GE 把 Add 融合到 LayerNorm 的 Kernel 中。
Scale + Add(偏置)。 GEMM 输出的 Scale 和 Bias Add 也作为 GEMM 的 Epilogue——GEMM 计算最后一个 Tile 后立即做 Scale 和 Bias,在 L1 上完成。
Dropout 的 Mul 融合。 Dropout 的 mask(ops-rand 生成)跟激活输出的 Mul 融合到上一个算子的 Epilogue 中——激活函数算完后立即跟 mask 相乘,不写 DDR。
更多典型融合案例
Softmax + Mul(Dropout) 的融合:Softmax 的输出是概率分布,Dropout 随机丢弃一些概率值。如果不融合,Softmax 的 [n,n] 输出写入 DDR,Dropout 从 DDR 读入再写出——32MB 的额外搬运。融合后 Dropout 的 mask 在 Softmax 的 exp 结果上直接相乘——Softmax 的分母求和、mask 的逐元素乘在同一个 Kernel 内完成。
Add + Mul + Add 的残差融合:Transformer Block 的 残差输出 = attention_output + input 加完后的结果经过乘法和下一次加法。这三个 ElementWise 向量计算可以合并为一个 Vector Kernel——一次 Launch、一次数据搬运。
实际推理中的收益
Ascend 910 上 LLaMA-7B 的 ElementWise 融合前后的对比:
| 配置 | 单 Block 的 Kernel 数 | Launch 开销 | Block 延迟 |
|---|---|---|---|
| 不融合 | 24 个独立 Kernel | 240μs | 2.85ms |
| ElementWise 融合 | 12 个融合 Kernel | 120μs | 2.65ms |
Launch 开销减半是直接收益。另外的间接收益是中间 Tensor 搬运减少——融合后只要原来一半的 ElementWise 中间结果写 DDR。
参考仓库
鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者,聚合华为+生态”的社区,内容涵盖鲲鹏、昇腾资源,帮助开发者快速获取所需的知识、经验、软件、工具、算力,支撑开发者易学、好用、成功,成为核心开发者。
更多推荐



所有评论(0)