写 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：

// 多 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%。

参考仓库

ops-math 数学算子库

ops-blas 线性代数算子库