问题描述

在 Torch 2.13 + TorchNPU 的 NPU FlexAttention 中，直接调用 torch.ops.higher_order.flex_attention_backward 并传入空的 kernel_options={} 时，Inductor 在渲染 mask-out backward 模板阶段失败：

torch._inductor.exc.InductorError:
LoweringException: UndefinedError: 'PRESCALE_QK' is undefined

影响范围

• PyTorch: 2.13.0+cpu
• TorchNPU: 2.13.0.dev20260907
• Triton: 3.6.0
• Triton Ascend: 3.6.0
• NPU FlexAttention direct backward / mask-out lowering

最小复现

运行社区测试：

pytest -vs test_flex_attention213.py::TestFlexAttentionPRIVATEUSE1::test_direct_backward_preserves_explicit_buffers_npu

该测试通过 higher-order op 直接传入空字典，绕过 public flex_attention API 的 _apply_kernel_options 默认值填充。

根因

torch_npu/_inductor/kernel/flex_attention.py 的 NPU backward lowering 会补齐 ROWS_GUARANTEED_SAFE、BLOCKS_ARE_CONTIGUOUS、FLOAT32_PRECISION 和 IS_DIVISIBLE，但未补齐模板需要的：

PRESCALE_QK = False
WRITE_DQ = True

其中 PRESCALE_QK 在 Jinja 条件中被引用，因此尚未生成 Triton tl.constexpr 前就触发 UndefinedError。上游 Torch 2.7.1 在 public API 层补齐这些默认值；Torch 2.13 还会在 Inductor backward lowering 层防御性补齐，支持绕过 public API 的内部调用。

期望行为

NPU backward lowering 应在模板路由和渲染前补齐公共 kernel option 默认值，使 direct higher-order-op 调用与 public API 调用保持一致。

验证结果

补齐 PRESCALE_QK=False 和 WRITE_DQ=True 后：

• 新增 source-level 回归测试通过；
• 相关 FlexAttention source tests：6 passed, 2 skipped；
• 原失败 NPU 用例：1 passed；
• 生成的 output_code.py 包含：

PRESCALE_QK : tl.constexpr = False
WRITE_DQ : tl.constexpr = True