已关闭
SwigluGroupGrad算子性能优化 #4502
小王!创建于  8月3日关闭于  8月11日
小王!
小王!成员
8月3日 创建

SIMD 模板的性能优化:

1、针对网络中的 BF16、H=2048、存在 weight 和 y_origin、无 group_index 重点场景,围绕固定长度归约、对齐搬运、输出转换和编译期分支裁剪增加多项快路径。其他 shape 仍通过通用路径处理,保证原有功能覆盖。

2、针对 grad_weight 的 NumPy pairwise 求和开销,新增固定长度 2048 的专用归约。将数据固定拆分为 16 个 128 元素叶子,并按照与通用实现一致的二叉树顺序合并,去除递归栈、标量循环和状态判断开销。

3、针对每个 128 元素叶子的八路累加,使用 SIMD 同时处理 8 个叶子,16 个叶子分两组完成。各 lane 仍按照原有顺序累加,并通过固定的 PairReduceElem 完成 8→4→2→1 合并,保持原有 FP32 加法结合顺序。

4、针对 BF16、FP16 路径中间结果多次落 UB 的问题,将 FP32 到输出类型的转换合并到主 SIMD 计算流程。梯度在寄存器中完成计算和转换后直接写入输出 buffer,减少 FP32 中间结果的一次 UB 写入和重新读取。

5、针对 HA_ == H_ 的完全对齐场景,将输入 x 的 gate、up 两部分由两次搬运合并为一次连续 2H 搬运,输出 grad_x 同样合并搬出,减少 DMA 指令数量和逐行搬运控制开销。H=2048 场景可直接命中该优化。

6、针对热路径中的无效操作进行编译期裁剪。无 group_index 时不再执行梯度末尾的乘 1 操作;存在 y_origin 时,将 grad_weight 所需的 grad_y * y_origin 提前计算和写入,使其能够与后续 Sigmoid 的长依赖计算更好地交叠。

7、合并 Cast 阶段的多个流水同步,并精简 SIMD 计算、归约和输出转换之间不必要的 PIPE_ALL。仅在存在真实数据依赖的位置保留 Barrier,减少流水串行等待,有利于 Vector 计算、UB 访问和数据搬运之间的重叠。

8、H 不等于 2048 时,grad_weight 继续使用原 NumPy pairwise 通用归约;大 H 场景仍保留 chunk 分支。通用路径继续支持动态 H、非对齐尺寸、clamp、group_index、FP16、FP32及特殊值场景。

likedislike
小王!小王!成员
8月10日 修改了issue 的描述
CANN-robotCANN-robot成员
8月11日 关闭了 issue
CANN-robotCANN-robot成员
8月11日 添加了label:resolved