已合并

[fa_v3]skip cpu check, register_sharding for npu_fusion_attention_v3 #30861

王超创建于 2月12日从wangchao430/pytorch： v2.10.0_fav3 合入到 Ascend/pytorch： v2.10.0

[fa_v3]skip cpu check, register_sharding for npu_fusion_attention_v3 #30861

已合并

王超创建于 2月12日

从wangchao430/pytorch v2.10.0_fav3合入到Ascend/pytorch v2.10.0

王

王超成员

2月12日

【合入来源】

如有社区issue，请关联issue链接
请勿携带内部流程信息（需求链接、问题单、内部issue等）

需求
问题单
issue/工单
重构优化
资料更新

【修改方案】

torch_npu.npu_fusion_attention_v3适配支持DTensor策略注册。
graph_tree中支持根据torch_npu._inductor.config.npugraph_trees.disable_cpu_input_check配置开启时跳过cpu Tensor的相关检测，从而支持torch_npu.npu_fusion_attention_v3入graph_tree

【资料变更】

不涉及

【接口变更】

不涉及

【功能验证】

新增DTensor测试用例

【CheckList】

PR提交人对以下CheckList自检项进行全量自检，自检通过或不涉及，均修改 [ ] 为 [x]

代码注释完备，正确记录错误日志
代码实现进行了返回值、空指针等校验
PR标题正确使用类型标签，如：feat、fix、refactor、docs、test等
PR持续集成流水线（CI）执行通过，代码检查无异常

Pull Request已成功合入, 合并人@ascend-robot

（感谢王超的贡献）

王王超成员

2月12日创建了 pull request，commit 47bbf549

AtlasAccount成员

2月12日评论：

openLiBing可接管PR, PR详情链接:
https://www.openlibing.com/apps/transitionPrDetail?owner=Ascend&repo=pytorch&number=30861&platform=gitcode&codeHostingPlatformFlag=gitcode

ascend-robot成员

2月12日添加了label：stat/needs-squash

ascend-robot成员

2月12日评论：

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

PR Approval Progress

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

Module Approval Details

module	lgtm status	approve status
repo-Ascend/pytorch	✅ zqwenn, liwei386 (2/2)	✅ liwei386, zqwenn (2/1)
test	✅ liwei386, zqwenn (2/2)	✅ zqwenn, liwei386 (2/1)
torch_npu/_inductor	✅ zqwenn, liwei386 (2/2)	✅ liwei386, zqwenn (2/1)
torch_npu/distributed	✅ zqwenn, liwei386 (2/2)	✅ zqwenn, liwei386 (2/1)
torch_npu/npu	✅ liwei386, zqwenn (2/2)	✅ liwei386, zqwenn (2/1)
torch_npu/utils	✅ liwei386, zqwenn (2/2)	✅ zqwenn, liwei386 (2/1)

💡 Tip:

Committer can comment /approve or /lgtm

Commenting /approve implies both code review (lgtm) and intent to merge (approve)

ascend-robot成员

2月12日添加了label：ascend-cla/yes

王

王超成员

2月12日评论：

compile

ascend-robot成员

2月12日添加了label：ci-pipeline-running

AtlasAccount成员2月12日进行代码检视1

torch_npu/_inductor/config.py

6685# (1) 'all' means try to fallback all kernel to fx graph call.

6786# (2) [1, 2, 10] means try to fallback kernel like triton_xxx_1, triton_xxx_2 and triton_xxx_10

6887force_fallback_kernel_id = []

6988

AtlasAccount2月12日评论：

代码逻辑和结构: 注释与代码逻辑不一致。注释中提到'set slow_path_cudagraph_asserts to True'，但实际代码设置为False。这会导致理解上的混淆，特别是当其他开发者阅读注释时，会误以为设置为True来跳过CPU检查，但实际上设置为False。需要确保注释准确反映代码意图。

问题类型: 代码逻辑和结构
文件路径: torch_npu/_inductor/config.py
行号: 66
问题代码:

        # When disable_cpu_input_check is True, set slow_path_cudagraph_asserts to True to skip the CPU check. 
        if value:
            torch._inductor.config.triton.slow_path_cudagraph_asserts = False

修改建议:

修正注释以准确描述代码行为。如果意图确实是设置为False，注释应说明原因。例如：'When disable_cpu_input_check is True, set slow_path_cudagraph_asserts to False to skip the CPU check.' 或者根据实际需求调整代码逻辑。

