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，调度开销分摊到多个操作上。

cpp 复制代码

// 不融合：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。

实际推理中的收益

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。

参考仓库

ops-elementwise 逐元素算子库

graph-autofusion 自动融合框架