变更摘要

本次 PR 主要修复了 Inductor 后端在多轴 reduction（特别是 RMS norm 权重梯度计算）场景下的代码生成问题。核心改动集中在 NPUIndexTritonKernel 和 ReductionAnalysis 两个类中：对多 reduction 轴且连续 reduction 的情况新增了专门的代码生成路径，包括循环体结构重组、store 写入时机调整、以及带 permute 的 reshape 处理。同时补充了对应的单元测试用例。

主要改动

• 多轴 reduction 的 golden var list 选择：在 kernel_analysis.py 的 ReductionAnalysis 中，当 reduction 轴多于 1 且为连续 reduction 时，若 golden_var_list 为空则主动调用 select_golden_varlist()，并使用 reduction 轴对应的 golden var 数量作为结果。

• 多轴 reduction 循环体代码生成重构：在 triton.py 的 NPUIndexTritonKernel 中新增 is_first_reduction_tiling 和 use_outer_reduction_post_loop 判断逻辑，区分首个 reduction tiling 轴与非首个轴，调整 post_loop_combine、post_loop_store 以及普通 stores 的写入时机与顺序，确保多轴 reduction 场景下代码生成正确。

• 带 permute 的 dense reshape 支持：在 _map_tuple_or_scalar 相关代码中，对多轴连续 reduction 场景计算 golden_var_list 的排列顺序，必要时生成 tl.reshape(v.permute(...), ...) 而非直接 tl.reshape(v, ...)，同时 mask 条件表达式也应用相同的 permute 变换。

• 多轴 reduction 的 promote_to_tensor reshape：在 promote_to_tensor 方法中增加对多轴连续 reduction 的处理，使用 golden_var_list 确定哪些轴参与 reduction 并将对应维度置为 "1"，生成正确的 reshape 字符串。

• 新增 RMS norm 权重梯度测试用例：在 test_var_mean_add_mul.py 中新增 rms_norm_weight_grad、check_rms_norm_weight_grad 辅助方法以及 test_rms_norm_weight_grad_head_seq_permute、test_rms_norm_weight_grad_batch_seq_permute 两个测试方法，覆盖 head 维度和 batch 维度 permute 的 RMS norm 权重梯度计算场景。

openLiBing可接管PR, PR详情链接:
https://www.openlibing.com/apps/transitionPrDetail?owner=Ascend&repo=pytorch&number=42071&platform=gitcode&codeHostingPlatformFlag=gitcode

代码审查

I've now completed the review of all three changed files. Let me summarize:

Closing Summary

Findings by Priority

• P1: 1 finding — 在 codegen_body 中，当多 reduction 轴场景下唯一的 reduction tiling 恰好也是最后一个 tiling 时，prefix 和所有 post-loop 代码（post_loop_combine、post_loop_store、stores、deferred_reduction_stores）被静默跳过，导致生成的 Triton kernel 不完整（triton.py:2701-2702）

• P2: 1 finding — 在 is_first_reduction_tiling 分支中使用 writeline 写入 IndentedBuffer._lines 已带缩进的内容，可能导致重复缩进（triton.py:2737-2740）

Files Reviewed

Overall Risk Judgment

中等偏高。P1 问题在特定条件下（多 reduction 轴被合并为单个 reduction tiling 且非 persistent 模式）会导致 prefix 和 post-loop 代码被完全跳过，产生不完整/无效的 Triton kernel。这正是 Issue #2543 所描述的 rms norm codegen error 场景。建议在合入前修复 P1 问题。

⛔ 需要修改

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

PR Approval Progress

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

Module Approval Details

💡 Tip:

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

CLA Signature Pass

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

Linking Issue Notice

@luqichao , the pull request must be linked to at least one issue.
If an issue has already been linked, but the needs-issue label remains, you can remove the label by commenting /check-issue .

当前仓库存在以下 保护分支 ：

评论 /sync <branch1> <branch2> ... 可将当前 PR 修改同步到其它分支（创建同步 PR）：
a) 如果当前 PR 是 Open 状态，同步操作将延迟到 PR 被合并时执行
b) 如果当前 PR 已经 Merged，将立即执行同步操作

注意：

1. /sync 命令可以指定同步到多个分支，仅最后一个 /sync 命令生效
2. 如果创建的同步 PR 不正确，可通过向同步 PR 的源分支提交轻量级 PR 完善，或使用 /close 命令关闭

compile

• compile、compile_inductor、compile_torchair : 运行流水线
• retry : 重试流水线所有失败子任务
• retry <任务名> : 仅重试指定失败子任务
• stop : 停止流水线

ascend docs pipeline is running...

ascend docs pipeline is running...

✅ 跳过 docs ci 检查，没有需要检查的文档文件

✅ 跳过 docs ci 检查，没有需要检查的文档文件

• compile、compile_inductor、compile_torchair : 运行流水线
• retry : 重试流水线所有失败子任务
• retry <任务名> : 仅重试指定失败子任务
• stop : 停止流水线

/lgtm

/lgtm
/approve

Pull Request 已合并或已关闭。

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