写 AI 训练代码时经常看到一个操作loss loss.sum()或者accuracy (pred label).sum()/len(label)。一个简单的归约求和CPU 上几微秒的事。在 NPU 上求和这个操作不简单——Tensor 分布在 DDR 的不同位置上把所有元素累加成一个标量需要精心设计的数据搬运策略。CANN 的 ops-math 仓库管理着所有数学类基础算子——归约、逐元素数学运算、类型转换。ReduceSum 是其中最常用的算子之一。ReduceSum 为什么常见AI 训练和推理中 ReduceSum 的出现频率极高损失计算cross_entropy_loss内部做log_softmax → nll_loss → sumNorm 计算LayerNorm 的mean是用 ReduceSum 除以元素数算出来的Attention Score 的 Softmax分母部分做exp(x).sum(axis-1)梯度裁剪grad.norm()需要对梯度 Tensor 做平方和归约Transformer 的每次前向推理中 ReduceSum 被调用几十次。它不在计算量上占大头但它的 Memory 访问模式决定了它很难被优化到跟 GEMM 一样的效率。Tensor 归约为什么会慢ReduceSum 的问题是计算量极小数据搬运量极大。对一个[4096, 4096]的 float16 Tensor 做 axis0 的归约求和一个标量数据读取量32MB整个 Tensor计算量16M 次加法计算/搬运比约 0.5 FLOPs/byte对比 GEMM52.5 FLOPs/byteReduceSum 的计算/搬运比差了 100 倍。这意味着 ReduceSum 的执行时间几乎 100% 花在数据搬运上——Vector Unit 算加法只需要几微秒但 32MB 数据从 DDR 搬到片上需要几百微秒。昇腾NPU如何做并行归约ReduceSum 不跑 Cube Unit矩阵乘用不到它在 Vector Unit 上执行。朴素的 ReduceSum 就是单线程扫描——从 DDR 读 128 个元素向量加法再读 128 个循环到读完。但对大 Tensor 来说单线程太慢。ops-math 用 Parallel Reduction 来加速。把 Tensor 分成 N 块每块分配给不同的 AI Core// 多 Core ReduceSum简化// 每个 Core 处理自己分到的数据块floatpartial_sum0;for(inticore_id*block_size;i(core_id1)*block_size;i128){float16 vec[128]Load(GM,i);partial_sumSum(vec);// Vector 指令做 128 元素求和}// 所有 Core 的 partial_sum 合并成最终结果AllReduce(partial_sum,final_sum);4 个 Core 一起做搬运量不变每个 Core 读自己的分片但计算时间降到 1/4。ops-math 的优化思路ops-math 内部针对不同归约场景做了专门的优化小 Tensor 归约元素数 4096。直接在单 Core 的 Vector Unit 上做不需要跨 Core 同步。跨 Core 同步的开销几十微秒比计算时间还长。中 Tensor 归约4096 元素数 1M。多 Core Parallel Reduction每个 Core 读一个分片算部分和最后 AllReduce 合并。大 Tensor 归约元素数 1M。多 Core Parallel Reduction 归约树合并。不用 AllReduce开销太大而是用树形归约——Core 0 和 Core 1 先合并Core 2 和 Core 3 先合并两层之后再合并。归约的精度也需要注意。FP16 的 ReduceSum 在元素数多时可能因累加顺序不同产生精度偏差。ops-math 支持 FP32 累加器——部分和在 Vector Unit 上用 FP32 累加最后再转回 FP16。大模型中的归约场景LLaMA-13B 推理的 Prefill 阶段一次前向计算涉及约 40 次 ReduceSum40 个 Decoder Block × 各 1 次 LayerNorm mean var 计算Attention Softmax 的分母求和每步一次Loss 计算训练场景单次 ReduceSum 的时间约 15-35μs。40 次合计约 1ms——占 Prefill 总延迟的 3-5%。单独看不大但 40 次独立 Kernel Launch 的调度开销累计后约 0.4ms占了近一半的归约时间。ops-math 的优化方向之一是批量归约——把多个独立的 ReduceSum 合并成一个 Kernel一次 Launch 算完所有归约。ReduceSum 的并行策略选择ops-math 针对不同情况选择不同的并行归约策略小数据量1024 元素单 Core 串行。并行带来的同步开销超过收益中数据量1024-1M 元素多 Core 并行每个 Core 算部分和最后用 Vector Unit 的 reduce 指令合并大数据量1M 元素多 Core 归约树用树形结构分批合并axis参数的选择也影响归约效率。sum(axis0)跨行归约时的数据访问模式是跨步的——DMA 无法连续读取带宽利用率下降。sum(axis-1)归约最后一维时数据连续带宽利用率高。ops-math 会在编译时检查 axis 参数如果 axis 导致跨步访问会先做一次 Transpose 把归约轴换到连续位置再执行归约。大模型中的归约场景LLaMA-70B 训练时ReduceSum 的使用场景包括Loss 函数计算batch 内的 loss 求和、梯度裁剪前的 norm 计算梯度 Tensor 的平方和、LayerNorm 的 mean 计算。每个训练 step 这些归约操作的总时间约 200-300μs——跟反向传播的几十毫秒相比可以忽略。但如果在推理场景中频繁调用独立的归约 Kernel调度开销可能占总延迟的 5-10%。参考仓库ops-math 数学算子库ops-blas 线性代数算子库