此评论由代码审查工具自动生成

AtlasAccount成员2月12日进行代码检视2

torch_npu/distributed/tensor/_attention.py

已过期

665+ )

666+ elif op_call == npu.npu_fusion_attention_grad.default or op_call == npu.npu_fusion_attention_grad_v3.default:

401667 local_kwargs = get_redistributed_local_kwargs(

402668 _infer_npu_fusion_attention_grad_kwargs_spec, op_info, output_sharding

AtlasAccount2月12日评论：

语法错误: 第304行存在与第303行相同的语法错误，'None, if actual_seq_kvlen is None else Replicate()' 中的逗号位置不正确。

问题类型: 语法错误
文件路径: torch_npu/distributed/tensor/_attention.py
行号: 304
问题代码:

            None, if actual_seq_kvlen is None else Replicate(),   # actual_seq_kvlen

修改建议:

将第304行修改为：'None if actual_seq_kvlen is None else Replicate(),   # actual_seq_kvlen'

此评论由代码审查工具自动生成

系统消息系统

2月12日评论：

changed this line on aa64687a view diff detail

AtlasAccount成员2月12日进行代码检视2

torch_npu/distributed/tensor/_attention.py

已过期

446718}

447719

448720old_handlers = DTensor._op_dispatcher._custom_op_handlers

AtlasAccount2月12日评论：

语法错误: 第347行存在与第303行相同的语法错误，'None, if actual_seq_qlen is None else Replicate()' 中的逗号位置不正确。

问题类型: 语法错误
文件路径: torch_npu/distributed/tensor/_attention.py
行号: 347
问题代码:

                None, if actual_seq_qlen is None else Replicate(),    # actual_seq_qlen

修改建议:

将第347行修改为：'None if actual_seq_qlen is None else Replicate(),    # actual_seq_qlen'

此评论由代码审查工具自动生成

系统消息系统

2月12日评论：

changed this line on aa64687a view diff detail

AtlasAccount成员2月12日进行代码检视2

torch_npu/distributed/tensor/_attention.py

已过期

446718}

447719

448720old_handlers = DTensor._op_dispatcher._custom_op_handlers

AtlasAccount2月12日评论：

语法错误: 第348行存在与第304行相同的语法错误，'None, if actual_seq_kvlen is None else Replicate()' 中的逗号位置不正确。

问题类型: 语法错误
文件路径: torch_npu/distributed/tensor/_attention.py
行号: 348
问题代码:

                None, if actual_seq_kvlen is None else Replicate(),   # actual_seq_kvlen

修改建议:

将第348行修改为：'None if actual_seq_kvlen is None else Replicate(),   # actual_seq_kvlen'

此评论由代码审查工具自动生成

系统消息系统

2月12日评论：

changed this line on aa64687a view diff detail

AtlasAccount成员2月12日进行代码检视2

torch_npu/distributed/tensor/_attention.py

已过期

491752 None # sink

492753 ]

493754 )

494755 strategies.append(dp_sharding_strategy)

AtlasAccount2月12日评论：

语法错误: 第383行存在与第303行相同的语法错误，'None, if actual_seq_qlen is None else Replicate()' 中的逗号位置不正确。

问题类型: 语法错误
文件路径: torch_npu/distributed/tensor/_attention.py
行号: 383
问题代码:

                None, if actual_seq_qlen is None else Replicate(),    # actual_seq_qlen

修改建议:

将第383行修改为：'None if actual_seq_qlen is None else Replicate(),    # actual_seq_qlen'

此评论由代码审查工具自动生成

系统消息系统

2月12日评论：

changed this line on aa64687a view diff detail

AtlasAccount成员2月12日进行代码检视2

torch_npu/distributed/tensor/_attention.py

已过期

492753 ]

493754 )

494755 strategies.append(dp_sharding_strategy)

495756

AtlasAccount2月12日评论：

语法错误: 第384行存在与第304行相同的语法错误，'None, if actual_seq_kvlen is None else Replicate()' 中的逗号位置不正确。

问题类型: 语法错误
文件路径: torch_npu/distributed/tensor/_attention.py
行号: 384
问题代码:

                None, if actual_seq_kvlen is None else Replicate(),   # actual_seq_kvlen

