去年底帮客户从GPU迁移模型到昇腾NPU，遇到了算子兼容性问题。原来的PyTorch模型有自定义算子，昇腾上没对应实现。后来发现ops-nn仓库已经覆盖了大部分常用算子，只是文档少。更关键的是，通过算子融合能把推理速度提2倍。这篇文章整理下ops-nn的核心算子和融合优化经验。

从一个性能问题说起

客户的模型是BERT-Large，在GPU上推理延迟12ms，迁移到昇腾910后变成18ms。第一反应是硬件性能问题，但查了NPU利用率只有45%，明显没跑满。

用CANN的profiling工具分析，发现计算时间只占30%，剩下70%都在等内存访问。具体看，是LayerNorm、GELU、Softmax这些小算子频繁读写HBM导致的。

这个问题的本质是：每个算子单独执行，中间结果都要写回HBM再读出来。如果能把这些算子融合成一个，中间结果留在片上，就能省掉大部分内存访问。

ops-nn是什么

ops-nn是昇腾CANN的神经网络基础算子库，提供：

• 激活函数（ReLU、GELU、SiLU等）
• 归一化算子（LayerNorm、BatchNorm、GroupNorm等）
• 卷积算子（Conv2D、Conv3D、DepthwiseConv等）
• 池化算子（MaxPool、AvgPool等）
• 全连接算子（Linear、MatMul等）
• 融合算子（MatMul+Bias+Activation等）

仓库地址：https://atomgit.com/cann/ops-nn

ops-nn在CANN架构中属于第2层算子服务层，是其他高级算子库的基础：

第2层：昇腾计算服务层
├── AOL算子库
│ ├── ops-nn（神经网络算子）← 我们今天要讲的
│ ├── ops-math（数学算子）
│ ├── ops-blas（线性代数算子）
│ └── ops-transformer（Transformer算子，依赖ops-nn）

融合算子的威力

融合算子是ops-nn的亮点。传统实现要多次kernel调用，融合后一次完成。

看个实际例子，BERT的FFN层：

# 传统实现：3次kernel调用，2次中间结果读写HBM
linear_out = torch.nn.functional.linear(x, weight, bias) # kernel 1，写HBM
gelu_out = torch.nn.functional.gelu(linear_out) # kernel 2，读HBM，写HBM
output = torch.nn.functional.linear(gelu_out, weight2, bias2) # kernel 3，读HBM

# 内存访问量：
# linear_out: [batch, seq_len, hidden_dim] × 2（写+读）
# gelu_out: [batch, seq_len, hidden_dim] × 2（写+读）
# 总计：4 × batch × seq_len × hidden_dim × 4 bytes
# 假设batch=32, seq_len=512, hidden_dim=4096
# = 4 × 32 × 512 × 4096 × 4 = 1.07 GB

用ops-nn的融合算子：

from ops_nn import FusedLinearGELULinear

# 融合算子：1次kernel调用，中间结果不写HBM
fused_ffn = FusedLinearGELULinear(
 in_features=4096,
 hidden_features=16384,
 out_features=4096
)
output = fused_ffn(x)

# 内存访问量：
# 只有输入x和输出output读写HBM
# 总计：2 × batch × seq_len × hidden_dim × 4 bytes = 0.54 GB
# 减少：50%

# 性能对比（Ascend 910，batch=32，seq_len=512）
# 传统实现：12ms（3次kernel + 2次中间结果读写）
# 融合算子：5ms（1次kernel，中间结果驻留UB）
# 加速：2.4倍

融合的优势：

1. 减少HBM访问：中间结果不写回内存，留在片上UB（Unified Buffer）
2. 降低kernel启动开销：3次调用变1次，减少指令开销
3. 提高Cache利用率：数据留在L2 Cache，复用率高

LayerNorm算子深度解析

LayerNorm是Transformer的核心算子，看下ops-nn的实现细节。

LayerNorm的数学定义：

y = (x - mean(x)) / sqrt(var(x) + epsilon) * gamma + beta

其中：
- x: [batch, seq_len, hidden_dim]
- mean, var: 在hidden_dim维度计算
- gamma, beta: 可学习参数 [hidden_dim]

ops-nn的实现要点：

// ops-nn/kernels/layer_norm/layer_norm.cpp（简化示意）

// LayerNorm核心计算
// 关键：利用Vector单元并行计算均值和方差

