关联 issue: #4384

【合入来源】

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

• 需求
• 问题单
• issue/工单
• 重构优化
• 资料更新

triton_experimental 后端两类问题：①softmax 族 Triton 融合在 Ascend 上性能劣化（带掩码注意力 softmax 融合物化 [B,H,S,S] 掩码矩阵、损失路径 log_softmax 拖带 gather 拆解、宽行 softmax
退化为单核串行扫描）；②npu_expand 降低链在 accuracy 流程无限递归（RecursionError）。

【修改方案】

涉及组件：torch_npu/_inductor/triton_experimental/（Python 层，lowering.py / config.py / __init__.py），与 Inductor 降低注册表（torch._inductor.lowering 的 lowerings/decompositions 表）交互，无 C++
改动。

1. softmax 三级路由（lowering.py::_register_softmax_aclnn_fallback，_activate 挂载）

上游 softmax 分解为三段式 amax→exp→÷sum（两次归约真依赖），调度上 ≤1024 行宽走持久归约（行驻留、可邻域融合），更宽行本应走 split_reductions 跨核拆分，但 TE 已显式关闭 split（历史 OOB 规避）→
宽行退化为每行单 program 串行扫描。msprof（910B2）实测宽行（30k/50k）Triton softmax 比 aclnnSoftmax 慢 2.5–2.8x。据此分三类路由：

1. aten._safe_softmax（掩码复合）：分解为 aten._softmax。两种旧路径均劣化——Triton 融合会把掩码物化进 kernel（TrOCR 实测 10.3 ms/iter 额外设备时间）；直接 fallback 则 eager 复合实现运行冗余 -inf
   守卫链 IsNegInf→SWhere→All（4.6 ms/iter，输入已被 masked_fill 预掩码）。分解后交行宽路由接管。
2. log_softmax 族：整体 aclnn 回退（损失路径 Triton log_softmax 慢于 aclnnLogSoftmax，且融合拖带 Gather_AsStrided 拆解 +2.9 ms/iter）。
3. 纯 softmax 行宽路由（softmax_aclnn_max_fuse_numel=256）：行宽 ≤256 保留 Triton 融合（邻域融合正收益，BertForMaskedLM 实证），>256 回退 aclnnSoftmax。行宽取 sympy 静态求值（size_hint
   兜底，异常时保守走 Triton）。

2. npu_expand 递归修复（RecursionError）

模块体先 from torch._inductor.lowering import expand as _orig_expand，随后 monkeypatch lowering.expand = npu_expand。模块二次执行时 from-import 绑定到已 patch 的 npu_expand → _orig_expand 自别名 →
npu_expand:385 无限自递归（accuracy 流程先跑 eager 改变加载时序触发；perf 流程及缓存命中不触发，长期掩盖）。修复双保险：

1. 原始上游句柄首次执行时暂存到永不被重载的上游模块属性 _npu_upstream_expand_saved，_orig_expand 恒取暂存；
2. tail-bcast 物化分支（Pointwise.create+realize 拖拽惰性输入求值可绕回 broadcast 链）加 threadlocal 重入哨兵，重入即返回上游结果。

3. 新增配置开关（config.py）：safe_softmax_aclnn_fallback / log_softmax_aclnn_fallback（默认 True）、softmax_aclnn_max_fuse_numel（默认 256，0 = 关闭路由恒走 Triton），均可独立关闭用于 A/B 回归。

【资料变更】

不涉及

【接口变更】

不涉及跨代码仓/客户面接口。仓内 triton_experimental.config 新增 3 个内部调优开关（见修改方案 3），仅影响本后端内部降低策略。

【功能验证】

测试场景：Ascend 910B2 + torch 2.13 + torch_npu(triton_experimental) + triton-ascend 3.2.2，torchbench/huggingface benchmark runner（--backend inductor --npu-backend triton --inference --disable-aclgraph），全部在清空缓存后全新编译降低条件下执行（排除缓存命中假阳性）。

测试方法：

1. 性能：--performance，30 迭代中位数（eager vs compile）；
2. 精度：--accuracy（tolerance 1e-4），并确认无 RecursionError；
3. 设备侧归因：torch_npu.profiler（msprof 集成）逐 kernel 设备时间对账，4.64 ms/iter 残差 100% 定位到 -inf 守卫链；
4. 尺寸效应：softmax 行宽 64→50265 扫描确立路由阈值。

结果：

回归：resnet18 1.66x / drq 1.23x / yolov3 持平（无 softmax 路径模型不受影响）；修改前 accuracy 流程 RecursionError 必现的 4 个 torchbench 模型（resnet18/drq/fastNLP_Bert/yolov3）修复后全部正常完成且 Rec=0。

UT：本 PR 未新增 UT，看护方式为端到端 benchmark + 精度流（上述）；后续可在 pytorch/test/_inductor/ 补 softmax 路由的降级断言用例（区分 _safe_softmax 分解、log_softmax 回退、宽行路由三分支）。

【CheckList】

• 代码注释完备，正确记录错误日志
• 代码实现进行了返回值、空指针等校验（行宽求值 try/except 兜底、getattr 判空、orig is not None 前置判断）
• PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等（建议标题：fix(triton_experimental): softmax aclnn 路由优化 + npu_expand 递归修复）
• PR持续集成流水线（CI）执行通过，代码检查无异常（分支 te-softmax-aclnn-routing 已推送，以流水线实际执行结果为准）