Merged
perf(fused_cross_entropy_loss_with_max_sum): 小V场景性能优化(平坦打包快路+预取流水+裸vdiv) #14270
AlfengYuanAdd on 3 days ago
perf(fused_cross_entropy_loss_with_max_sum): 小V场景性能优化(平坦打包快路+预取流水+裸vdiv) #14270
Merged
AlfengYuanAdd on 3 days ago
AlfengYuan
AlfengYuanCollaborators
3 days ago

关联 issue

#8653 aclnnFusedCrossEntropyLossWithMaxSum 小V场景性能问题

根因

大 V 通路按行广播(每行一次 DIST_BRC_B32),小 V(V∈{1,2,3,4,7};不达标用例中 200/299 条 V=1)时每行仅 1~8 有效 lane,广播与 2-D 小搬运的固定开销占满流水,有效带宽退化约两个数量级。issue 中 49 条单 case 性能不达标用例全部为小 V 场景。

修改内容(4 个提交)

  1. 小V多行平坦打包快路:vocab 整块 1-D 搬入、b16 预 cast fp32、softmax G 行打包 G*v lane、Gather(l/V) 展开 max/invSum、V=1 免 Gather;tiling 新增 aPerLoop。
  2. 预取流水:原实现 tile 内 DeQue 串起"搬入→计算→搬出"全串行(MTE/VEC 零重叠);改为 prologue 预取 + 循环内"先发本块 VEC、再预取下一块、最后等输出",max/sum/pred/loss 四队列深 2;aPerLoop 6912→4736(每行标量预算 20B→36B)。提速 1.12~1.22x。
  3. 裸 vdiv + b16 内联 cast:①inv_sum 的 PRECISION_1ULP_FTZ_FALSE 补偿除改裸 vdiv(次正规 sum 时 inv_sum 本就 inf,补偿白做);②b16 vocab 免整 tile 预转,加载时 DIST_UNPACK_B16 内联 cast(bit 级等价;V∈{5,7} 组内偏移 32B 对齐不满足,保留预转)。再提速 1.04~1.99x(中位 1.39x)。
  4. 评审修复:小V预取流水用的 4 个标量队列声明深度 1→2 对齐 TQue 契约;vocab 源缓冲(vocabQue_ 双槽 + vocabFp32Buf_)各加 1 个 VL padding,收敛 V=3/5/6/7(Gv<VL)满 tile 末组向量满载读的越界(越界 lane 本就掩码、不影响结果;padding 开销由 host 预算既有 2V_ALIGN*(dtypeSize+4B) 余量覆盖,aPerLoop 不变)。另有 1 个纯格式化提交(pre-commit clang-format)。

验证

  • 49 条原不达标用例:精度+性能 49/49 PASS。有效带宽 730~835 →(预取流水)→ 1164~1309 GB/s,追平 torch copy 实测上限(1230 GB/s);相对原始基线累计约 25~45x。
  • 全量 299 回归:性能 299/299、精度 298/299。唯一 FAIL(L1_296)为偶发边界失配,与本次修改无关:同种子(--seed 12345)同输入下 baseline 与本分支两个输出 bit 级完全一致(V=3789 不走小V快路);不固定种子时该用例输入每次随机,rmse 比值在阈值 2.0 两侧波动(复跑 9 次 7 次 PASS)。机制:该用例 logits_max∈(-1000,-10) 与 vocab 恒 6.1e-5 独立随机构造、互不自洽,softmax 输出踩 fp32 表示上限,rmse 为绝对误差被最大元素单点主导,NPU 向量 exp 与 torch CPU exp 在溢出边缘的 ulp 级差异之比随数据抽取波动。
  • ophost UT:8/8 PASS(含新增 V=3/fp32 tiling 边界用例,aPerLoop=2560)。
  • 评审修复回归:49 条用例精度+性能 49/49(性能与修复前持平,中位 1.00x);新增 V=3/V=5 fp32(DIST_NORM 路径)与 V=6 fp16(内联 cast 路径)满 tile 边界用例 3/3 PASS。pre-commit 全绿。
  • 180 条 vocab=None 用例依赖 !14169 的 golden 修复(已合入,本分支基于其上)。