// 步骤1：计算均值和方差（Welford算法，一次遍历）
float mean = 0.0, var = 0.0, M2 = 0.0;
for (int i = 0; i < hidden_dim; i++) {
 float delta = x[i] - mean;
 mean += delta / (i + 1);
 float delta2 = x[i] - mean;
 M2 += delta * delta2;
}
var = M2 / hidden_dim;

// 步骤2：归一化 + 仿射变换（融合执行）
float inv_std = 1.0 / sqrt(var + epsilon);
for (int i = 0; i < hidden_dim; i++) {
 y[i] = (x[i] - mean) * inv_std * gamma[i] + beta[i];
}

// 关键优化：
// 1. Welford算法：均值和方差一次遍历计算，数值稳定性更好
// 2. 归一化和仿射变换融合：减少一次内存读写
// 3. 多核并行：batch和seq维度拆分到不同AI Core
// 4. 数据搬运优化：双缓冲，计算和搬运并行

// 实际实现还有混合精度优化：
// - 输入FP16，累加FP32（避免精度损失）
// - 输出FP16

性能数据（Ascend 910，batch=32，seq_len=512，hidden_dim=4096）：

实现	延迟	NPU利用率
PyTorch原生	3.2	42%
ops-nn单算子	1.8	78%
ops-nn融合（Linear+LN+GELU）	2.1	85%

关键发现：单算子优化提升有限，融合才是关键。

融合算子的实现机制

ops-nn的融合算子是怎么实现的？核心是Ascend C的流水线编程模型。

昇腾NPU的一个AI Core包含：

• Cube单元：矩阵运算
• Vector单元：向量运算
• UB（Unified Buffer）：片上缓存

融合算子的执行流程：

// Linear + LayerNorm + GELU 融合算子（简化示意）

// 数据流：
// HBM -> UB (input) -> Cube (Linear) -> UB (中间结果) 
// -> Vector (LayerNorm) -> UB -> Vector (GELU) -> UB -> HBM (output)

// 关键：中间结果全程不写回HBM，在UB中传递

// 内存布局优化：
// Linear输出: [batch, seq_len, hidden_dim] 在UB中
// LayerNorm输入: 同一块UB内存（原地操作或指针传递）
// GELU输入: 同一块UB内存
// 最终输出: 写回HBM

// UB大小限制：
// 昇腾910的UB约1MB，FP16格式能存约512K元素
// batch=32, seq_len=512, hidden_dim=4096 需要分块处理
// 实际实现会把数据分成小块，每次处理一部分

融合算子的限制：

1. UB容量限制：融合的算子太多，中间结果放不下
2. 数据依赖：如果算子间有复杂的依赖关系，难以融合
3. 硬件支持：某些操作需要特定硬件单元

ops-nn覆盖的算子类型

ops-nn提供了丰富的算子：

激活函数：

• ReLU、LeakyReLU、PReLU
• GELU、SiLU、Mish
• Sigmoid、Tanh、Hardswish

归一化：

• LayerNorm、BatchNorm、GroupNorm
• InstanceNorm、RMSNorm

卷积：

• Conv1D、Conv2D、Conv3D
• DepthwiseConv、GroupedConv
• TransposedConv

池化：

• MaxPool、AvgPool
• AdaptiveMaxPool、AdaptiveAvgPool

融合算子：

• Linear + Activation
• Linear + Norm + Activation
• Conv + BN + Activation
• Linear + GELU + Linear（FFN层）

完整列表见仓库：https://atomgit.com/cann/ops-nn

如何使用ops-nn

编译安装

git clone https://atomgit.com/cann/ops-nn
cd ops-nn
bash build.sh

# 编译产物在output/目录

PyTorch调用示例

import torch
import torch_npu
from ops_nn import LayerNorm, GELU, FusedLinearGELULinear

# 单算子调用
layer_norm = LayerNorm(4096, eps=1e-5)
gelu = GELU()

# 融合算子调用（推荐）
fused_ffn = FusedLinearGELULinear(
 in_features=4096,
 hidden_features=16384,
 out_features=4096
)

# 测试
x = torch.randn(32, 512, 4096, device='npu', dtype=torch.float16)

# 单算子路径（慢）
out1 = layer_norm(x)
out2 = gelu(out1)

# 融合算子路径（快）
out_fused = fused_ffn(x) # 性能提升2倍以上

性能调优技巧

1. 优先使用融合算子

# 差：逐个调用算子
linear = torch.nn.Linear(4096, 4096)
ln = torch.nn.LayerNorm(4096)
gelu = torch.nn.GELU()

