Reduce 算子优化【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills适用于需要聚合多个值的归约操作适用算子基础归约: sum, mean, max, min, prod归一化: softmax, logsoftmax, layernorm, batchnorm统计: variance, std通用归约策略1. 块内归约 原子操作triton.jit def reduction_kernel(input_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): pid tl.program_id(0) offsets pid * BLOCK_SIZE tl.arange(0, BLOCK_SIZE) mask offsets n_elements # 加载数据 data tl.load(input_ptr offsets, maskmask, other0.0) # 块内归约 block_sum tl.sum(data, axis0) # 原子操作写回全局内存 tl.atomic_add(output_ptr, block_sum)2. 减少规约精度损失关键: 如果需要在 FP16 或 BF16 的数据上执行计算性规约除了max, min的规约计算应在规约计算前将其强制转换为 FP32以避免低精度累加带来的数值误差。# 错误直接用 fp16/bf16 累加精度损失大 data tl.load(input_ptr offsets, maskmask, other0.0) # data 为 fp16/bf16 block_sum tl.sum(data, axis0) # 低精度累加 carry carry block_sum # 低精度累加 # 正确在执行累加计算前转为 fp32在 fp32 上完成规约 data tl.load(input_ptr offsets, maskmask, other0.0) data data.to(tl.float32) # 强制提升为 fp32 block_sum tl.sum(data, axis0) # 高精度累加 carry carry block_sum # 高精度累加 # 如果输出要求 fp16/bf16在最终 store 前转回 tl.store(output_ptr, block_sum.to(input_ptr.dtype.element_ty))原则在执行规约操作前.to(tl.float32)如果涉及多次规约累积多次规约结果的累加器对象精度应为tl.float32涉及计算的规约操作除了max, min的规约操作均在 FP32 上执行在最后tl.store前按需转回原始数据类型3. 数值稳定性处理关键: 对于涉及 exp 的操作softmax、logsoftmax必须减去最大值防止溢出。# 错误错误直接 exp 可能溢出 scores tl.math.exp2(x) # 正确正确减去最大值 max_val tl.max(x, axis0) scores tl.math.exp2(x - max_val)【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考