环境信息

• 硬件：Ascend_910B
• Python：3.11.15
• PyTorch：2.10.0+cpu
• torch_npu：2.10.0.post5.dev20260828
• CANN：9.1.0

问题描述

torch.ops.aten.std.dim(x, [1], True, False) 经 Inductor Triton 后端编译时存在两个 NPU 代码生成问题：

1. 静态路径使用 NPUTritonKernel。上游 _has_constant_mask 基于 Triton block 为 2 的幂的假设删除尾块 mask，但 NPU autotune 可以选择非 2 的幂 block。例如 xnumel=262144 时选择 XBLOCK=4688，最后一个 block 存在 384 个越界 lane，可能引发 MTE 越界或越界写。
2. 动态路径使用 NPUIndexTritonKernel。上游 TritonPrinter 将 shape 相关的 sympy.Float/ToFloat 打印为 tl.float64，而 Triton-Ascend 无法合法化对应的 f64 向量运算，最终报 MLIRCompilationError/NoTritonConfigsError。

问题 DSL：

tmp13 = (tl.full([], -1.00000000000000, tl.float64)) + ks0
tmp14 = tmp13.to(tl.float32)

后续转换到 float32 发生在 f64 加法之后，无法避免编译失败。

复现方式

import torch


def fn(x):
    return torch.ops.aten.std.dim(x, [1], True, False)


x = torch.randn((256, 8192, 1024), device="npu", dtype=torch.float32)
static_fn = torch.compile(fn, dynamic=False)
dynamic_fn = torch.compile(fn, dynamic=True)
static_fn(x)
dynamic_fn(x)

预期行为

• NPU 非固定 autotune 配置保留尾块 mask；固定配置仍可复用上游优化。
• 仅 NPU Inductor Triton kernel 的符号表达式使用 float32，不修改社区 printer 或其他后端行为。
• 静态和动态 std(dim=1) 均可正确编译和执行。

验证结果

• Ascend_910B 上动态 std(dim=1) 编译、执行成功，结果与 eager 在 atol=1e-2, rtol=1e-2 下相符。

• 全新 Inductor 缓存生成 6 个 Python 文件，tl.float64 命中数为 0。

• 修复后关键 DSL 为：

  tmp13 = (tl.full([], -1.00000000000000, tl.float32)) + ks0