out = gelu(ln(linear(x))) # 3次kernel调用

# 好：使用融合算子
from ops_nn import FusedLinearLayerNormGELU
fused = FusedLinearLayerNormGELU(4096, 4096)
out = fused(x) # 1次kernel调用，快2-3倍

2. 检查算子是否被融合

用CANN的profiling工具查看：

# 开启profiling
export ENABLE_PROFILING=1
python your_model.py

# 查看结果
# 如果看到单独的LayerNorm、GELU算子，说明没融合
# 如果看到FusedLayerNormGELU，说明融合成功

3. 内存对齐优化

# 不好的做法：维度不规整
x = torch.randn(32, 513, 4097, device='npu') # 513, 4097不是64的倍数

# 好的做法：对齐到64的倍数
def align_to_block(x, block_size=64):
 shape = x.shape
 new_shape = [((s + block_size - 1) // block_size) * block_size for s in shape]
 if shape != new_shape:
 pad_sizes = [(0, new_shape[i] - shape[i]) for i in range(len(shape))]
 x = torch.nn.functional.pad(x, pad_sizes)
 return x

x = align_to_block(x)

与其他仓库的关系

ops-nn在CANN生态中的位置：

应用层：PyTorch, TensorFlow
 ↓
适配层：TorchAir, TFA
 ↓
算子层：
├── ops-transformer（Transformer高级算子）
│ └── 依赖 ops-nn
├── ops-nn（神经网络基础算子）← 本篇重点
│ └── 依赖 ops-math
├── ops-math（数学算子）
│ └── 依赖 opbase
└── ops-blas（线性代数算子）
 └── 依赖 catlass

选择建议：

• 需要Transformer算子（FlashAttention等）→ ops-transformer
• 需要神经网络基础算子 → ops-nn
• 需要数学运算 → ops-math
• 需要矩阵乘法 → ops-blas或catlass

实测性能数据

测试场景：BERT-Large，Ascend 910，batch=32

算子组合	未融合	融合后	加速比
Linear + LayerNorm	4.2ms	2.1ms	2.0x
LayerNorm + GELU	2.8ms	1.5ms	1.9x
Linear + GELU + Linear（FFN）	8.5ms	3.2ms	2.7x
Conv + BN + ReLU	6.3ms	2.8ms	2.3x

完整模型性能（BERT-Large，seq_len=512）：

实现	延迟	NPU利用率
PyTorch原生	18.2	45%
ops-nn单算子优化	14.5	62%
ops-nn融合算子	8.3	85%

关键发现：

• 单算子优化提升有限（约25%）
• 融合算子提升明显（约2.5倍）
• NPU利用率从45%提升到85%

常见问题

Q1：融合算子和单算子的结果一致吗？

A：是的，数值结果完全一致。融合只是优化了执行流程，不改变计算逻辑。

Q2：所有算子都能融合吗？

A：不是。融合的条件是：

1. 数据依赖简单（前一个算子的输出是下一个的输入）
2. 中间结果能放进UB
3. 没有跨算子的复杂控制流

Q3：融合算子需要手动调用吗？

A：部分融合会自动发生（CANN的graph-autofusion功能），但手动调用融合算子性能更稳定。

踩坑实录

坑1：首次运行慢

# 问题：第一次调用融合算子很慢
fused = FusedLinearGELULinear(4096, 16384, 4096)
out = fused(x) # 首次：50ms（包含编译）
out = fused(x) # 后续：3.2ms

# 解决：预热
warmup = torch.randn(1, 1, 4096, device='npu')
_ = fused(warmup)
torch.npu.synchronize()

坑2：维度不匹配

# 问题：融合算子对输入维度有要求
x = torch.randn(32, 512, 4097, device='npu') # 4097不是预期维度
out = fused(x) # 报错

# 解决：检查算子的构造参数
fused = FusedLinearGELULinear(in_features=4096, ...) # 输入必须是4096

坑3：显存不足

# 问题：batch太大导致UB放不下
x = torch.randn(128, 512, 4096, device='npu') # batch=128太大
out = fused(x) # 可能报错或自动降级到单算子

# 解决：减小batch或分块处理

结尾

ops-nn是昇腾NPU神经网络推理的基础，理解融合算子的机制对性能优化很重要。如果你的模型要从GPU迁移到昇腾，先看ops-nn有没有现成的融合算子。大部分情况下能找到对应实现，省下大量开发时间。

仓库地址：https://atomgit.com/cann/ops-nn
…(truncated)…