【合入来源】

如有社区issue，请关联issue链接
请勿携带内部流程信息（需求链接、问题单、内部issue等）

【合入来源】

• 需求
• 问题单
• issue/工单: https://gitcode.com/Ascend/pytorch/issues/4375
• 重构优化
• 资料更新

【修改方案】

问题现象：npu_backend=triton_experimental 时，输出为单个标量的全量归约融合 kernel（CausalLM 的 log_softmax + nll_loss(ignore_index) 路径，典型如 triton_unk_fused_clone_nll_loss_forward_slice_view_*）编译失败：

ValueError('Cannot broadcast, rank mismatch: [1, 1, 1], [1, 1]')
tl.store(in_out_ptr0 + (tl.full([1, 1], 0, tl.int32).broadcast_to(XBLOCK, 1)), tmp15, None)

HuggingFace runner 的 ElectraForCausalLM / RobertaForCausalLM / XGLMForCausalLM 三模型均因同一 kernel 模板失败。

根因分析（上游 torch/_inductor 与 triton_experimental 的交互）：

1. 上游 store codegen（torch/_inductor/codegen/triton.py 的 integer-index store 分支）会手动追加 .broadcast_to(<value.shape>)，其中 value.shape 取自发射时刻 CSE 记录的形状——此刻归约值与索引秩一致（均为 2），代码正确；
2. triton_experimental 随后做 r 树提升（promoted r-tree）改写，把归约值表达式整形为 [None, None, :]（秩 3），但 store 行中手动追加的 broadcast_to 参数与基座 tl.full 仍停留在提升前的秩 2；
3. Triton 要求 tl.store 两侧秩相等，[1,1,1] vs [XBLOCK,1] 直接编译失败。

修复内容：在既有的 store 回改函数 _rewrite_reduction_store_shape（torch_npu/_inductor/triton_experimental/codegen/triton.py）内补全分支——当 broadcast_to 参数数小于 real_ndim（陈旧秩）时重建参数：r 槽强制置 1（常量索引 ⟹ 全量归约 ⟹ 值在 r 轴为 singleton，类型上保证）、旧的非 1 参数按序映射到其余槽位、基座 tl.full 的 shape 列表补齐到 real_ndim。写入地址、数值与冗余写次数均不变，仅修正索引表达式的形状元数据，语义严格等价；参数计数不匹配时保持原行不改（防御性 bail）。

【资料变更】

不涉及

【接口变更】

不涉及

【功能验证】

验证环境：Ascend NPU（aarch64）+ torch 2.13.0 + torch_npu（triton_experimental 后端）。

1. 新增 UT：test/_inductor/test_triton_experimental_rtree_store.py，覆盖 CausalLM shift + log_softmax + nll_loss(ignore_index) 触发模式（双级 r 树提升 + 标量整数索引 store）。修复前该 UT 失败，报错与线上三模型完全一致（rank mismatch: [1, 1, 1], [1, 1]）；修复后通过（含与 eager 数值一致性断言）。
2. 既有冒烟回归：test/_inductor/test_triton_experimental_enable.py 13/13 通过，无回归。
3. 端到端验证（HuggingFace benchmark runner，triton_experimental 后端）：
   • ElectraForCausalLM：修复前编译失败 → 修复后 1.155x 加速、pass_accuracy；
   • RobertaForCausalLM：修复前编译失败 → 修复后 1.318x、pass_accuracy；
   • XGLMForCausalLM：修复前编译失败 → 修复后 1.502x、pass_accuracy；
   • 原正常模型（resnet18 / fastNLP_Bert 等）无劣化。

【CheckList】

• 代码注释完备，正确记录错误日志
• 代码实现进行了返回值、空指针等校验
• PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等
• PR持续集成流水线（CI）执行通过，代码检查无异常