ops-elementwise：小算子的融合艺术

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 次 Launchz=x+y;// Launch Addw=z*scale;// Launch Mulu=w+bias;// Launch Add// 融合：1 次 Launchu=(x+y)*scale+bias;// Launch Fusion

融合后的 Launch 开销减少了 66%。Vector Unit 在 L1 上流水线执行加法和乘法——x + y的结果在 Vector 寄存器中直接传给* scale，不需要写 DDR 再读。

昇腾NPU如何减少 Kernel 启动

ops-elementwise 的融合策略：

1. 连续 ElementWise 融合。检测连续出现的 Add、Mul、Scale、Bias 等操作，合并为一个 Composite Kernel
2. 激活函数融合。GELU 这种需要多个 Vector 指令的激活函数也可以跟前面的 ElementWise 合并——x * scale + bias → GELU
3. 与 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 融合前后的对比：

Launch 开销减半是直接收益。另外的间接收益是中间 Tensor 搬运减少——融合后只要原来一半的 ElementWise 中间结果写 DDR。

参考仓库

ops-elementwise 逐元素算子库

graph-autofusion 自动融合框架