已合并
fix(inductor): support Flex Attention mask_out cppwrapper #45565
Xuan Peng创建于 25 天前
fix(inductor): support Flex Attention mask_out cppwrapper #45565
已合并
Xuan Peng创建于 25 天前
Xuan Peng
25 天前

【合入来源】

Fixes https://gitcode.com/Ascend/pytorch/issues/4431

Forward-port of https://gitcode.com/Ascend/pytorch/pull/45507 from v2.7.1 to master.

【修改方案】

  1. 将 Triton kernel metadata 中的 workspace_size、lock_num 和 lock_init_val 从 NPU autotuner 传递到 NPU C++ wrapper,并适配新版本的 grouped autotune variant metadata 路径。
  2. C++ launcher 根据 metadata 分配 runtime workspace 和同步锁,初始化锁内存,并检查 aclrtMemcpy 返回值。
  3. master 的 Flex Attention 模板已原生使用符号化 grid,因此不再回迁 v2.7.1 专属的 SymbolicGridFn backport。

【资料变更】

不涉及。

【接口变更】

不涉及跨仓或客户可见接口变更;仅调整 Inductor NPU 内部 kernel metadata 传递和 C++ wrapper 代码生成。

【功能验证】

  • python3 -m py_compile torch_npu/_inductor/codegen/cpp_wrapper_npu.py torch_npu/_inductor/runtime/triton_heuristics.py torch_npu/_inductor/kernel/flex_attention.py torch_npu/_inductor/kernel/flexattention_template.py:通过。
  • git diff --check:通过。
  • 隔离代码生成检查:通过;确认生成的 C++ 包含 runtime workspace、同步锁分配与初始化,并确认 pure-SIMT 路径不会生成这些分配。
  • 未运行完整 UT 或 NPU 实机用例:本地 Python 环境未安装 torch,且无 NPU 运行环境;PR CI 待执行。

【CheckList】

likedislike
Pull Request已成功合入, 合并人@ascend-robot
(感谢 Xuan Peng 的贡献)
XXuan Peng
25 天前 创建了 pull request,commit e745f85b
XXuan Peng
25 天前 关联了issue:[Bug]: Flex Attention cppwrapper fails in mask_out backward on v2.7.1
atomgit-bot
atomgit-bot
25 天前 评论:

变更摘要

本 PR 为 torch_npu Inductor NPU 后端补齐 Flex Attention mask_out 的 C++ wrapper 支持(由 v2.7.1 的 #45507 前向移植到 master):将 Triton kernel metadata 中的运行时资源信息(workspace_size、lock_num、lock_init_val)从 NPU autotuner 传入 C++ wrapper,由生成的 C++ launcher 在运行时分配 workspace 与同步锁并完成锁内存初始化。核心改动位于 torch_npu/_inductor/codegen/cpp_wrapper_npu.py 与 torch_npu/_inductor/runtime/triton_heuristics.py;master 的 Flex Attention 模板已原生使用符号化 grid,因此无需回迁 v2.7.1 的 SymbolicGridFn backport。

主要改动

  • kernel metadata 透传: 在 NPUCachingAutotuner 与 NPUSymbolicGroupedAutotuner 的 params 中新增 lock_num、lock_init_val、workspace_size 三个字段(通过 getattr 从 metadata 读取并转 int,缺省为 0),适配新版本 grouped autotune variant 的 metadata 路径。
  • workspace 运行时分配: CppWrapperNpu.generate_args_decl 新增 kernel_params 参数,在非 pure-SIMT 且 workspace_size > 0 时生成按 grid_0 * grid_1 * grid_2 缩放大小的 allocate_workspace 调用,并将 workspace_addr 指向其 storage 数据。
  • 同步锁分配与初始化: 在非 pure-SIMT 且 lock_num > 0 时分配同步锁内存,构造 lock_init_val 初值向量,通过 aclrtMemcpy(ACL_MEMCPY_HOST_TO_DEVICE)初始化,并检查返回值,失败时抛出包含错误码的 std::runtime_error。
  • wrapper 调用链调整: DeferredNpuTritonCallWrapper 改为先调用 generate_args_decl(传入 is_triton_kernel=True、is_pure_simt、kernel_params=params)再拼装 launch_call lambda;同时 add_device_include 新增 torch_npu/csrc/core/npu/NPUWorkspaceAllocator.h 头文件引用。
  • pure-SIMT 路径隔离: 生成的 workspace 分配与同步锁分配均以 not is_pure_simt 为前置条件,pure-SIMT 路径不会生成这些运行时分配代码,保持原有行为不变。
likedislike
不准确?
atomgit-bot
atomgit-bot
25 天前 评论:

代码审查

✅ 未发现问题

likedislike
不准确?
ascend-robotascend-robot成员
25 天前 添加了label:ascend-cla/yes
ascend-robot
ascend-robot成员
25 天前 评论:

Thanks for your pull-request.
The full list of commands accepted by me can be found at here.
You can get sig-info at here.
You can self-configure the PR merge rules for this repository. For more details, please refer to Here.


PR Approval Progress

✅ Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

module lgtm status approve status
torch_npu/_inductor ✅ weizhan4, crazyDannyBoy (2/2) ✅ weizhan4, crazyDannyBoy (2/1)

💡 Tip:

  • Committer can comment /approve or /lgtm
  • Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

HinPeng, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
ascend-robot
ascend-robot成员
25 天前 评论:

当前仓库存在以下 保护分支 :

Protected Branch Version Release
master
v2.7.1-26.1.0
v2.10.0
v2.9.0
v2.7.1
v2.11.0
v2.12.0
v2.9.0-26.1.0
v2.12.0-26.1.0
v2.11.0-26.1.0
v2.10.0-26.1.0
ci-test

评论 /sync <branch1> <branch2> ... 可将当前 PR 修改同步到其它分支(创建同步 PR):
a) 如果当前 PR 是 Open 状态,同步操作将延迟到 PR 被合并时执行
b) 如果当前 PR 已经 Merged,将立即执行同步操作

