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

变更摘要

此 PR 为 MoE（Mixture of Experts）token 反置换操作新增了 DTensor 分布式分片策略支持，主要针对 _npu_moe_token_unpermute 和 npu_moe_token_unpermute_grad_v2 两个算子注册了分片策略函数。

主要改动

• 新增 _npu_moe_token_unpermute_strategy 策略函数：为 npu._npu_moe_token_unpermute.default 算子注册 DTensor 分片策略，提供全复制（Replicate）和隐藏维度分片（Shard(1)）两种策略，其中隐藏维度分片策略对输出张量和输入 permuted_tokens 沿第 1 维进行分片。
• 新增 npu_moe_token_unpermute_grad_v2_strategy 策略函数：为 npu.npu_moe_token_unpermute_grad_v2.default 算子注册 DTensor 分片策略，同样支持全复制和隐藏维度分片两种策略，隐藏维度分片策略中梯度输出 permuted_tokens 沿第 1 维分片，probs 梯度使用 Partial() 归约。
• 可选张量参数处理：两个新增策略函数均对可选的 probs 和 permuted_tokens 参数进行了判空处理，当参数为 None 时对应策略位置填 None。

代码审查

经过完整的审查，我对变更文件做了以下检查：

1. _npu_moe_token_unpermute_strategy（行 101-123）：与已有的 npu_moe_token_unpermute_strategy（行 77-98）对比，新增函数多了一个输出（permuted_tokens_for_backward），两个输出的 sharding 均与输入的 permuted_tokens 保持一致，参数数量与函数注释完全匹配。

2. npu_moe_token_unpermute_grad_v2_strategy（行 151-177）：与已有的 npu_moe_token_unpermute_grad_strategy（行 126-148）对比，v2 版本将必选的 permuted_tokens 张量替换为 permuted_tokens_size_0（int）和 permuted_tokens_dtype（ScalarType），并在末尾增加了可选的 permuted_tokens。输出 sharding 语义与 v1 完全一致，参数数量（8 个）与函数注释匹配。该模式亦与文件中已有的 npu_moe_token_permute_grad_v2_strategy（行 54-74）一致。

3. 未发现逻辑错误、安全漏洞、可靠性问题或破坏性变更。

总结

• torch_npu/distributed/tensor/_moe_ops.py：已审查，发现 1 个 P3 级别问题（缺少对新增策略的测试覆盖），无其他问题。

整体风险判断：低。两个新增的分片策略函数严格遵循文件中已有的模式，分片语义正确、参数数量匹配。唯一的缺口是缺少对应的自动化测试，建议后续补充。

⚠️ 已识别出整体风险，但无法提取行内评论，请参考整体评估。

CLA Signature Pass

ascend-ds-bot, thanks for your pull request. All authors of the commits have signed the CLA. 👍

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

💡 Tip:

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

ascend docs pipeline is running...

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

• compile、compile_inductor、compile_torchair : 运行流水线
• retry : 重试流水线所有失败子任务
• retry <任务名> : 仅重试指定失败子任务
• stop : 停止流水线

/approve

/approve