写 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 → sum - Norm 计算 ：LayerNorm 的
mean是用 ReduceSum 除以元素数算出来的 - Attention Score 的 Softmax ：分母部分做
exp(x).sum(axis=-1) - 梯度裁剪 ：
grad.norm()需要对梯度 Tensor 做平方和归约
Transformer 的每次前向推理中 ReduceSum 被调用几十次。它不在计算量上占大头，但它的 Memory 访问模式决定了它很难被优化到跟 GEMM 一样的效率。
Tensor 归约为什么会慢
ReduceSum 的问题是：计算量极小，数据搬运量极大。
对一个 [4096, 4096] 的 float16 Tensor 做 axis=0 的归约求和一个标量：
- 数据读取量：32MB（整个 Tensor）
- 计算量：16M 次加法
- 计算/搬运比：约 0.5 FLOPs/byte
对比 GEMM（52.5 FLOPs/byte），ReduceSum 的计算/搬运比差了 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：
cpp
// 多 Core ReduceSum（简化）
// 每个 Core 处理自己分到的数据块
float partial_sum = 0;
for (int i = core_id * block_size; i < (core_id+1) * block_size; i += 128) {
float16 vec[128] = Load(GM, i);
partial_sum += Sum(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 次 ReduceSum：
- 40 个 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(axis=0) 跨行归约时的数据访问模式是跨步的------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%。