CANN-ops-nn和ops-transformer-昇腾NPU两个算子仓库怎么分工

CANN-ops-nn和ops-transformer-昇腾NPU两个算子仓库怎么分工

刚接触 CANN 的人最常问：ops-nn 和 ops-transformer 到底谁管什么？我在 PyTorch 里写F.linear，底层走的是 ops-nn 还是 ops-transformer？答案是：走 ops-nn，除非 ATB 或 graph-autofusion 把你的调用链改造成了融合路径。

分工逻辑

ops-nn：通用神经网络算子。不管你跑什么模型——CNN、RNN、Transformer、扩散模型——都用 ops-nn 的基础算子。覆盖广，单个算子性能优化到位，但不会替你做跨算子的融合编排。

ops-transformer：Transformer 专属融合算子。只管 Transformer 架构里的热点计算模式。覆盖窄，但每个算子都是高度定制的融合实现。

用一句话概括：ops-nn 是工具箱，ops-transformer 是定制家具。工具箱里的锤子螺丝刀通用但需要你自己组合；定制家具直接放到指定位置就能用，但只适合特定的房间格局。

上游依赖关系

ops-transformer 的融合算子内部调用了 ops-nn 的基础算子。以 FlashAttention 为例：

FlashAttention kernel 内部： → Cube 单元：MatMul（调用 ops-blas 的 GEMM 指令） → Vector 单元：Softmax（调用 ops-math 的 Softmax 实现） → Cube 单元：MatMul（再次调用 GEMM） → Vector 单元：RoPE（调用 ops-nn 的 RotaryEmbedding 基础实现）

ops-transformer 不是从零造轮子。它把 ops-nn/ops-math/ops-blas 的基础算子按照 FlashAttention 的计算模式编排成一个 kernel，但底层的基本计算单元还是复用的。

依赖链：

opbase → ops-blas/ops-math/ops-nn → ops-transformer → ATB ↑ 基础组件 ↑ 原子算子 ↑ 融合编排 ↑ 高层API

实际调用的分界线

当你在昇腾NPU上跑 PyTorch 模型，torch_npu 的算子分发逻辑：

1. 检查 ops-transformer 是否有对应的融合算子
2. 如果有，检查输入是否满足融合条件（维度对齐、dtype 匹配等）
3. 满足 → 走 ops-transformer；不满足 → fallback 到 ops-nn
4. 如果 ops-transformer 没有对应算子 → 走 ops-nn

这个分发过程对用户透明。你写F.scaled_dot_product_attention(q, k, v)，torch_npu 自动选择 FlashAttention（ops-transformer）或标准 Attention（ops-nn）。

什么时候只靠 ops-nn 就够

• 你的模型不是标准 Transformer 架构
• 你在做模型调试，需要逐步检查中间结果
• 你的输入形状不满足融合算子的对齐要求
• 你在训练小模型（融合的 kernel launch 节省对小 batch 不明显）

什么时候必须用 ops-transformer

• 大模型推理（Llama 7B 以上），Attention 和 FFN 是绝对瓶颈
• 长序列场景（序列长度 > 2K），FlashAttention 的 O(N) 显存优势不可替代
• MoE 模型，MergedMatMul 和 MC2 是刚需
• 分布式训练，MC2 的通算融合是唯一能让通信不占 40% 时间的方法

一个容易搞混的例子

torch.nn.functional.linear在昇腾NPU上调的是 ops-nn 的 MatMul + Bias。但如果你在 ATB 里加载一个 Llama 模型，ATB 会把连续的 Q Linear + K Linear + V Linear 替换成 ops-transformer 的 MergedMatMul + RotaryEmbedding 融合算子。

同一个F.linear调用，走 ATB 路径和走 PyTorch eager 路径，底层跑的是不同的算子实现。性能差异可能达到 2-3 倍。

这就是为什么 ATB 的推理服务比纯 PyTorch 快得多——不是 ATB 有什么魔法，是它把 ops-nn 的基础算子替换成了 ops-transformer 的融合算子。

理解了 ops-nn 和 ops-transformer 的分工，遇到性能问题就知道该查哪一层。基础算子慢→查 ops-nn 的 tiling 和融合接口；融合不够→看 ops-transformer 有没有覆盖你的计算模式。两个仓库：

https://atomgit.com/cann/ops-nn