注意:

  1. /sync 命令可以指定同步到多个分支,仅最后一个 /sync 命令生效
  2. 如果创建的同步 PR 不正确,可通过向同步 PR 的源分支提交轻量级 PR 完善,或使用 /close 命令关闭
likedislike
AtlasAccountAtlasAccount成员
25 天前 添加了label:ci-pipeline-running
ascend-robot
ascend-robot成员
25 天前 评论:

ascend docs pipeline is running...

likedislike
ascend-robotascend-robot成员
25 天前 添加了label:docs-ci-pipeline-running
ascend-robot
ascend-robot成员
25 天前 评论:

✅ 跳过 docs ci 检查,没有需要检查的文档文件

likedislike
ascend-robotascend-robot成员
25 天前 删除了label:docs-ci-pipeline-running
ascend-robotascend-robot成员
25 天前 添加了label:docs-ci-pipeline-success
AtlasAccountAtlasAccount成员
25 天前 删除了label:ci-pipeline-running
AtlasAccountAtlasAccount成员
25 天前 添加了label:ci-pipeline-passed
AtlasAccount
AtlasAccount成员
25 天前 评论:
流水线 PR-pipeline_pytorch#64929 [ commitID:b085a5c9 ] 已完成
>>>代码风格自动修复执行成功(无修复内容)
阶段 任务名 状态 详情
编译构建 Build_X86 ✅ >>>
Build_ARM ✅ >>>
Build_X86_torchair 🛑 >>>
Build_ARM_torchair 🛑 >>>
patch_test 🛑 >>>
Build_X86_213 ✅ >>>
Build_ARM_213 ✅ >>>
恶意代码检查 Antipoison ✅ >>>
编码安全与规范检查 codecheck_pre-commit ✅ >>>
check_error ✅ >>>
lintrunner ✅ >>>
开源片段检查 SCA ✅ >>>
开发者测试 UT_ARM_A3_Part_01 🛑 >>>
UT_ARM_A3_Part_02 🛑 >>>
UT_ARM_A2_Part_01 ✅ >>>
UT_ARM_A2_Part_02 ✅ >>>
UT_ARM_A2_Part_03 ✅ >>>
UT_inductor_Part_01 🛑 >>>
UT_inductor_Part_02 🛑 >>>
UT_inductor_Part_03 🛑 >>>
UT_inductor_Part_04 🛑 >>>
UT_DIST_ARM_Part_01 🛑 >>>
UT_DIST_ARM_Part_02 🛑 >>>
UT_DIST_ARM_Part_03 🛑 >>>
UT_DIST_ARM_Part_04 🛑 >>>
UT_ARM_A2_Select_Part_01 ✅ >>>
UT_ARM_A2_Select_Part_02 ✅ >>>
UT_ARM_A2_Part_213 ✅ >>>
UT_inductor_Part_213 🛑 >>>
UT_DIST_ARM_Part_213 🛑 >>>
UT_ARM_A2_Select_Part_213 ✅ >>>
流水线 PR-pipeline_pytorch ✅ >>>
此流水线已支持下列评论快捷指令,仅PR创建者和白名单成员[wujinyuan1, huangjingwei, liangsongwei, yashi999, culechan, Dring, wuyouqi1, L1919_snow, qq_52711437, WhiteNight12, nomiz, xiu_21, ffmh, wanglijun55, hss-shuai, husichao, smallsilly, lanshaozuishuai, jimmyisme1, lzy0920232, alpha-junh, Sunshine_Youngster, wei_zhuoyi, zhangyihuiben, zyw-hw, zzzkeke, rmch, yangch0324, LucciC, AACAES, renyujin, wjlflyer, senzhen-town, pengjingyou, qsc97, limuan, yule100, xiaoqi-zhou, kuhn7, chenxingying, hanye02, zichun_ye, anyrenwei, kkjocker, wangzili121, Lu_G, yvjc, puddingfjz, HandsoemLemon, bigprestigee1, huawuyi, zhenyu10, dairenjie, du-jin-hang, zou-jieyu, adelaideliu, TrHan, wanlinan, Windwindzzz, pengqihw, kisnwang, yuheng_wang, honghao_wang, jizewei, zhangguoguang, sunyu-xuan, chenrayray, hbhu_bin, liujunzhu, c_34, LiNuoh, maoyuanpeng1, zzhongmin, zhaoyu65, bellatan, jiabaolin, zhuofanshen, wencaiwen, lu_zhuge, caoshuyang, molly12, lyx324521, LQ1206, gitcode-chenjiao, cai-weiwei1989, CHDong, ogqin, yuanlipingGit, xuqinglin1, lqz2, zouwei1, chaoluoa, paradox325, jackzhang1116, yaoyao, akh, yujiacheng, dengjie0116, Hubert11111, Shine_Ws, wslhj555, longqiand, OYtao666, JiaqingQiang, luyyyy, Kingbelial, zhanghaiyu0101, wenxp1018, yanliu-luoluo, ksun_sekiro, liyong328, wgzheng, tangky, vivi_is_coding, aoiaoisola, weixin_44494597, wangmengmengwang65667, hid57809721, qq_35468730, comeonup, C547032, gcw_m5OQChA4, yao_yao_ling_xian, cnnbwcy, szqfes_12, cora_19, cann_lilin, can, shawnylee233, fanglanyue0916, hhz0, LiNuohang, taohuoquan, Jesse, WSs_321, SCh_zx]评论有效
  • compile、compile_inductor、compile_torchair : 运行流水线
  • retry : 重试流水线所有失败子任务
  • retry <任务名> : 仅重试指定失败子任务
  • stop : 停止流水线
