已合并
fix sharding strategy for npu_fusion_attention #27780
fix sharding strategy for npu_fusion_attention #27780
已合并
jizewei创建于 2025年12月10日
jizewei成员
2025年12月10日

What type of PR is this?

Uncomment only one /kind <> line, hit enter to put that in a new line, and remove leading whitespaces from that line:

/kind bug
/kind task
/kind feature

/kind bug

What does this PR do / why do we need it:

fix npu_fusion_attention strategy:

  1. https://github.com/pytorch/pytorch/pull/168249 supports register_sharding for operation with Tensor kwargs. Apply a patch for previous pytorch version.
  2. The previous strategy followed query's placements, which may lead to inconsistent results compared single-device computation if sharing the dim of S/D. The new strategy focuses on SDPA scenario and supports DP/TP.
  3. For TP, the head dim of qkv local tensors does not match the head_num arg, so head_num needs to be reset before local computation.
  4. For npu_fusion_attention_backward, redistribution handling of DTensors in kwargs is added(official pytorch ops rarely have Tensor in kwargs, so most DTensor processing only considers args and ignores kwargs).

Which issue(s) this PR fixes:

Fixes #

Special notes for your reviewers:

likedislike
Pull Request已成功合入, 合并人@ascend-robot
(感谢 jizewei 的贡献)
Jjizewei成员
2025年12月10日 创建了 pull request,commit 7df1c8a4
ascend-robot
ascend-robot成员
2025年12月10日 评论:

Thank your for your pull-request.

The full list of commands accepted by me can be found at here.

You can get sig-info at here

likedislike
AtlasAccount
AtlasAccount成员
2025年12月10日 评论:
ascend-robotascend-robot成员
2025年12月10日 添加了label:stat/needs-squash
ascend-robot
ascend-robot成员
2025年12月10日 评论:

以下是根据您提交的修改文件推荐的Reviewer和Committer序列,需各模块评审通过后方可合入

Module List Reviewers Committers
repo-Ascend/pytorch chenhao_1209, qinchuanyu, guoguanghao, yuhaiyan, zhaozhijie shaoyf, wasd1111222, wangqiang160, chenhao_1209, yanpengquan
likedislike
此处折叠了69条消息 查看更多
ascend-robotascend-robot成员
2025年12月15日 添加了label:MindSpore-Component
王超成员
2025年12月15日 评论:

/approve

likedislike
ascend-robotascend-robot成员
2025年12月15日 添加了label:approved
ascend-robot
ascend-robot成员
2025年12月15日 评论:

Review Guide

This Pull-Request Passes Review.
Committers who writed a comment of /approve are: zqwenn, wangchao430.
Reviewers who writed a comment of /lgtm are: wangchao430, zqwenn.

likedislike
ascend-robotascend-robot成员
2025年12月15日 合入了pull request