Merged
修复dliklg算子初始化越界 #15406
llwy0320Add on 4 days ago
修复dliklg算子初始化越界 #15406
Merged
Merged fromllwy0320/ops-transformer6847_issueintocann/ops-transformermaster
Total: 1 files changed+22-13
| @@ -714,19 +714,28 @@ __aicore__ inline void DenseLightningIndexerKLLossGradKernelBase<CubeBlockType, | |||
| 714 | softmaxOutPreBlockTotal = tilingData->multiCoreParams.softmaxOutPreBlockTotal; | 714 | softmaxOutPreBlockTotal = tilingData->multiCoreParams.softmaxOutPreBlockTotal; |
| 715 | softmaxOutPreBlockTail = tilingData->multiCoreParams.softmaxOutPreBlockTail; | 715 | softmaxOutPreBlockTail = tilingData->multiCoreParams.softmaxOutPreBlockTail; |
| 716 | 716 | ||
| 717 | - initdqSize = vBlockIdx == qPreBlockTotal - 1 ? qPreBlockTail : qPreBlockFactor; | 717 | + if (vBlockIdx < softmaxOutPreBlockTotal) { |
| 718 | - initdkSize = vBlockIdx == kPreBlockTotal - 1 ? kPreBlockTail : kPreBlockFactor; | 718 | + initsoftMaxSize = |
| 719 | - dqOffset = (uint64_t)vBlockIdx * qPreBlockFactor; | 719 | + vBlockIdx == softmaxOutPreBlockTotal - 1 ? softmaxOutPreBlockTail : softmaxOutPreBlockFactor; |
| 720 | - dkOffset = (uint64_t)vBlockIdx * kPreBlockFactor; | 720 | + softMaxOffset = (uint64_t)vBlockIdx * softmaxOutPreBlockFactor; |
| 721 | - initdweightSize = vBlockIdx == weightPreBlockTotal - 1 ? weightPreBlockTail : weightPreBlockFactor; | 721 | + InitOutput<float>(softmaxOutGm[softMaxOffset], initsoftMaxSize, 0); |
| 722 | - dweightOffset = (uint64_t)vBlockIdx * weightPreBlockFactor; | 722 | + } |
| 723 | - initsoftMaxSize = vBlockIdx == softmaxOutPreBlockTotal - 1 ? softmaxOutPreBlockTail : softmaxOutPreBlockFactor; | 723 | + if (vBlockIdx < weightPreBlockTotal) { |
| 724 | - softMaxOffset = (uint64_t)vBlockIdx * softmaxOutPreBlockFactor; | 724 | + initdweightSize = vBlockIdx == weightPreBlockTotal - 1 ? weightPreBlockTail : weightPreBlockFactor; |
| 725 | - InitOutput<float>(softmaxOutGm[softMaxOffset], initsoftMaxSize, 0); | 725 | + dweightOffset = (uint64_t)vBlockIdx * weightPreBlockFactor; |
| 726 | - InitOutput<float>(dWeightGm[dweightOffset], initdweightSize, 0); | 726 | + InitOutput<float>(dWeightGm[dweightOffset], initdweightSize, 0); |
| 727 | - InitOutput<OUT_T>(dqGm[dqOffset], initdqSize, 0); | 727 | + } |
| 728 | - InitOutput<OUT_T>(dkGm[dkOffset], initdkSize, 0); | 728 | + if (vBlockIdx < qPreBlockTotal) { |
| 729 | - InitOutput<T>(bmm3Res[dkOffset], initdkSize, 0); | 729 | + initdqSize = vBlockIdx == qPreBlockTotal - 1 ? qPreBlockTail : qPreBlockFactor; |
| 730 | + dqOffset = (uint64_t)vBlockIdx * qPreBlockFactor; | ||
| 731 | + InitOutput<OUT_T>(dqGm[dqOffset], initdqSize, 0); | ||
| 732 | + } | ||
| 733 | + if (vBlockIdx < kPreBlockTotal) { | ||
J | |||
| 734 | + initdkSize = vBlockIdx == kPreBlockTotal - 1 ? kPreBlockTail : kPreBlockFactor; | ||
| 735 | + dkOffset = (uint64_t)vBlockIdx * kPreBlockFactor; | ||
| 736 | + InitOutput<OUT_T>(dkGm[dkOffset], initdkSize, 0); | ||
| 737 | + InitOutput<T>(bmm3Res[dkOffset], initdkSize, 0); | ||
| 738 | + } | ||
| 730 | } | 739 | } |
| 731 | SyncAll<false>(); | 740 | SyncAll<false>(); |
| 732 | } | 741 | } |
【review】 问题分类：可测 位置：attention/dense_lightning_indexer_kl_loss_grad/op_kernel/arch35/dense_lightning_indexer_kl_loss_grad_kernel_base.h:733，代码：if (vBlockIdx < kPreBlockTotal) { 问题：这次用每个缓冲区自己的 PreBlockTotal 限制初始化写入。触发场景是 kernel 实际启动的向量块数大于某个缓冲区的 PreBlockTotal，或者尾块只剩 PreBlockTail 个元素。现有算子测试只有 Host InferShape 和 Tiling，后者验证 tiling key、数据与 workspace，不执行这段 InitOutput，因此移除该边界条件后现有断言仍能通过。 可能导致的结果：初始化越界或最后一个有效块漏清零都无法由现有测试识别，设备侧的越界写或梯度污染可能在后续修改中回归。 建议的修改方法：增加一个实际执行 kernel 的回归用例，选用能让 K 梯度的 PreBlockTotal 小于实际向量块数且有非整尾块的形状。用设备内存越界检测或缓冲区哨兵验证额外块不写入，并验证最后一个有效块只处理 PreBlockTail。