当前位置：首页 > news >正文

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

news 2026/6/11 0:21:46

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 指令同时比较多个元素：

// 昇腾 Vector Unit 上的 ReduceMax Kernel（伪代码）// 输入：x (GM 地址), n = 元素数// 输出：max_value (GM 地址)__vector__voidreduce_max_kernel(...){float16 local_max=-FLT_MAX;// 每次处理 128 个元素for(inti=0;i<n;i+=128){float16 vec[128]=load_gm_to_local(x+i);// DDR→L1local_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。