已合并
add flexattention score_mod support zeros_and_scatter op #45609
stonexxx创建于 7 天前
add flexattention score_mod support zeros_and_scatter op #45609
已合并
Pull Request已成功合入, 合并人@ascend-robot
（感谢 stonexxx 的贡献）
atomgit-bot
7 天前 评论：
7 天前 评论：
变更摘要
本 PR 为 NPU 上的 flex attention score_mod 增加对 zeros_and_scatter 算子的支持：在 torch_npu/_inductor/kernel/flex_attention.py 中为其注册附加 lowering，并将 pointwise 子图 lowering 过程中新建的 buffer 注册到 V.graph；同时将用户传入的 kernel_options（如 BLOCK_M、BLOCK_N）透传到 FlexAttentionConfigGenerator 及前向/反向候选配置生成函数，并剥离无意义的 GPU 专用 BACKEND 选项。配套地删除了不再需要的 patch_extract_read_writes 补丁（含 torch_npu/_inductor/dependencies.py 整个文件），从 fallback 列表移除已由自定义 lowering 覆盖的右移与 alias 算子，并新增对应测试用例。
主要改动
- 支持
zeros_and_scatter的 lowering：在_get_flex_attention_additional_lowerings中注册torch.ops.flex_lib.zeros_and_scatter.default对应的zeros_and_scatter_lowering，并在_build_subgraph_buffer_with_additional_lowerings中把 pointwise 子图 lowering 期间新建的 buffer 统一注册到V.graph.register_buffer，以支持 subgraph 内的 buffer 创建/变更。 kernel_options透传与校验：FlexAttentionConfigGenerator、generate_fwd_candidate_configs、generate_bwd_candidate_configs新增kernel_options参数，按 FWD/BWD/BWDDQ/BWDDKDV 模式解析对应BLOCK_M/BLOCK_N键，校验其为正 2 的幂、多个键取值一致且与 sparse block size 兼容；各调用点（前向、反向、BWDDQ/BWDDKDV）同步传入该参数。- 剥离 GPU 专用选项：在
_register_npu_inductor_flex_attention的前向与反向注册路径中移除kernel_options里的"BACKEND"键，避免其泄漏到 Triton constexpr 参数。 - 清理不再需要的补丁与 fallback：删除
torch_npu/_inductor/dependencies.py（原强制normalize=False的patch_extract_read_writes）及torch_npu/_inductor/__init__.py中对应的导入与调用；从lowering_fallback_list.py移除aten.__rshift__系列算子与aten.alias.default。 - 新增测试覆盖：新增
test/_inductor/test_flexattnetion.py，验证 epilogue（如.cos()）在 flex attention 模板中融合，以及kernel_options={"BLOCK_M": 16}能正确生效于生成的 Triton 代码。
不准确？
atomgit-bot
7 天前 评论：
7 天前 评论：
7 天前 添加了label：stat/needs-squash
7 天前 添加了label：needs-issue
7 天前 添加了label：ascend-cla/yes
ascend-robot
7 天前 评论：
7 天前 评论：
Linking Issue Notice
@stonexxx , the pull request must be linked to at least one issue.
If an issue has already been linked, but the needs-issue label remains, you can remove the label by commenting /check-issue .
ascend-robot
7 天前 评论：
7 天前 评论：
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 |
|---|---|---|
| test | ✅ HinPeng, rain-666 (2/2) | ✅ HinPeng (1/1) |
| torch_npu/_inductor | ✅ HinPeng, rain-666 (2/2) | ✅ HinPeng (1/1) |
💡 Tip:
- Committer can comment
/approveor/lgtm- Commenting
/approveimplies both code review (lgtm) and intent to merge (approve)
CLA Signature Pass
stonexxx, thanks for your pull request. All authors of the commits have signed the CLA. 👍
ascend-robot
7 天前 评论：
7 天前 评论：
当前仓库存在以下 保护分支 ：
| Protected Branch | Version | Release |
|---|---|---|
| master | ||
| v2.12.0 | ||
| v2.11.0 | ||
| v2.7.1 | ||
| v2.10.0 | ||
| v2.9.0 | ||
| v2.10.0-26.1.0 | ||
| v2.11.0-26.1.0 | ||
| v2.12.0-26.1.0 | ||
| v2.9.0-26.1.0 | ||
| v2.7.1-26.1.0 | ||
| ci-test |
评论 /sync <branch1> <branch2> ... 可将当前 PR 修改同步到其它分支（创建同步 PR）：
a) 如果当前 PR 是 Open 状态，同步操作将延迟到 PR 被合并时执行
b) 如果当前 PR 已经 Merged，将立即执行同步操作
注意：
- /sync 命令可以指定同步到多个分支，仅最后一个 /sync 命令生效
- 如果创建的同步 PR 不正确，可通过向同步 PR 的源分支提交轻量级 PR 完善，或使用 /close 命令关闭
7 天前 修改了pull request 的描述
7 天前 添加了label：ci-pipeline-running
7 天前 修改了pull request 的描述
7 天前 修改了pull request 的描述
ascend-robot
7 天前 评论：
7 天前 评论：
ascend docs pipeline is running...
7 天前 添加了label：docs-ci-pipeline-running
ascend-robot
7 天前 评论：
7 天前 评论：
✅ 跳过 docs ci 检查，没有需要检查的文档文件
7 天前 删除了label：docs-ci-pipeline-running
7 天前 添加了label：docs-ci-pipeline-success
7 天前 删除了label：ci-pipeline-running
7 天前 添加了label：ci-pipeline-failed
AtlasAccount
7 天前 评论：
7 天前 评论：
流水线 PR-pipeline_pytorch#65074 [ commitID：c5bbdd28 ] 运行失败
>>>代码风格自动修复执行成功（无修复内容）
| 阶段 | 任务名 | 状态 | 详情 |
|---|---|---|---|
| 编译构建 | Build_X86 | 🟣 INIT | >>> |
| Build_ARM | 🟣 INIT | >>> | |
| Build_X86_torchair | 🟣 INIT | >>> | |
| Build_ARM_torchair | 🟣 INIT | >>> | |
| patch_test | 🟣 INIT | >>> | |
| Build_X86_213 | 🟣 INIT | >>> | |
| Build_ARM_213 | 🟣 INIT | >>> | |
| 恶意代码检查 | Antipoison | ✅ COMPLETED | >>> |
| 编码安全与规范检查 | codecheck_pre-commit | ✅ COMPLETED | >>> |
| check_error | ✅ COMPLETED | >>> | |
| lintrunner | ❌ FAILED | >>> | |
| 开源片段检查 | SCA | ✅ COMPLETED | >>> |
| 开发者测试 | UT_ARM_A3_Part_01 | 🟣 INIT | >>> |
| UT_ARM_A3_Part_02 | 🟣 INIT | >>> | |
| UT_ARM_A2_Part_01 | 🟣 INIT | >>> | |
| UT_ARM_A2_Part_02 | 🟣 INIT | >>> | |
| UT_ARM_A2_Part_03 | 🟣 INIT | >>> | |
| UT_inductor_Part_01 | 🟣 INIT | >>> | |
| UT_inductor_Part_02 | 🟣 INIT | >>> | |
| UT_inductor_Part_03 | 🟣 INIT | >>> | |
| UT_inductor_Part_04 | 🟣 INIT | >>> | |
| UT_DIST_ARM_Part_01 | 🟣 INIT | >>> | |
| UT_DIST_ARM_Part_02 | 🟣 INIT | >>> | |
| UT_DIST_ARM_Part_03 | 🟣 INIT | >>> | |
| UT_DIST_ARM_Part_04 | 🟣 INIT | >>> | |
| UT_ARM_A2_Select_Part_01 | 🟣 INIT | >>> | |
| UT_ARM_A2_Select_Part_02 | 🟣 INIT | >>> | |
| UT_ARM_A2_Part_213 | 🟣 INIT | >>> | |
| UT_inductor_Part_213 | 🟣 INIT | >>> | |
| UT_DIST_ARM_Part_213 | 🟣 INIT | >>> | |
| UT_ARM_A2_Select_Part_213 | 🟣 INIT | >>> | |
| 流水线 | PR-pipeline_pytorch | ❌ FAILED | >>> |
- compile、compile_inductor、compile_torchair : 运行流水线
- retry : 重试流水线所有失败子任务
- retry <任务名> : 仅重试指定失败子任务
- stop : 停止流水线
stonexxx7 天前 update merge request[project id: 7404318, iid: 45609, commit_id: 91f897b648ff037566bf6e7a55323ad2091ad7bd] virtual merging success
7 天前 update merge request[project id: 7404318, iid: 45609, commit_id: 91f897b648ff037566bf6e7a55323ad2091ad7bd] virtual merging success
此处折叠了5条事件消息 查看更多
7 天前 添加了label：ci-pipeline-running
ascend-robot
7 天前 评论：
7 天前 评论：
ascend docs pipeline is running...
7 天前 删除了label：docs-ci-pipeline-success
7 天前 添加了label：docs-ci-pipeline-running
ascend-robot
7 天前 评论：
7 天前 评论：
✅ 跳过 docs ci 检查，没有需要检查的文档文件
7 天前 删除了label：docs-ci-pipeline-running
7 天前 添加了label：docs-ci-pipeline-success
6 天前 删除了label：ci-pipeline-running
6 天前 添加了label：ci-pipeline-passed
AtlasAccount
6 天前 评论：
6 天前 评论：
流水线 PR-pipeline_pytorch#65074 （重试第1次） [ commitID：c5bbdd28 ] 已完成
⚠️ 流水线执行结果已过期，当前 PR 最新 commit 为 4455228d，请重新运行流水线
>>>代码风格自动修复执行成功（无修复内容）
| 阶段 | 任务名 | 状态 | 详情 |
|---|---|---|---|
| 编译构建 | Build_X86 | ✅ COMPLETED | >>> |
| Build_ARM | ✅ COMPLETED | >>> | |
| Build_X86_torchair | ⚪ IGNORED | >>> | |
| Build_ARM_torchair | ⚪ IGNORED | >>> | |
| patch_test | ⚪ IGNORED | >>> | |
| Build_X86_213 | ✅ COMPLETED | >>> | |
| Build_ARM_213 | ✅ COMPLETED | >>> | |
| 恶意代码检查 | Antipoison | ✅ COMPLETED | >>> |
| 编码安全与规范检查 | codecheck_pre-commit | ✅ COMPLETED | >>> |
| check_error | ✅ COMPLETED | >>> | |
| lintrunner | ✅ COMPLETED | >>> | |
| 开源片段检查 | SCA | ✅ COMPLETED | >>> |
| 开发者测试 | UT_ARM_A3_Part_01 | ⚪ IGNORED | >>> |
| UT_ARM_A3_Part_02 | ⚪ IGNORED | >>> | |
| UT_ARM_A2_Part_01 | ✅ COMPLETED | >>> | |
| UT_ARM_A2_Part_02 | ✅ COMPLETED | >>> | |
| UT_ARM_A2_Part_03 | ✅ COMPLETED | >>> | |
| UT_inductor_Part_01 | ⚪ IGNORED | >>> | |
| UT_inductor_Part_02 | ⚪ IGNORED | >>> | |
| UT_inductor_Part_03 | ⚪ IGNORED | >>> | |
| UT_inductor_Part_04 | ⚪ IGNORED | >>> | |
| UT_DIST_ARM_Part_01 | ⚪ IGNORED | >>> | |
| UT_DIST_ARM_Part_02 | ⚪ IGNORED | >>> | |
| UT_DIST_ARM_Part_03 | ⚪ IGNORED | >>> | |
| UT_DIST_ARM_Part_04 | ⚪ IGNORED | >>> | |
| UT_ARM_A2_Select_Part_01 | ✅ COMPLETED | >>> | |
| UT_ARM_A2_Select_Part_02 | ✅ COMPLETED | >>> | |
| UT_ARM_A2_Part_213 | ✅ COMPLETED | >>> | |
| UT_inductor_Part_213 | ⚪ IGNORED | >>> | |
| UT_DIST_ARM_Part_213 | ⚪ IGNORED | >>> | |
| UT_ARM_A2_Select_Part_213 | ✅ COMPLETED | >>> | |
| 流水线 | PR-pipeline_pytorch | ✅ COMPLETED | >>> |
- compile、compile_inductor、compile_torchair : 运行流水线
- retry : 重试流水线所有失败子任务
- retry <任务名> : 仅重试指定失败子任务
- stop : 停止流水线
6 天前 添加了label：approved
6 天前 修改了pull request 的描述
6 天前 删除了label：needs-issue
rain-666
6 天前 评论：
6 天前 评论：
/lgtm
6 天前 添加了label：lgtm
6 天前 合入了pull request
AtlasAccount
6 天前 评论：
6 天前 评论：
流水线 pytorch_gitcode_PR_multiVersion#14805 [ commitID：4455228d ] 已完成
https://gitcode.com/Ascend/pytorch/issues/4496
【合入来源】
【修改方案】
【资料变更】
【接口变更】
【功能验证】
【CheckList】