ops-nn：昇腾 NPU 上的神经网络算子库，你每天都在用它

很多人用昇腾 NPU 跑模型，调的是 ATB，用的是 PyTorch NPU 后端，但真正干苦力活的底层算子，大部分来自一个叫 ops-nn 的仓库。

这个名字很朴素——"nn"就是 neural network，神经网络。它不花哨，不性感，就是老老实实提供 matmul、activation、softmax 这些最基础的神经网络算子。但恰恰是这些"不性感"的算子，撑起了整个昇腾 CANN 生态的底层算力。

今天把 ops-nn 说清楚。

---

ops-nn 是什么，在 CANN 里排第几

先说位置。昇腾异构计算架构的第一层是 AscendCL（应用开发接口），第二层是 AOL 算子库，ops-nn 就是 AOL 里最核心的那个库。

具体来说，昇腾 CANN 开源社区里，算子相关的仓库有好几个：

仓库名	定位
ops-nn	神经网络类基础算子（matmul、activation、softmax 等）
ops-math	数学类算子（conversion、random 等）
ops-blas	线性代数基础算子（高性能 GEMM）
ops-cv	计算机视觉类算子
ops-transformer	Transformer 类大模型算子
ops-tensor	张量操作类算子

ops-nn 的特点是**“基础"和"高频”**——几乎每个神经网络模型都会用到它里面的算子。matmul（矩阵乘法）就不用说了，activation（ReLU、GELU、SiLU）也是每层都有的。

如果你打开一个 PyTorch 模型跑在昇腾 NPU 上，底层调用链是这样的：

PyTorch NPU 后端
  └─ AscendCL（算子调用接口）
      └─ ops-nn 的 matmul / softmax / activation 算子
          └─ catlass（GEMM 模板，ops-nn 依赖它）
              └─ Ascend C（算子编程语言）

看出来了吗？ops-nn 是"前台"，catlass 是"后台"。ops-nn 里的 matmul 算子，底层就是基于 catlass 的 GEMM 模板实现的。这就是为什么我之前说"理解 catlass 和 ops-nn 的关系很重要"——它们是上下游。

---

ops-nn 里到底有哪些算子

ops-nn 的算子按功能分几大类，我列一下最常用的：

矩阵运算类

• MatMul：标准矩阵乘法，支持 float16 / bfloat16 / float32
• BatchMatMul：批量矩阵乘法，Transformer 的 QKV 投影用这个
• MatMulV2：MatMul 的增强版，支持广播、转置等更多特性

激活函数类

• ReLU / ReLUGrad：最基础的激活函数
• GELU / FastGELU：大模型最常用的激活，Llama 系列全用 GELU
• SiLU / Swish：SiLU 是另一个大模型常用激活（Qwen 系列用这个）
• Softmax / LogSoftmax：Attention 里的标配

归一化类

• LayerNorm / RMSNorm：Transformer 里的归一化，RMSNorm 是 Llama 用的
• BatchNorm：CNN 模型常用

融合算子类（这是 ops-nn 的性能杀手锏）

• MatMulLeakyRelu：MatMul + LeakyReLU 融合
• MatMulAdd：MatMul + Add 融合
• SoftmaxWithCrossEntropy：Softmax + 交叉熵融合

融合算子的意义之前说过：减少显存读写次数。分开跑 MatMul 和 ReLU，中间结果要写回显存；融合成一个算子，中间结果留在 SRAM 里直接算完，省一次显存读写。

---

实际用一下：从 PyTorch 到 ops-nn 底层

大多数时候你不需要直接调 ops-nn 的接口，PyTorch NPU 后端会自动把 torch.matmul 映射到 ops-nn 的 MatMul 算子。但有些场景下，你需要直接调 ops-nn 的接口——比如你要写一个自定义算子，或者你要验证某个算子的性能。

ops-nn 提供了 AscendCL 接口和 Python 接口两套调用方式。大多数开发者用 Python 接口就够了。

# 方式一：通过 PyTorch 自动映射（大多数场景）
import torch

q = torch.randn(1, 32, 4096, 128, dtype=torch.float16, device="npu")
k = torch.randn(1, 32, 4096, 128, dtype=torch.float16, device="npu")
# PyTorch 的 matmul 会自动调用 ops-nn 的 MatMul 算子
attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(128)

# 方式二：直接调用 ops-nn 的 Python 接口（自定义或调试场景）
from ops_nn import MatMul, Softmax