修改建议:

将第384行修改为：'None if actual_seq_kvlen is None else Replicate(),   # actual_seq_kvlen'

此评论由代码审查工具自动生成

系统消息系统

2月12日评论：

changed this line on aa64687a view diff detail

AtlasAccount成员2月12日进行代码检视2

torch_npu/distributed/tensor/_attention.py

已过期

447719}

448720

449721old_handlers = DTensor._op_dispatcher._custom_op_handlers

AtlasAccount2月12日评论：

语法错误: 第427行存在语法错误，'None, if seed is None else Replicate()' 中的逗号位置不正确。

问题类型: 语法错误
文件路径: torch_npu/distributed/tensor/_attention.py
行号: 427
问题代码:

            None, if seed is None else Replicate(),        # seed

修改建议:

将第427行修改为：'None if seed is None else Replicate(),        # seed'

此评论由代码审查工具自动生成

系统消息系统

2月12日评论：

changed this line on aa64687a view diff detail

AtlasAccount成员2月12日进行代码检视2

torch_npu/distributed/tensor/_attention.py

已过期

447719}

448720

449721old_handlers = DTensor._op_dispatcher._custom_op_handlers

AtlasAccount2月12日评论：

语法错误: 第428行存在语法错误，'None, if offset is None else Replicate()' 中的逗号位置不正确。

问题类型: 语法错误
文件路径: torch_npu/distributed/tensor/_attention.py
行号: 428
问题代码:

            None, if offset is None else Replicate(),      # offset

修改建议:

将第428行修改为：'None if offset is None else Replicate(),      # offset'

此评论由代码审查工具自动生成

系统消息系统

2月12日评论：

changed this line on aa64687a view diff detail

AtlasAccount成员2月12日进行代码检视2

torch_npu/distributed/tensor/_attention.py

已过期

447719}

448720

449721old_handlers = DTensor._op_dispatcher._custom_op_handlers

AtlasAccount2月12日评论：

语法错误: 第430行存在语法错误，'None, if actual_seq_qlen is None else Replicate()' 中的逗号位置不正确。

问题类型: 语法错误
文件路径: torch_npu/distributed/tensor/_attention.py
行号: 430
问题代码:

            None, if actual_seq_qlen is None else Replicate(),  # actual_seq_qlen

修改建议:

将第430行修改为：'None if actual_seq_qlen is None else Replicate(),  # actual_seq_qlen'

此评论由代码审查工具自动生成

系统消息系统

2月12日评论：

changed this line on aa64687a view diff detail

AtlasAccount成员2月12日进行代码检视2

torch_npu/distributed/tensor/_attention.py

已过期

447719}

448720

449721old_handlers = DTensor._op_dispatcher._custom_op_handlers

AtlasAccount2月12日评论：

语法错误: 第431行存在语法错误，'None, if actual_seq_kvlen is None else Replicate()' 中的逗号位置不正确。

问题类型: 语法错误
文件路径: torch_npu/distributed/tensor/_attention.py
行号: 431
问题代码:

            None, if actual_seq_kvlen is None else Replicate(), # actual_seq_kvlen

修改建议:

将第431行修改为：'None if actual_seq_kvlen is None else Replicate(), # actual_seq_kvlen'

此评论由代码审查工具自动生成

系统消息系统

2月12日评论：

changed this line on aa64687a view diff detail

AtlasAccount成员

2月12日评论：

问题/功能描述

本次PR旨在为npu_fusion_attention_v3算子提供完整的分布式训练支持与图优化增强。主要解决了两个核心问题：一是现有分布式张量（DTensor）框架缺少对该新版融合注意力算子前向与反向传播的分片策略支持，导致其无法在分布式环境中正确运行；二是现有NPU图优化流程对CPU输入节点的检查过于严格，在某些特定场景下会错误地跳过图优化，影响性能与功能适用性。

修改方案描述

修改方案包含三个核心部分：首先，在分布式策略注册模块中新增了npu_fusion_attention_v3_strategy和npu_fusion_attention_grad_v3_strategy两个分片策略函数，并更新统一处理函数以支持V3算子，确保其在数据并行与模型并行场景下的正确分片计算。其次，在配置系统中新增了npugraph_trees.disable_cpu_input_check开关，并在图树管理、设备节点检查等多个关键模块中集成此配置，当启用时允许跳过对CPU输入节点的严格检查，从而避免图优化被不当中断。最后，为增强调试能力，注册了新的aclgraph日志类别，并将相关日志记录器升级为更灵活的标准logger，同时优化了日志格式化方式。此外，补充了针对新算子的分布式测试用例，以验证其在不同输入布局和稀疏模式下的正确性。

