ops-math 踩坑记：那些年我们算过的张量

心中有国也有家

217人浏览 · 2026-05-21 10:36:31

心中有国也有家 · 2026-05-21 10:36:31 发布

ops-math 踩坑记：那些年我们算过的张量

第一次在昇腾NPU上跑 Transformer 推理，精度对不上。不是差很多，就是小数点后三四位的问题。

定位了两天，最后发现是 softmax 那步的数值稳定性问题——CPU上能容忍的写法，在NPU上会被放大。

这才认真看 ops-math 这个仓库。

它到底干了什么

ops-math 是昇腾CANN生态里的数学类基础算子库。跟 ops-nn（神经网络算子）、ops-blas（线性代数）并列，但专注的是更底层的数学运算：conversion类、math类、random类。

从架构上看，它在昇腾计算服务层的AOL算子库里，跟其他 ops-* 系列仓库一样，底层依赖 opbase 提供的基础组件。

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

三个核心模块

1. conversion 类——数据类型转换

昇腾NPUprefer 用 float16 或 bfloat16 计算（快），但输入可能是 float32，输出又要转回 float32。

conversion 类算子就是干这个的。看似简单，但NPU上的转换跟CPU不一样——它要考虑向量计算的对齐问题。

实际踩过的坑：从 float32 转到 float16，如果直接调 Cast，某些边界值（比如大于65504的）会被截断。后来发现 ops-math 里有个 SafeCast 的实现，会先做个clip再转。

# 错误写法（直接转换，可能溢出）
output = acl.ops.Cast(input_tensor, dtype=acl.dtype.float16)

# 正确写法（先clip，再转换）
clipped = acl.ops.ClipByValue(input_tensor, min=-65504, max=65504)
output = acl.ops.Cast(clipped, dtype=acl.dtype.float16)

2. math 类——基础数学函数

这一块是整个算子库最"杂"的部分。Exp、Log、Sqrt、Pow、ReduceSum… 都是看起来简单，但要在NPU上跑得快并不容易的东西。

关键点在于向量化。NPU的达芬奇架构有专门的向量计算单元（Vector Core），math类算子就是最大限度利用这个单元。

举个例子，ReduceSum 在CPU上可能写个for循环就完了。在NPU上，要用到多核并行+向量指令，还要处理不好整除的尾巴数据。

仓库里对应的实现在 ops_math/reduce/ 目录下，核心逻辑是用 acl::kernel::Reduce 这个底层接口。

实测在Ascend 910上，一个 shape 为 [8192, 8192] 的 float16 矩阵做 ReduceSum，ops-math 的实现比手写for循环快约18倍。

3. random 类——随机数生成

这个模块比较特殊。NPU上生成随机数，不能用CPU的 std::rand() 或者 Python 的 random 模块——那样数据要在CPU生成再拷到NPU，慢。

random 类算子是在NPU上直接生成随机数。支持的分布包括均匀分布、正态分布、伯努利分布等。

有个细节：NPU上的随机数生成是可以确定性复现的（给定相同的seed）。这对调试很重要。

import acl

# 在NPU上生成随机数（不需要先把数据拷到CPU）
acl.init()
context, _ = acl.create_context(0)

# 生成正态分布随机数，mean=0, std=1
output_tensor = acl.ops.RandomNormal(shape=(1024, 1024), mean=0.0, std=1.0, seed=42)

# 同样的seed，每次生成的结果一样（便于调试）

跟其他算子库的关系

ops-math 不是孤立的。实际使用中，它经常作为"底层支撑"被调用：

ops-nn 里的激活函数（Softmax、LogSoftmax）会调用 ops-math 的 Exp 和 Log
ops-blas 的矩阵运算前的数据预处理，会调用 ops-math 的 conversion 类
框架适配器（PyTorch Adapter、MindSpore Adapter）在算子映射时，会把部分数学运算路由到 ops-math

这也就是为什么 opbase 要在最底层——所有 ops-* 仓库都依赖它提供的通用组件（内存管理、数据类型定义等）。

一个完整的实战例子

假设要在昇腾NPU上实现一个自定义的损失函数，它包含以下步骤：

对预测值做 Softmax
计算交叉熵
加一个正则化项（L2范数）

如果纯用Python写，可能这样：

import torch
import torch.nn as nn

def custom_loss(pred, target):
    # Softmax + 交叉熵
    log_softmax = torch.nn.functional.log_softmax(pred, dim=-1)
    loss = torch.nn.functional.nll_loss(log_softmax, target)
    
    # L2正则化
    l2_norm = torch.norm(pred, p=2)
    loss = loss + 0.01 * l2_norm
    
    return loss

在NPU上跑，建议改成直接调用 ops-math 的底层算子（避免Python overhead）：

import acl
import acl.ops as ops

def custom_loss_npu(pred_tensor, target_tensor):
    # 直接调 ops-math 的 Exp 和 ReduceSum 实现 Softmax
    max_val = ops.ReduceMax(pred_tensor, dim=-1, keep_dims=True)
    shifted = ops.Sub(pred_tensor, max_val)
    exp_shifted = ops.Exp(shifted)
    sum_exp = ops.ReduceSum(exp_shifted, dim=-1, keep_dims=True)
    softmax_output = ops.Div(exp_shifted, sum_exp)
    
    # 交叉熵（用 Log 避免数值下溢）
    log_softmax = ops.Log(softmax_output)
    loss = ops.Neg(ops.GatherNd(log_softmax, target_tensor))
    
    # L2正则化
    l2_norm = ops.Pow(pred_tensor, 2)
    l2_norm = ops.ReduceSum(l2_norm)
    l2_norm = ops.Sqrt(l2_norm)
    
    loss = ops.Add(loss, ops.Mul(0.01, l2_norm))
    
    return loss

第二段代码看起来更长，但在NPU上跑实际更快——因为避免了Python层面的多次临时张量创建。

一些踩坑经验

数值精度问题：float16 上做 Log 和 Exp，注意输入范围。如果不确定，先转 float32 算，再转回 float16。
内存对齐：NPU上的向量计算喜欢对齐的地址。如果自己管理内存，注意用 acl.rt.malloc 而不是普通的 malloc。
随机数seed：如果训练过程中用了NPU生成的随机数（比如Dropout），记得设置固定的seed，否则每次跑结果不一样，不好调试。
算子融合机会：Softmax = Exp + ReduceSum + Div，这三个算子可以跟前后的MatMul融合。用ATB或者graph-autofusion可以自动做这个事。