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 的融合策略:

  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 融合前后的对比:

配置 单 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 自动融合框架

Logo

鲲鹏昇腾开发者社区是面向全社会开放的“联接全球计算开发者,聚合华为+生态”的社区,内容涵盖鲲鹏、昇腾资源,帮助开发者快速获取所需的知识、经验、软件、工具、算力,支撑开发者易学、好用、成功,成为核心开发者。

更多推荐