ascend-robot成员

2月12日删除了label：ci-pipeline-running

ascend-robot成员

2月12日添加了label：ci-pipeline-failed

ascend-robot成员

2月12日评论：

流水线 PR-pipeline_pytorch#7273 运行失败

阶段	任务名	状态	详情
编译构建	Build_X86	🕚	>>>
	Build_LibTorch	🕚	>>>
	Build_ARM	🕚	>>>
	Build_ARM_inductor	🕚	>>>
	Build_X86_py311	🕚	>>>
	Build_ARM_py311	🕚	>>>
恶意代码检查	Antipoison	✅	>>>
编码安全与规范检查	CodeCheck	❌	>>>
编码安全与规范检查	check_error	✅	>>>
开源片段检查	SCA	✅	>>>
开发者测试	UT_DIST_X86	🕚	>>>
	UT_X86_Part_01	🕚	>>>
	UT_X86_Part_02	🕚	>>>
	UT_ARM_A2_Part_01	🕚	>>>
	UT_ARM_A2_Part_02	🕚	>>>
	UT_inductor_Part_01	🕚	>>>
	UT_inductor_Part_02	🕚	>>>
	UT_inductor_Part_03	🕚	>>>
	UT_inductor_Part_04	🕚	>>>
流水线	PR-pipeline_pytorch	❌	>>>

王王超成员

2月12日 update merge request[project id: 7404318, iid: 30861, commit_id: 33b481cc5001559e47d62e6d30016a0de1ea3c56] virtual merging success

王王超成员

2月12日推送 1 个提交:aa64687a-fix with strategy register

王王超成员

2月12日 update merge request[project id: 7404318, iid: 30861, commit_id: 9948d29de46b72a1ed19bd0d2de04942877b2dc6] virtual merging success

ascend-robot成员

2月12日评论：

CLA Signature Pass

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

AtlasAccount成员

2月12日删除了label：ci-pipeline-failed

AtlasAccount成员

2月12日添加了label：ci-pipeline-failed

王

王超成员

2月12日评论：

compile

ascend-robot成员

2月12日删除了label：ci-pipeline-failed

ascend-robot成员

2月12日添加了label：ci-pipeline-running

ascend-robot成员

2月12日删除了label：ci-pipeline-running

ascend-robot成员

2月12日添加了label：ci-pipeline-passed

ascend-robot成员

2月12日评论：

流水线 PR-pipeline_pytorch#7279 已完成

阶段	任务名	状态	详情
编译构建	Build_X86	✅	>>>
	Build_LibTorch	✅	>>>
	Build_ARM	✅	>>>
	Build_ARM_inductor	✅	>>>
	Build_X86_py311	🛑	>>>
	Build_ARM_py311	🛑	>>>
恶意代码检查	Antipoison	✅	>>>
编码安全与规范检查	CodeCheck	✅	>>>
编码安全与规范检查	check_error	✅	>>>
开源片段检查	SCA	✅	>>>
开发者测试	UT_DIST_X86	✅	>>>
	UT_X86_Part_01	✅	>>>
	UT_X86_Part_02	✅	>>>
	UT_ARM_A2_Part_01	🛑	>>>
	UT_ARM_A2_Part_02	🛑	>>>
	UT_inductor_Part_01	✅	>>>
	UT_inductor_Part_02	✅	>>>
	UT_inductor_Part_03	✅	>>>
	UT_inductor_Part_04	✅	>>>
流水线	PR-pipeline_pytorch	✅	>>>

zhangqiongwen成员

2月14日评论：

/approve

ascend-robot成员

2月14日添加了label：approved

liwei386成员

2月14日评论：

/approve

ascend-robot成员

2月14日添加了label：lgtm

ascend-robot成员

2月14日评论：

Review Guide

This pull-request passes review.
Committers who wrote a comment of /approve are: liwei386, zqwenn.
Reviewers who wrote a comment of /lgtm are: liwei386, zqwenn.

ascend-robot成员

2月14日合入了pull request