# 直接调 ops-nn 算子，跳过 PyTorch 的一些封装开销
q_npu = npu_tensor(q)
k_npu = npu_tensor(k)
attn_scores = MatMul()(q_npu, k_npu.transpose(-2, -1))
attn_probs = Softmax()(attn_scores, axis=-1)

哪种方式好？大多数场景用方式一，简单、稳定、PyTorch 生态兼容。但如果你在做什么性能调优，发现 PyTorch 的封装开销占比显著（比如小 batch、小矩阵的高频调用），可以试方式二，直接调 ops-nn 的 Python 接口，能省一点开销。

---

ops-nn 和 catlass 的关系（再强调一次）

之前写过 catlass 的文章，这里再强调一次关系，因为这个关系太重要了：

ops-nn 的 matmul 算子，底层实现依赖 catlass 的 GEMM 模板。

具体来说：

• catlass 提供高性能 GEMM 模板（自动 tiling、双缓冲 DMA、指令调度）
• ops-nn 基于 catlass 模板，封装成用户可调用的 matmul 算子，加上输入输出校验、dtype 转换、边界处理等
• 最终用户通过 AscendCL 或 PyTorch 调用 ops-nn 的算子

所以如果你发现 ops-nn 的某个 matmul 算子性能不够好，有两个方向可以挖：

1. 算子层面：ops-nn 的封装有没有额外开销（比如不必要的 dtype 转换）
2. 模板层面：catlass 的 GEMM 模板参数有没有调优（BlockM/BlockN/BlockK 是否合理）

大部分场景下，catlass 的默认模板参数已经够好了。但如果你在做极致的性能调优，这两个层都要看。

---

ops-nn 和 ops-blas 的区别

又是一个容易混淆的点。ops-nn 里有 matmul，ops-blas 里也有 GEMM，它们有什么区别？

简单说：ops-nn 面向深度学习场景，ops-blas 面向科学计算场景。

维度	ops-nn	ops-blas
主要场景	深度学习模型训练/推理	科学计算、线性代数
支持的 dtype	float16 / bfloat16 / float32	float32 / float64 / complex
优化重点	和深度学习框架适配（PyTorch 等）	和 BLAS 标准接口兼容
典型算子	MatMul、Softmax、LayerNorm	GEMM、GEMV、AXPY

如果你在跑深度学习模型，用 ops-nn 就够了，PyTorch 后端会自动映射。如果你在做什么科学计算（比如矩阵分解、特征值计算），那 ops-blas 更合适。

---

几个踩坑经验

坑一：activation 算子的 dtype 要匹配。 ops-nn 的 GELU / SiLU 这些激活算子，输入 dtype 是 float16 的话，输出也是 float16。但有些模型（特别是用 float32 训练出来的）期望激活函数的输出是 float32。这时候如果直接调 ops-nn 的 GELU，输出会被截断成 float16，精度掉了。解决办法是在调激活之前显式做类型转换：

# 正确做法：先转 dtype，再调激活
x = x.to(torch.float16)
x = torch.nn.functional.gelu(x)  # 调 ops-nn 的 GELU
x = x.to(torch.float32)          # 转回 float32

坑二：融合算子的数值稳定性。 ops-nn 里有一些融合算子（比如 SoftmaxWithCrossEntropy），为了性能做了一些数值优化，在极端输入下数值稳定性和标准实现可能有细微差异。如果你在训练对数值精度要求极高的模型（比如做科学计算的神经网络），建议分开调 Softmax 和 CrossEntropy，不要用融合版本。

坑三：BatchMatMul 的 batch 维度别搞错。 torch.matmul 在输入是 3D 张量的时候，会自动做 batch matmul。但 ops-nn 的 BatchMatMul 对 batch 维度的位置有要求——必须是第 0 维。如果你的张量是 (heads, batch, seq, dim) 这种格式，直接调 BatchMatMul 会报错或者结果不对。解决办法是先用 transpose 把 batch 维换到第 0 维。

---

结尾

ops-nn 是昇腾 CANN 生态里最基础的神经网络算子库，它不性感，但每个跑在昇腾 NPU 上的模型都依赖它。

理解 ops-nn 的价值不在于"怎么直接调它的接口"（大多数时候你不需要），而在于理解你的模型底层在调哪些算子、这些算子的性能特征是什么、有没有融合版本可以用。

如果你在做模型性能调优，建议先用 profiler 工具看看你的模型时间花在哪些算子上了，再针对性地看 ops-nn 里有没有对应的融合算子可以用。源码在 https://atomgit.com/cann/ops-nn