likedislike
zhucehw成员
24 天前 评论:

/lgtm

likedislike
ascend-robot
ascend-robot成员
24 天前 评论:

The following users do not have permission to comment /lgtm or /approve on any module in this PR:
zhucehw

likedislike
weizhan4成员
24 天前 评论:

/lgtm
/approve

likedislike
ascend-robot
ascend-robot成员
24 天前 评论:

The following users do not have permission to comment /lgtm or /approve on any module in this PR:
zhucehw

likedislike
ascend-robotascend-robot成员
24 天前 添加了label:approved
ascend-robot
ascend-robot成员
24 天前 评论:

The following users do not have permission to comment /lgtm or /approve on any module in this PR:
zhucehw

likedislike
dezheng889成员
24 天前 评论:

/lgtm
/approve

likedislike
ascend-robot
ascend-robot成员
24 天前 评论:

The following users do not have permission to comment /lgtm or /approve on any module in this PR:
zhucehw

likedislike
ascend-robotascend-robot成员
24 天前 添加了label:lgtm
ascend-robotascend-robot成员
24 天前 合入了pull request
AtlasAccount
AtlasAccount成员
24 天前 评论:
流水线 pytorch_gitcode_PR_multiVersion#14781 [ commitID:b085a5c9 ] 已完成
likedislike