环境信息

以下信息整理自 PR #47833 中记录的排查与验证：

• 仓库 / 分支：Ascend/pytorch master，修复基线 7eb2fe2f7。
• 硬件：Ascend950PR。
• CANN：9.1.0；PyTorch：2.13.0。
• 后端：TORCHINDUCTOR_NPU_BACKEND=dvm。
• 输入：BF16，Q/K/V 形状均为 [8, 128, 2048, 128]。

🐛 问题描述

使用 DVM 后端，通过 torch.compile(..., fullgraph=True) 编译 causal FlexAttention 时，原实现禁止 DVM 使用 mask-out，进入 mask-in 路径后编译失败，最终报 NoValidChoicesError，无法完成该场景的前向和反向计算。

启用 mask-out 后，反向还会调用 sum(axis=-1)；公共 MLIR reduction lowering 的 sum_() 未将整数 axis 规范化为列表，负轴处理尝试遍历整数时会抛出 TypeError。

复现步骤

1. 在导入后端之前设置 TORCHINDUCTOR_NPU_BACKEND=dvm，加载 torch_npu._inductor。
2. 构造形状为 [8, 128, 2048, 128] 的 BF16 Q/K/V，并启用梯度。
3. 使用 q_idx >= kv_idx 构造 causal block mask。
4. 对调用 flex_attention 的函数执行 torch.compile(..., backend="inductor", fullgraph=True)，运行前向及反向。

实际结果：未修复时在编译阶段报 NoValidChoicesError；仅启用 mask-out 后还需处理反向 reduction 的整数 axis。

预期结果与修复方向

该 causal 场景能够完成编译、前向和反向，输出及 Q/K/V 梯度均为有限值，梯度形状与输入一致。

• 允许 DVM 使用已有的 FlexAttention mask-out 路径。
• 在公共 MLIR sum_() lowering 入口将整数 axis 转为列表，再执行负轴处理。

关联 PR 与验证范围

修复 PR：https://gitcode.com/Ascend/pytorch/pull/47833

PR 已记录该大形状 causal 前向/反向通过，以及现有 MLIR reduction 回归通过。本 Issue 聚焦上述编译问题；自定义 score 的精度问题、进程内后端切换问题和未使用 torch.compile 时的 eager HBM OOM 不作为本 Issue 的验收项。