ops-reduce：ReduceMax 与 ReduceMean 的并行优化

Transformer 推理中 Reduce 类算子无处不在：LayerNorm 需要算 mean 和 var（ReduceMean + 方差），Softmax 需要算 max（ReduceMax）和 sum（ReduceSum），Attention 的归一化因子需要跨序列维度做归约。

ops-reduce 是 CANN 管理所有归约算子的仓库------ReduceMax、ReduceMin、ReduceMean、ReduceProd。它们共享同一个核心问题：计算量极小但搬运量极大。

为什么 Reduce 类算子容易成为瓶颈

以 ReduceMax 为例------在一个 [4096, 4096] 的 Tensor 上沿行方向找最大值：

数据量：32MB
计算量：4096×4095 = 16M 次比较
计算/搬运比：约 0.5 FLOPs/byte

对比 GEMM 的 52.5 FLOPs/byte------Reduce 的瓶颈非常明显：时间主要花在从 DDR 读数据上，Vector Unit 的比较运算只需要几微秒。

ReduceMax 与 ReduceMean 的典型场景

ReduceMax------Attention 中的 Softmax。 Softmax 的第一步需要对 Score 矩阵的每一行找最大值，用于数值稳定性（减去最大值防止 exp 溢出）。这个 ReduceMax 在朴素 Attention 中需要对 32MB 的 Score 矩阵做一次完整扫描。

ReduceMean------LayerNorm。 LayerNorm 的第一步 mean = sum(x) / n 需要对输入 Tensor 做 ReduceSum 再除以 n。LLaMA 的每个 Decoder Block 做两次 LayerNorm------每次都要对隐藏层的 [B, n, d] Tensor 做 mean 和 var。

昇腾NPU如何做并行归约

ops-reduce 在 Vector Unit 上做并行归约。一个 Core 处理一部分数据，用 SIMD 指令同时比较多个元素：

cpp 复制代码

// 昇腾 Vector Unit 上的 ReduceMax Kernel（伪代码）
// 输入：x (GM 地址), n = 元素数
// 输出：max_value (GM 地址)

__vector__ void reduce_max_kernel(...) {
    float16 local_max = -FLT_MAX;
    
    // 每次处理 128 个元素
    for (int i = 0; i < n; i += 128) {
        float16 vec[128] = load_gm_to_local(x + i); // DDR→L1
        local_max = vec_max(local_max, vec);         // SIMD 比较
    }
    
    // 如果启动了多个 Core，做跨 Core 归约
    float16 global_max = warp_reduce_max(local_max);
    
    store_local_to_gm(global_max, max_value);        // L1→DDR
}

多 Core 场景中每个 Core 先算自己的局部最大值，然后通过一次跨 Core 的 reduce 操作合并为全局最大值。这个跨 Core reduce 在 Vector Unit 上用 warp_reduce_max 指令完成------比写 DDR 再读回来快两个数量级。

性能分析

在 Ascend 910 上对不同大小 Tensor 做 ReduceMax 的实测：

Tensor 形状	元素数	搬运量	延迟	带宽利用率
[1024]	1K	2KB	3μs	15%
[1,4096]	4K	8KB	5μs	28%
[4096,4096]	16M	32MB	85μs	72%
[8,4096,4096]	128M	256MB	610μs	78%

小 Tensor 的带宽利用率低是因为 DMA 启动开销占比大。大 Tensor 的利用率接近硬件上限。

Transformer 中的归约链路

LLaMA-7B 的两次 LayerNorm 的归约链路：

复制代码

输入 [B, n, d] → ReduceMean → (x - mean)² → ReduceMean → var
                       ↑                          ↑
                  一次归约                     一次归约
                  mean 缓存                    var 缓存

ops-reduce 在 LayerNorm 场景中的优化是：mean 和 var 的归约 Kernel 共享同一个 Tensor 的搬运------mean 归约时从 DDR 读一次 x，var 归约时复用 L1 上的 x 数据，不需要重新搬运。优化后归约的搬运量从 2 次减到 1 次。

ReduceMean 的融合执行

ReduceMean 可以跟下游算子融合。LayerNorm 的 (x - mean) / sqrt(var + eps) 中，ReduceMean 和后面的计算共享同一个 x 的搬运------x 搬入 L1 后先做 ReduceMean，mean 的结果留存在 L1 中，后续的 x - mean 和 ReduceVar 直接使用------x 不需要重新搬运。

融合后 LayerNorm 的总搬运量从 2 次降到 1 次。实测中 LLaMA-7B 的 LayerNorm 延迟从 12μs 降到 7μs。

性能分析表的补充

Ascend 910 上不同 Tensor 大小的 ReduceMean（axis=-1）性能：

Tensor 形状	元素数	延迟	瓶颈环节
[1, 4096]	4K	4μs	DMA 启动
[4, 4096, 4096]	64M	210μs	DDR 带宽
[8, 4096, 4096]	128M	395μs	DDR 带宽
[1, 4096, 4096]	16M	56μs	DDR 带宽

Batch 维度增加时延迟近线性增长------ops-reduce 在 Batch 维度上可以并行，每个 Batch 分给不同 Core 处理，Batch=8 时 4 个 Core 并行能把延迟从 395μs 降到 110μs。

参考仓库

ops-reduce 归约算子库

ops-math 数学算子库