变更摘要

本 PR 针对设置 TORCHINDUCTOR_NPU_BACKEND=dvm 后编译 BF16 causal FlexAttention（B=8, H=128, S=2048, D=128）因原 mask-in 路径的 UB 分配/同步问题报 NoValidChoicesError 的问题，改为让 DVM 使用现有 mask-out 路径，从而支持该形状的前向与反向计算。修改仅涉及 DVM 注册入口、DVM wrapper 与 DVM 测试，Triton 后端代码及共享 FlexAttention kernel/template 保持不变。

主要改动

• 新增 DVM FlexAttention 注册入口: 新增 torch_npu/_inductor/dvm/flex_attention.py 的 register_flex_attention()，在 flex_attention（options_index=6, graph_index=3）与 flex_attention_backward（options_index=11, graph_index=7）两个 lowerings 上包裹 _wrap_lowering，从而在 DVM lowering 期间启用 mask-out。
• 启用 mask-out 并绑定 lowering helper: _wrap_lowering 通过 patch.object(flex, "_supports_flex_attention_mask_out", lambda: True) 强制走 mask-out，并结合 _is_score_mod_identity_graph 判断是否保留自定义 score 的 mask-in 路径，同时把 lowerings、empty、empty_strided、to_dtype 替换为 inductor_lowering 的版本。
• 补充同步与 buffer reuse 编译配置: 针对 flex_attention_template 与 flex_attention_backward_template 复制模板并覆写 NPUTemplateCompileOption，新增 sync_solver、disable_tightly_coupled_buffer_reuse 两个编译选项。
• 规范化 DVM reduction 的整数 axis: 对 torch.ops.aten.sum 的 lowering 包裹 sum_，当 axis 为 int 时转换为 [axis] 再调用原实现。
• DVM wrapper 复用 runtime helper 生成机制: NpuDvmWrapperCodeGen 新增 runtime_helper_definitions 初始化、write_async_compile_wait 中追加 _RuntimeHelperDefinitionsLine，并提供 write_runtime_helper_once 转调 NPUPythonWrapperCodeGen。
• 切换 DVM 后端的注册调用与回归用例更新: _load_dvm_backend() 改为调用 register_flex_attention()；测试将大尺寸用例形状调整为 (1, 2, 256, 128) 并收紧断言阈值，新增 test_large_causal_attention_backward 覆盖 (8, 128, 2048, 128) causal 前反向，并在后端切换回归中改用环境变量启用 DVM。

主要改动

• DVM FlexAttention 注册入口: 新增 torch_npu/_inductor/dvm/flex_attention.py 的 register_flex_attention()，包裹 flex_attention 与 flex_attention_backward 的 lowerings。
• mask-out 与 helper 绑定: _wrap_lowering 通过 patch _supports_flex_attention_mask_out 返回 True 启用 mask-out，并按 _is_score_mod_identity_graph 判断保留自定义 score 的 mask-in 路径，同时覆写 lowerings/empty/empty_strided/to_dtype。
• 模板编译选项调整: 为 flex_attention_template 与 flex_attention_backward_template 复制并设置 NPUTemplateCompileOption 的 sync_solver 与 disable_tightly_coupled_buffer_reuse。
• aten.sum axis 规范化: 对 torch.ops.aten.sum 包裹 sum_，将 int 类型的 axis 统一为列表形式。
• wrapper 与入口调用更新: NpuDvmWrapperCodeGen 新增 runtime helper 相关方法与 _RuntimeHelperDefinitionsLine 输出，_load_dvm_backend() 改为调用 register_flex_attention()；测试新增 test_large_causal_attention_backward 并调整后端切换用例。

代码审查

✅ 未发现问题

/approve

/check-pr