大模型推理的性能瓶颈通常被归结为 Attention 和 FFN，但实际上 ReduceMean、Softmax、Cast 这些 ops-math 算子如果走了非优化路径，能吃掉 10-15% 的推理时间。这篇从 Llama2-7B 的推理链路出发，逐步检查每个 ops-math 算子的性能状态。

热点分析

用 CANN 的 Profiler 工具采集 Llama2-7B 的推理 profile：

export PROFILING_MODE=1
python infer_llama.py

# 分析结果
python -m msquickstart.profile.analyze --profiling-dir ./profiling_data

典型结果中 ops-math 算子的时间占比：

加起来 8.9%。融合后可以降到 2-3%。

Softmax：确认走的是优化路径

标准 Softmax 需要 3 次 HBM 读取。ops-math 的优化版本把减最大值、指数、归一化压进一次 HBM 读取：

import torch_npu

# 优化版本：单次 HBM 读取
probs = torch.nn.functional.softmax(logits, dim=-1)
# 在昇腾NPU上 torch_npu 自动走 ops-math 优化路径

判断是否走优化路径的方法：用 Profiler 看 Softmax kernel 的执行时间。seq_len=4096 时优化版本约 0.03ms，非优化版本约 0.08ms。如果你的数据超过 0.06ms，检查输入 tensor 是否 .contiguous()。

ReduceMean：融合到 LayerNorm

ReduceMean 在 RMSNorm 里计算 mean(x²) 作为分母。CANN 8.5 的 fused_add_rms_norm 把 ReduceMean 和后续的 Sqrt、Div 融合了：

# 融合前：ReduceMean + Sqrt + Div = 3 个 kernel
normed = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + eps)

# 融合后：1 个 kernel
normed = torch_npu.npu.rms_norm(x, weight, eps)

32 层 × 2 次 = 64 个 ReduceMean kernel 被消除，省约 1.5ms。

Cast：消除不必要的转换

训练框架经常做多余的 Cast。bf16 权重加载后转 fp16，推理时又转回来：

# 常见错误：加载时已经是 fp16，又转了一次
model = AutoModel.from_pretrained("model_id", torch_dtype=torch.float16)
model = model.half()  # 多余！已经是 fp16

# 正确写法
model = AutoModel.from_pretrained("model_id", torch_dtype=torch.float16)

每次多余的 Cast 约 0.02ms，32 层里如果每层多一次就是 0.64ms。

端到端对比

Llama2-7B 推理，ops-math 优化前后：

decode 延迟降 18%。收益来自 Softmax 走优化实现、ReduceMean 被融合、消除多余 Cast。

数学算子优化不是改一行代码的事，需要用 Profiler 定位热点，逐个确认是否走了融合和优化路径。仓库在这里：

https://atomgit.com/cann/ops-math