说明(一个被否掉的方案)

曾评估单阶段 softmax = exp(x-max-logSum)(可再快 ~40%),但 cross_check 特殊位要求与 golden(fp64 promote)/third_party(fp32) 之一精确一致,两者均以 exp(x-max) 中间量溢出定义 inf;单阶段在 x-max∈(88.72, 88.72+logSum) 窗口产出有限值,导致 16/49 精度 FAIL(dump 三方对照实测失配点全在该窗口)。故弃用,保留两阶段 exp(x-max)·invSum 结构,per-element Sub+Exp+Mul 已到底。

🤖 Generated with Claude Code

likedislike
Pull Request has been successfully merged by CANN-robot.
(Thank you for AlfengYuan's contribution.)
AlfengYuanAlfengYuanCollaborators
3 days ago Created pull request, commit da51cbcb
AlfengYuanAlfengYuanCollaborators
3 days ago Associated issue: [Question|问题咨询]: 项目构建耗时较长
atomgit-bot
atomgit-bot
3 days ago Comments:

变更摘要

本 PR 针对 aclnnFusedCrossEntropyLossWithMaxSum 在小 vocab(V≤8)场景下按行广播导致带宽退化的问题,在 regbase 通路中新增“小 V 平坦打包”快路:将多行 vocab/softmax 打包为连续 1-D 向量搬运与计算,并按行数而非行内 lane 重排 UB 预算。同时在 kernel 侧引入跨 tile 预取流水(MTE/VEC 重叠)、inv_sum 使用裸 vdiv、b16 vocab 加载时内联 cast,并同步更新 aPerLoop 分块字段、tiling 逻辑与 UT 期望。

主要改动

  • 新增 tiling 字段 aPerLoop 与小V预算计算:FusedCrossEntropyLossWithMaxSumRegBaseTilingData 增加 aPerLoop(大V通路恒为 FUSED_CE_MAX_SUM_A_PER_LOOP,省显存通路置 0);DoTilingFull 在 v <= FUSED_CE_MAX_SUM_SMALL_V_MAX 时按每行字节开销(vocab/softmax 双缓冲、fp32 副本、9 个行向量 buffer)与 ubSize_ 反算并对齐 128 的行数上限 8192,UB 不足时报错返回失败。
  • 新增小V判定常量:fused_cross_entropy_loss_with_max_sum_tiling_data.h 引入 FUSED_CE_MAX_SUM_SMALL_V_MAX = 8,作为“多行打包 + Gather 展开 max/inv”行主序平坦通路的启用阈值。
  • 小V主流程与预取流水:kernel 新增 smallV_ 标志及 ProcessFullSmallV/PrefetchSmallVTile/SmallVTileRows,按“先下发本块 VEC 计算、再预取下一块输入、最后等本块输出”的排布实现 prologue 预取与跨 tile MTE/VEC 重叠;对应 buffer 初始化改为按 VL_FP32 对齐的 tile 与深 2 的 logitsMaxQue_/sumExpScalarQue_/predictedScalarQue_/lossScalarQue_ 行向量队列,并新增 vocabFp32Buf_。
  • 扁平 softmax 与裸 vdiv 计算:新增 ComputeSoftmaxFlat/SoftmaxFlatLoop,按 G 行打包成 G*v 条 lane,V=1 免 Gather 直接对位,V>1 通过寄存器内构建 idx[l]=floor(l/v)(Arange→Cast→Div→CAST_FLOOR)驱动 Gather 展开 max/inv;ComputeLossInvSumLoop 中 inv_sum 改用新增的 castDivIntrinsic(INTRINSIC)裸 vdiv。
  • b16 内联 cast 与加载模板化:LoadVocabToFp32 改为按源类型 S 模板化,新增 CastVocabToFp32;b16 vocab 在组内偏移满足 32B 对齐时用 DIST_UNPACK_B16 加载即内联 cast,仅 V∈{5,7} 等不满足对齐的场景回退整 tile 预转 fp32。
  • UT 补充:test_fused_cross_entropy_loss_with_max_sum_arch3510_tiling.cpp 新增 fp16 小V(v=1)用例,校验 vPerLoop、vChunk、aPerLoop=4736 等 tiling 输出。

🧩 本次 Review 已使用自定义 Skill:

likedislike
Incorrect?
atomgit-bot
atomgit-bot
3 days ago Comments:

代码审查

审查结论

各变更文件审查结果

文件 结论
op_host/fused_cross_entropy_loss_with_max_sum_tiling_arch35.cpp no issues(DoTilingFull 小V预算与 kernel 侧 9 行向量 / 双缓冲 / fp32 副本的分配逐项核对一致;aPerLoop 追加在结构体末尾、DoTilingForMemory 置 0、int64 乘法链无 32 位溢出、bytesPerRow>=52 无除零;smallV 判定与 kernel 门控等价;V≤8 时 SplitV 必然 vCores=1,故小V通路不写重复 loss)
op_kernel/arch35/fused_cross_entropy_loss_with_max_sum_ar.h 4 处(1×P2、3×P3)
op_kernel/arch35/fused_cross_entropy_loss_with_max_sum_tiling_data.h no issues(新常量/字段追加在末尾,未改动历史字段顺序;V_ALIGN=64 与 kernel VL_FP32 假设一致)
tests/ut/op_host/test_fused_cross_entropy_loss_with_max_sum_arch3510_tiling.cpp 1 处(P3,注释推导错误)

问题统计

  • P0/P1:0
  • P2:1(深度 1 的 TQue 传入 2 个 buffer)
  • P3:4(裸 vdiv 的次正规契约分歧、热循环内 Duplicate 可外提、g 对齐判定重复计算、UT 注释预算算错)

整体风险判断

小V平坦快路的整体设计(1-D 整块搬入 + Gather(l/V) 展开 max/inv + 跨 tile 预取 + b16 内联 cast 的对齐判定)经逐条推演在数值与访存边界上是自洽的:组内字节偏移、rowPad/tileElems 的 VL 对齐、尾块掩码、host↔kernel UB 预算三方一致,未发现越界写或漏处理尾块。主要遗留风险是 4 个深度 1 的队列被按深度 2 使用(可能触发 Ascend C 队列深度断言/状态越界,量级取决于框架实现),以及小V通路把补偿除换成裸 vdiv 造成的两条通路数值契约分歧(需次正规 sum_exp 才能触发)。建议合入前先确认 TQue 深度约束并至少补一次 v∈{5,7}(预转路径)与 V=1/大V 的输出比对。

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

💬 仅评论


🧩 本次 Review 已使用自定义 Skill:

likedislike
Incorrect?
CANN-robotCANN-robotCollaborators
3 days ago Added label: stat/needs-squash
CANN-robotCANN-robotCollaborators
3 days ago Added label: cann-cla/yes
CANN-robot
CANN-robotCollaborators
3 days ago Comments:

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.
You can self-configure the PR merge rules for this repository. For more details, please refer to 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
loss ✅ 汤平川, 陈琦 (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

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

likedislike
CANN-robotCANN-robotCollaborators
3 days ago Set zengjuan,u010470851,chenfeng61,chaotang233,crystalhu,yu-xinjie62,yangyang016,fanqirui,su-yueming,yourealize,chenqi317,wang-xing001,liubo75,tangweiwei2,east_yang,pingchuantang,qianzehong as approver
CANN-robotCANN-robotCollaborators
3 days ago Set u010470851,chenfeng61,yu-xinjie62,su-yueming,yourealize,wang-xing001,east_yang,pingchuantang,qianzehong as approver
AlfengYuanAlfengYuanCollaborators
3 days ago Deleted associated issue: [Question|问题咨询]: 项目构建耗时较长
atomgit-bot
atomgit-bot3 days agoStart code review1
loss/fused_cross_entropy_loss_with_max_sum/op_kernel/arch35/fused_cross_entropy_loss_with_max_sum_ar.h
@@ -152,0 +194,4 @@
194+ pipe_->InitBuffer(logitsMaxQue_, 2, rowPad * sizeof(float));
195+ pipe_->InitBuffer(sumExpScalarQue_, 2, rowPad * sizeof(float));
196+ pipe_->InitBuffer(predictedScalarQue_, 2, rowPad * sizeof(float));
197+ pipe_->InitBuffer(lossScalarQue_, 2, rowPad * sizeof(float));
atomgit-bot
atomgit-bot3 days agoComments:

🟡 Medium Priority

变更行:fused_cross_entropy_loss_with_max_sum_ar.h:194-197 在小V快路里对 logitsMaxQue_ / sumExpScalarQue_ / predictedScalarQue_ / lossScalarQue_ 调用 pipe_->InitBuffer(que, 2, rowPad * sizeof(float)),但这 4 个队列在 142-145 行声明为 TQue<QuePosition::VECIN, 1>(lossScalarQue_ 为 TQue<QuePosition::VECOUT, 1>),即模板深度参数 bufNum=1。

契约/影响:Ascend C 的 TQue 第二个模板参数是队列深度(可缓存的 Tensor 个数),TPipe::InitBuffer(TQue&, num, len) 的 num 语义上是 [1, bufNum] 范围内的 buffer 数,TQue 内部状态(BufState 数组/事件索引)按 depth 分配。传 num=2 > depth=1 属于越界使用:轻则在框架断言处直接失败/只分配 1 个 buffer 而代码按 2 个流水理解,重则状态数组越界写破坏相邻成员,导致 kernel 编译期报错或运行期数据错乱。

一致性证据:本仓所有双缓冲用法都把深度与 num 对齐——loss/smooth_l1_loss_grad/op_kernel/arch35/smooth_l1_loss_grad.h:142-145 声明 TQue<QuePosition::VECIN, 2> 后才 214-217 行 InitBuffer(..., 2, ...);loss/sigmoid_focal_loss/op_kernel/arch35/sigmoid_focal_loss_kernel.h:778 声明深度 2 后 391 行传 2。本文件同一函数的 else 分支(200-208 行)也是深度 1 配 num=1。只有新增的小V分支出现深度 1 配 2 的单边偏离。

likedislike
Incorrect?
AlfengYuanAlfengYuanCollaborators
3 days ago Modify description of pull request
johan_leeCollaborators3 days agoStart code review1
loss/fused_cross_entropy_loss_with_max_sum/op_kernel/arch35/fused_cross_entropy_loss_with_max_sum_ar.h
@@ -366,0 +670,4 @@
670+ (AscendC::Reg::RegTensor<uint32_t>&)idxReg, mask);
671+ AscendC::Reg::Gather<float, float, uint32_t>(invReg, invAddr + r0,
672+ (AscendC::Reg::RegTensor<uint32_t>&)idxReg, mask);
673+ LoadVocabToFp32(src, xReg, mask, off);
johan_lee3 days agoComments:

【review】
问题分类:正确性
位置:loss/fused_cross_entropy_loss_with_max_sum/op_kernel/arch35/fused_cross_entropy_loss_with_max_sum_ar.h:673,代码:LoadVocabToFp32(src, xReg, mask, off);
问题:FP32 且 V=3 时 g=16,组步长为 48 个 float。当前 aPerLoop 是 128 的倍数,满 tile 的 vocab 缓冲长度已经整除 64,按 VL 向上对齐不增加尾部空间。最后一组从 tileElems-48 开始,LoadVocabToFp32 仍用不带掩码的 DIST_NORM 读取 64 个 float。
可能导致的结果:最后一次向量搬入会越过当前 vocab 缓冲区 16 个 float,读到相邻 UB 区域。这里的 mask 只约束后续计算和写回,无法约束这次读。只覆盖 V=1 的新增 tiling 用例不会触发该边界。
建议的修改方法:按最后一组的实际起点为每个输入槽预留一个完整向量的可读范围,并把额外空间计入 Host UB 预算,或为不足 VL 的组使用受掩码约束的搬入。增加 V=3、每核至少一个满 tile 的边界用例。

likedislike
AlfengYuanAlfengYuanCollaborators
3 days ago Pre merge succeeded (commit_id: 5793d9e80140fc6b1ab93eba825d959e856fb05f)
AlfengYuanAlfengYuanCollaborators
3 days ago Pushed 1 commits:1024dce7-fix(fused_cross_entropy_loss_with_max_sum): 修复小V快路队列深度越约与末组向量读越界
AlfengYuanAlfengYuanCollaborators
3 days ago Pre merge succeeded (commit_id: 5c17334d58cc923e347a7d5c5eabcdde2b61fc66)
CANN-robotCANN-robotCollaborators
3 days ago Removed label: cann-cla/yes
CANN-robotCANN-robotCollaborators
3 days ago Added label: cann-cla/yes
AlfengYuan
AlfengYuanCollaborators
3 days ago Comments:

感谢评审,两条意见均已确认成立并在 1024dce70 修复:

1. TQue 深度 1 配 InitBuffer num=2(P2,@atomgit-bot):属实,系小V预取流水改造时漏改声明。已将 logitsMaxQue_/sumExpScalarQue_/predictedScalarQue_/lossScalarQue_ 四个队列声明深度 1→2,与 InitBuffer(que,2,...) 对齐;大V分支 num=1 仍在 [1,bufNum] 合法区间内,行为不变。

2. 小V末组向量读越界(@johan_lee):分析完全正确。V>1 且 Gv<VL 时(V=3:G=16、组步长 48<64;V=5/6/7 同理),满 tile 末组 DIST_NORM/UNPACK_B16 满载读越过 tileElems(tileElems=aPerLoopv 恰被 VL 整除),掩码只约束计算与写回、不约束这次读。修复:vocab 源缓冲(vocabQue_ 双槽 + vocabFp32Buf_)各多分配一个 VL,越界读收敛在自有 buffer 内;该开销已由 host 预算中 2V_ALIGN(dtypeSize+4B) 的余量覆盖,aPerLoop 不变。同时按建议新增了 V=3/fp32 的 tiling UT 边界用例(aPerLoop=2560),并在 TTK 侧补跑了 V=3/V=5 fp32(DIST_NORM 路径)与 V=6 fp16(内联 cast 路径)的满 tile 边界用例 3/3 PASS。

顺带处理了两条 P3 建议:inv_sum 热循环内常量 1 向量的 Duplicate 已外提;SoftmaxFlatLoop 的 G 改为调用方统一计算传入,消除重复推导。UT 预算注释推导错误(286→264(dtypeSize+4B))已修正。

回归:ophost UT 8/8(含新增 V=3 用例);49 条原不达标用例精度+性能 49/49 PASS,性能与修复前持平(中位 1.00x,带宽仍 1150~1310 GB/s)。

关于裸 vdiv 的次正规契约分歧(P3):是有意取舍——次正规 sum 时 inv_sum 本就为 inf,补偿除的白做;两通路在该极端输入下的语义差异与 golden/third_party 的溢出定义一致,维持现状。

likedislike
AlfengYuanAlfengYuanCollaborators
3 days ago Pre merge succeeded (commit_id: 97247911d264a4c9d46eb0be301f6a46399aa2de)
5 event messages are collapsed here。 See more
AlfengYuanAlfengYuanCollaborators
3 days ago Modify description of pull request
AlfengYuan
AlfengYuanCollaborators
3 days ago Comments:

关于"裸 vdiv 的次正规契约分歧"(P3)的详细说明与取舍论证:

分歧点

inv_sum = 1/sum_exp_logits 在两条通路用了不同的向量除法实现:

通路 实现 次正规行为
大V通路 ComputeLossInvSum DivAlgo::PRECISION_1ULP_FTZ_FALSE(补偿除,≤1ulp,无 FTZ) 1/次正规 经补偿序列给出正确结果
小V快路 ComputeLossInvSumLoop DivAlgo::INTRINSIC(裸 vdiv,硬件直出) 次正规除数按 FTZ 当 0 处理 → 1/次正规 = inf

即同一输入走不同路径结果可能不同,故评审称"两条通路数值契约分歧"。

精确的分歧窗口

fp32 次正规范围为 (0, 1.175e-38)(最小正规数 2^-126)。对 1/s 细分:

  • s < 2.9e-39(次正规下半段):1/s > 3.4e38 超出 fp32 最大值,补偿除同样得 inf——两通路一致,无分歧;
  • s ∈ [2.9e-39, 1.175e-38)(次正规上半段):补偿除给出有限大值(如 1/1e-38 ≈ 1e38),裸 vdiv 给 inf——真正的分歧只有这一段。

(更正我上一条回复的措辞:"次正规 sum 时 inv_sum 本就为 inf"只对下半段严格成立;上半段补偿除确实能给出有限值。结论不变。)

为什么判断为可接受的取舍

  1. 合法输入触发不了:sum_exp = Σ exp(x_i - max),行最大值项贡献 exp(0) = 1,故真实数据恒有 sum_exp ≥ 1,不可能落入次正规区间。只有绕过 max/sum 自洽关系的病态构造数据才能触发。现有 299 条用例的 sum_exp 最小取值为 0.001(正正规数),回归全量通过。
  2. 分歧数值无实用语义:即使触发,争议结果是 1e38 量级 vs inf——该场景下 softmax 输出 exp(x-max) * 1e38 本就饱和溢出,两种结果均无可用含义。
  3. loss 主输出未受影响:log(sum_exp) 两条通路都保留补偿实现(castLogFt32False),因为 golden 对 log(次正规) 给有限值 -87.5,必须对齐;只有 inv_sum 这个中间量换成了裸 vdiv。
  4. 收益在热路径:小V快路每 G 行一次该除法,裸 vdiv 省掉补偿序列,是第 3 个提交 1.04~1.99x(中位 1.39x)提速的来源之一。

残留风险与回退方式

若未来测试集加入 sum_exp 次正规的对抗用例,小V路径会在 cross_check special 位规则上报 FAIL(NPU=inf,golden/third_party 为有限大值),大V路径不受影响。届时如需堵上,将 castDivIntrinsic 换回 castDivFt32False 即可(单行改动,代价是放弃该部分提速)。

综上,维持裸 vdiv 现状。

likedislike
AlfengYuan
AlfengYuanCollaborators
3 days ago Comments:

compile

likedislike
AlfengYuanAlfengYuanCollaborators
3 days ago Pre merge succeeded (commit_id: a91a5368d1210491f6e261718187122fd2703a93)
CANN-robot
CANN-robotCollaborators
3 days ago Comments:
🚀 CI 流水线已启动
📋 执行详情: 点击查看流水线
likedislike
CANN-robotCANN-robotCollaborators
3 days ago Added label: ci-pipeline-running
CANN-robotCANN-robotCollaborators
2 days ago Removed label: ci-pipeline-running
CANN-robotCANN-robotCollaborators
2 days ago Added label: ci-pipeline-passed
yuning_chen
yuning_chenCollaborators
2 days ago Comments:

您好,pr合入已准备就绪,请尽快联系committer进行检视,谢谢!
@su-yueming @wang-xing001 @east_yang @liubo75 @chenfeng61 @qianzehong @u010470851

likedislike
yuning_chen
yuning_chenCollaborators
2 days ago Comments:

您好,pr合入已准备就绪,请尽快联系committer进行检视,谢谢!
@su-yueming @wang-xing001 @east_yang @liubo75 @chenfeng61 @qianzehong @u010470851

likedislike
chenqi317Collaborators
2 days ago Comments:

/lgtm

likedislike
AlfengYuan
AlfengYuanCollaborators
2 days ago Comments:
likedislike
TangPC
TangPCCollaborators
2 days ago Comments:

/lgtm
/approve

likedislike
CANN-robotCANN-robotCollaborators
2 days ago Added label: lgtmapproved
CANN-robotCANN-robotCollaborators
2 days ago Closed associated issue
CANN-robotCANN-robotCollaborators
2 days ago Merged pull request
CANN-robot
CANN-robotCollaborators
2 days ago Comments:

Pull Request 已合并或已关闭。

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

likedislike