之前有个同事问我：“ops-nn 和 ops-math 有什么区别？激活函数放哪个仓库？”
我答不上来。翻了一圈文档才发现，CANN 的算子仓库划分逻辑不是按"功能"来的，是按"用途"来的。ops-math 存基础数学运算，ops-nn 存神经网络专用的算子。但有些算子两边都有，比如 GELU——ops-math 里有，ops-nn 里也有。
我花了两天把 ops-nn 的算子按场景捋了一遍，这篇整理出来，免得你再走弯路。
ops-nn 的定位：神经网络专用算子
ops-nn 是 CANN 核心算子仓库之一，存放的是神经网络训练和推理中高频使用的算子。
和 ops-math 的边界：
| 仓库 | 定位 | 典型算子 |
|---|---|---|
| ops-math | 通用数学运算 | Add/Mul/ReduceSum/Log/Exp |
| ops-nn | 神经网络专用 | MatMul/Activation/Conv/BatchNorm/Softmax |
| ops-blas | 线性代数专用 | GEMM/GEMV |
| ops-transformer | 大模型专用 | FlashAttention/MoE |
重叠区域：GELU、Softmax 这类算子在 ops-math 和 ops-nn 都有。区别是 ops-nn 的版本做了神经网络场景的优化（比如和 Conv/BatchNorm 的融合），ops-math 的版本更通用。
简单判断：如果你的算子只用在神经网络前向/反向传播里，用 ops-nn 的版本。如果通用计算也要用，用 ops-math 的版本。
按场景选算子
场景1：模型推理——激活函数怎么选？
ops-nn 里最常见的激活函数：
| 算子 | 适用场景 | 显存占用 | 计算速度 |
|---|---|---|---|
| ReLU | CNN 模型 | 最低 | 最快 |
| GELU | Transformer 模型 | 中等 | 中等 |
| SiLU/Swish | LLaMA 等 LLM | 中等 | 中等 |
| FastGELU | GELU 的近似实现 | 最低 | 最快 |
关键区别：GELU 和 FastGELU。
GELU 的精确公式是x * Φ(x)，其中 Φ 是标准正态分布的累积分布函数。计算时要用到 erf 函数，开销比较大。FastGELU 用一个多项式近似替代 erf，精度损失很小但速度快很多。
import torch import time def bench_activation(fn, x, warmup=5, runs=100): # 预热 for _ in range(warmup): _ = fn(x) torch.npu.synchronize() times = [] for _ in range(runs): start = time.time() _ = fn(x) torch.npu.synchronize() times.append(time.time() - start) return sum(times) / len(times) * 1000 # 测试数据 x = torch.randn(1, 4096, 4096, dtype=torch.float16).npu() # GELU（精确版） gelu_time = bench_activation(torch.nn.functional.gelu, x) # FastGELU（近似版，ops-nn 里的实现） # PyTorch 里通过近似模式调用 fast_gelu_time = bench_activation( lambda x: torch.nn.functional.gelu(x, approximate='tanh'), x ) # SiLU silu_time = bench_activation(torch.nn.functional.silu, x) # ReLU relu_time = bench_activation(torch.nn.functional.relu, x) print(f"{'算子':<12} {'延迟(ms)':<12} {'相对ReLU':<10}") print("-" * 34) print(f"{'ReLU':<12} {relu_time:.3f}{' '*6} {'1.00x':<10}") print(f"{'GELU':<12} {gelu_time:.3f}{' '*6} {gelu_time/relu_time:.2f}x") print(f"{'FastGELU':<12} {fast_gelu_time:.3f}{' '*6} {fast_gelu_time/relu_time:.2f}x") print(f"{'SiLU':<12} {silu_time:.3f}{' '*6} {silu_time/relu_time:.2f}x")实测结果：
| 算子 | 延迟(ms) | 相对 ReLU |
|---|---|---|
| ReLU | 0.082 | 1.00x |
| FastGELU | 0.119 | 1.45x |
| SiLU | 0.143 | 1.74x |
| GELU | 0.198 | 2.41x |
结论：如果你的模型用 GELU，换成 FastGELU 几乎不损失精度，但能快 40%。LLaMA 系列用的 SiLU 没法替换，但性能也还好。
场景2：训练加速——BatchNorm 和 Conv 怎么融合？
ops-nn 最重要的优化之一是Conv + BatchNorm + Activation 三合一融合。
推理时，BatchNorm 的参数可以提前融合到 Conv 的权重里，不需要单独计算。这是老套路了，但 ops-nn 做了更进一步：融合后的算子在 Cube 单元上一次完成，不需要中间结果落回显存。
import torch import torch.nn as nn import time # ========== 未融合版本 ========== class ConvBnRelu(nn.Module): def __init__(self, in_ch, out_ch, kernel_size=3): super().__init__() self.conv = nn.Conv2d(in_ch, out_ch, kernel_size, padding=1, bias=False) self.bn = nn.BatchNorm2d(out_ch) self.relu = nn.ReLU(inplace=True) def forward(self, x): x = self.conv(x) x = self.bn(x) x = self.relu(x) return x # ========== 融合版本（推理专用） ========== class FusedConvBnRelu(nn.Module): def __init__(self, conv_bn_relu_module): super().__init__() # 把 BN 的参数融合到 Conv 权重里 conv = conv_bn_relu_module.conv bn = conv_bn_relu_module.bn # 融合公式：W_fused = W * (gamma / sqrt(var + eps)) # b_fused = beta - gamma * mean / sqrt(var + eps) gamma = bn.weight.data beta = bn.bias.data mean = bn.running_mean.data var = bn.running_var.data eps = bn.eps scale = gamma / torch.sqrt(var + eps) self.weight = nn.Parameter(conv.weight.data * scale.reshape(-1, 1, 1, 1)) self.bias = nn.Parameter(beta - mean * scale) self.relu = nn.ReLU(inplace=True) def forward(self, x): # 融合后只需要一次 Conv + ReLU x = torch.nn.functional.conv2d(x, self.weight, self.bias, padding=1) x = self.relu(x) return x # ========== 性能对比 ========== model_unfused = ConvBnRelu(64, 128).eval().npu() model_fused = FusedConvBnRelu(model_unfused).eval().npu() x = torch.randn(1, 64, 224, 224).npu() def benchmark(model, x, warmup=10, runs=50): for _ in range(warmup): _ = model(x) torch.npu.synchronize() times = [] for _ in range(runs): start = time.time() _ = model(x) torch.npu.synchronize() times.append(time.time() - start) return sum(times) / len(times) * 1000 unfused_time = benchmark(model_unfused, x) fused_time = benchmark(model_fused, x) print(f"未融合: {unfused_time:.3f}ms") print(f"融合后: {fused_time:.3f}ms") print(f"提升: {(1 - fused_time/unfused_time)*100:.1f}%") # 验证精度 with torch.no_grad(): out_unfused = model_unfused(x) out_fused = model_fused(x) max_diff = (out_unfused - out_fused).abs().max().item() print(f"最大精度差异: {max_diff:.6f}") # 应该 < 1e-4实测结果：
| 配置 | 延迟(ms) | 提升比例 |
|---|---|---|
| Conv + BN + ReLU（未融合） | 1.82ms | 基线 |
| Conv + BN + ReLU（融合后） | 1.24ms | 32% |
| 最大精度差异 | 0.000031 | 可忽略 |
场景3：显存优化——Softmax 的 FP16 陷阱
Softmax 看起来简单，但在 FP16 下有个经典陷阱：指数溢出。
import torch # FP16 Softmax 溢出演示 x = torch.tensor([[1.0, 2.0, 100.0]], dtype=torch.float16).npu() # 直接 softmax（FP16 下会溢出） try: result = torch.softmax(x, dim=-1) print(f"FP16 softmax: {result}") # 可能出现 NaN 或不正确的分布 except Exception as e: print(f"Error: {e}") # 正确做法：先减最大值再 softmax（ops-nn 的实现已经内置了这个） x_safe = x - x.max(dim=-1, keepdim=True).values result_safe = torch.softmax(x_safe.float(), dim=-1).half() print(f"Safe softmax: {result_safe}")ops-nn 的 Softmax 实现已经内置了数值稳定处理，不需要你手动减最大值。但如果你自己写 kernel，必须注意这个问题。
在昇腾 NPU 上，ops-nn 的 Softmax 还做了一个额外优化：分块计算。长序列的 Softmax 不需要把整个向量加载到本地内存，可以分块求 max、分块求 exp 再归一化。这个优化和 FlashAttention 的分块思想一样。
# 长序列 Softmax 性能测试 import time def bench_softmax(seq_len, warmup=5, runs=50): x = torch.randn(1, 32, seq_len, dtype=torch.float16).npu() for _ in range(warmup): _ = torch.softmax(x, dim=-1) torch.npu.synchronize() times = [] for _ in range(runs): start = time.time() _ = torch.softmax(x, dim=-1) torch.npu.synchronize() times.append(time.time() - start) return sum(times) / len(times) * 1000 seq_lengths = [512, 1024, 2048, 4096, 8192] print(f"{'序列长度':<12} {'延迟(ms)':<12}") print("-" * 24) for s in seq_lengths: t = bench_softmax(s) print(f"{s:<12} {t:.3f}")实测结果：
| 序列长度 | 延迟(ms) |
|---|---|
| 512 | 0.015 |
| 1024 | 0.028 |
| 2048 | 0.053 |
| 4096 | 0.101 |
| 8192 | 0.196 |
延迟和序列长度基本线性关系，说明分块优化生效了。如果没有分块，8192 的延迟应该是 512 的 256 倍（N²），实际只有 13 倍。
完整代码：ops-nn 算子选型速查
把上面的测试串起来，做一个自动选型脚本：
""" ops-nn 算子选型工具 根据模型类型和场景推荐算子组合 """ import torch def recommend_ops(model_type, scenario): """根据模型类型和场景推荐算子""" recommendations = { ("transformer", "inference"): { "activation": "FastGELU（比 GELU 快 40%，精度损失 < 0.1%）", "softmax": "ops-nn Softmax（内置分块 + 数值稳定）", "normalization": "LayerNorm（推理时权重可预融合）", "attention": "ops-transformer FlashAttention（不是 ops-nn，但必须提）", "tip": "GELU 换 FastGELU 是最划算的单项优化" }, ("transformer", "training"): { "activation": "GELU（训练需要精确梯度，不要用 FastGELU）", "softmax": "ops-nn Softmax", "normalization": "LayerNorm", "attention": "FlashAttention（训练和推理都建议用）", "tip": "训练时精度优先，推理时速度优先" }, ("cnn", "inference"): { "activation": "ReLU（最快，CNN 不需要 GELU）", "softmax": "ops-nn Softmax", "normalization": "BatchNorm（权重预融合到 Conv 里）", "conv_bn_relu": "融合算子（省 30% 延迟）", "tip": "BatchNo
