已合并
优化SwigluGroupGrad性能 #8544
优化SwigluGroupGrad性能 #8544
已合并
小王!创建于 8月11日
小王!
小王!成员
8月11日

描述

优化方案如下:
一、SIMT 模板性能优化:
回退并保留库上稳定 SIMT 基线,仅新增严格互斥的 H=1 快路径;H=1 时直接完成单元素反向与可选权重输出,去除通用 hidden 循环、pairwise 判断和多层可选分支,避免中小 H 场景回归。
二、SIMD small kernel 模板性能优化:
路由层将 T≤8、H≥262144 且同时存在 weight、y_origin 的超宽场景转入 fixed-tree 专用模式;任务域从“行”扩展为“行 × hidden 子树”,按 hidden 切分填满多核,并保持 gradWeight 的 FP32 pairwise 归约顺序。跨核 partial 通过 DMA 传递,避免缓存行伪共享;同时统一 blockDim 与任务契约,确保实际开核数与 Kernel 分工一致。

关联的Issue

https://gitcode.com/cann/ops-nn/issues/4832

测试

本地批跑500条性能用例均无功能性问题,优化后的性能单case均达标,二级冒烟和obp冒烟均无异常

文档更新

不涉及,仅性能优化修改了代码

类型标签

AI/Agent生成声明

likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 小王! 的贡献)
小王!小王!成员
8月11日 创建了 pull request,commit 93a5450f
atomgit-bot
atomgit-bot
8月11日 评论:

变更摘要

本次 PR 针对 SwigluGroupGrad 算子的 RegBase 和 SIMT 内核进行了多项性能优化。主要围绕 Tiling 策略重构:引入隐藏维度多核拆分(forceHiddenCoreSplit)、优化 UB 缓冲区分配(yOrigin/weight 双缓冲以支持 DMA 预取与批处理)、移除不再需要的 dwAccum 中间缓冲区,以及重新定义 blockFactor 语义以支持「行 × 隐藏块」二维任务划分。同时,SIMT 路径移除了此前为匹配 NumPy 逐对求和语义而实现的复杂递归归约树,改为在主循环内联执行普通 FP32 累加,并新增了 splitHidden + weight 场景下的 workspace 分配与调度模式设置。

主要改动

  • Tiling 策略新增隐藏维度多核拆分: 在 CalcTilingStrategy() 中新增 forceHiddenCoreSplit 逻辑,当行数不足以填满所有 AIV 核心且 H_ > 2048 时,将隐藏轴切分为多个 chunk,形成「行 × 隐藏块」的二维任务网格,使更多核心参与计算。
  • UB 缓冲区布局重构以支持 DMA 双缓冲: yOriginBytes 从单缓冲改为双缓冲(2 * rows * HA * dt),weightBytes 改为 2 倍 weightSingleBufferBytes,dwOutBytes 改为单缓冲;同时移除了 splitHidden 路径下不再需要的 dwAccumBytes,释放 UB 空间。
  • blockFactor 语义变更与调度核心数计算: blockFactor 不再存储 rFB(每核行数),而是直接存储 usedCoreNum;PostTiling() 直接使用该值设置 BlockDim,并在 splitHidden + weight 场景下通过 SetScheduleMode(1) 启用多核间规约调度。
  • 新增 gradWeight 部分结果 workspace: 在 GetWorkspaceSize() 中,当 splitHidden == 1 且 isWeight_ == 1 时,分配大小为 totalRows_ * numChunksPerRow * FP32_BYTES 的 workspace,用于存放各核心的 gradWeight 部分和。
  • SIMT 内核 gradWeight 归约简化: 移除 SimtDwPairwiseSum / SimtDwPairwiseLeaf 等完整的 NumPy 逐对求和树实现及其辅助函数(SimtFp32Mul、SimtFp32Add、SimtDwProductAt),改为在 H 循环内通过 SimtDwAccum 执行普通 FP32 顺序累加,简化代码并消除递归树带来的性能开销。
likedislike
不准确?
atomgit-bot
atomgit-bot
8月11日 评论:

代码审查

审查总结

审查了全部 3 个变更文件:

文件 审查结果
quant/swiglu_group_grad/op_host/arch35/swiglu_group_grad_regbase_tiling.cpp 发现 1 个 P2 问题:GetWorkspaceSize 在 tiling 决策前读取 splitHidden 字段,导致 hidden-split + weight 场景下 workspace 仅分配 0 字节,kernel 将越界写入 GM
quant/swiglu_group_grad/op_kernel/arch35/swiglu_group_grad_regbase.h 未发现问题。DMA 双缓冲、UB 分配、同步(SyncAll)、workspace 偏量和 reduction 逻辑均与 tiling host 侧一致,kernel 内 overflow / 越界检查充分
quant/swiglu_group_grad/op_kernel/arch35/swiglu_group_grad_simt.h 发现 1 个 P3 问题:Copyright 年份 2026→2025 疑似误改;gradWeight 计算从 golden-compatible pairwise sum 改为顺序累加是刻意的性能权衡,非 bug

按严重度统计:P2 × 1,P3 × 1

整体风险评估:中等风险。P2 的 workspace 分配不足问题在 hidden-split + weight 场景下会导致确定性的 GM 越界写入,可能引发数据损坏或设备异常,建议优先修复。P3 为低风险合规问题。其余大量变更(DMA 批量预取、barrier 消除、BF16 fast path、H=2048 SIMD reduction 路径)逻辑自洽,未发现正确性问题。

类型 数量
🔴 阻塞 0
🟡 建议 1

💬 仅评论

likedislike
不准确?
CANN-robotCANN-robot成员
8月11日 添加了label:cann-cla/yes
CANN-robot
CANN-robot成员
8月11日 评论:

Thanks for your pull-request.
The full list of commands accepted by me can be found at here.
You can get sig-info at here.
For more, you also can visit HICANN.


PR Approval Progress

✅ Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

module lgtm status approve status
quant ✅ 唐玮玮, 陈风 (2/2) ✅ 陈风 (1/1)

💡 Tip:

  • Committer can comment /approve or /lgtm
  • Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

qq_52056150, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
此处折叠了180条消息 查看更多
tangweiwei2成员
8月18日 评论:

/lgtm
/approve

likedislike
CANN-robotCANN-robot成员
8月18日 添加了label:lgtm
CANN-robotCANN-robot成员
8月18日 关闭了关联的issue
CANN-robotCANN-robot成员
8月18日 合入了pull request
CANN-robot
CANN-robot成员
8月18日 评论:

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.

likedislike