[Operator Mechanism] Reject token_dim == 0 in moe_permute / moe_unpermute

[cangtianhuang]

PR Category

User Experience

PR Types

Bug fixes

Description

paddle.nn.functional.moe_permute / moe_unpermute 在 PaddleAPITest 0-size 覆盖测试中暴露出两类出错 case：

• moe_permute(Tensor([10022, 0], "float8_e4m3fn"), ...)
  → zipped_expertwise_rowmap 输出全 0，与 API 文档说明的 "-1 = 未分配" 不一致
• moe_unpermute(Tensor([12288, 0], "bfloat16"), ...)
  → zipped_probs_topk 输出全 NaN，未初始化被 nan/inf check 拦下

(PreconditionNotMet) There are NAN or INF (num_nan=40764, num_inf=0,
 num_zero=0) in [device=gpu:0, op=moe_unpermute, tensor=, dtype=fp32].

原因是两个 kernel 都有 if (X.numel() == 0) return; / if (unzipped_tokens.numel() == 0 || total_zipped_tokens_num == 0) return; 的早退分支，但 rowmap / probs / scale 等输出的初始化 cudaMemsetAsync 都放在早退之后，导致 token_dim == 0 时输出未被初始化，值不确定。

由于 hidden_states.shape[1]（token_dim，即 hidden size）在真实 MoE 场景中始终为正整数（典型值 4096 / 7168 / 8192）。token_dim == 0 意味着每个 token 没有任何特征，路由 / 加权 / 聚合等都没有数学意义。

因此本 PR 意图收紧参数约束， moe_permute / moe_unpermute 不再接受 token_dim == 0 的输入。

修改：

• paddle/phi/infermeta/multiary.cc
  • MoePermuteInferMeta：X.dims()[1] 的检查 >= 0 → > 0
  • MoeUnpermuteInferMeta：unzipped_tokens.dims()[1] 同上
• paddle/phi/kernels/gpu/moe_permute_kernel.cu
  • MoePermuteKernel 入口新增 PADDLE_ENFORCE_GT(cols, 0, ...)
• paddle/phi/kernels/gpu/moe_unpermute_kernel.cu
  • MoeUnpermuteKernel 入口新增 PADDLE_ENFORCE_GT(cols, 0, ...)

是否引起精度变化

否

[liuhao2638]

复查新提交后，仍有会阻塞合入的 CI 问题，具体位置已放在行级评论中。

优先级：P2

• 非行级：这次收紧的是 token_dim == 0 的回归行为，但 PR 还没有增加覆盖该输入的单测。建议在 test/legacy_test/test_moe_permute_unpermute.py 的 invalid input 用例中补充 moe_permute 的 [seq_len, 0] hidden_states 和 moe_unpermute 的 [unzipped_rows, 0] hidden_states_unzipped，断言抛出包含 dims()[1] 或 positive 的错误。
  处理要求：请针对该评论进行回复（同意并已修改请回复 Done，不同意请说明理由）。

^{Powered by Nyanpasu with gpt-5.5 xhigh, please check the suggestions carefully.}

[PaddlePaddle-bot]

🤖 Paddle-CI-Agent | pr_review | 2026-07-14 16:28:50

📋 Review 摘要

PR 概述：收紧 moe_permute / moe_unpermute 的输入约束，拒绝 token_dim == 0。
变更范围：paddle/phi/infermeta/multiary.cc、GPU moe_permute / moe_unpermute kernel。
影响面 Tag：[Operator Mechanism] [User Experience]

问题

级别	文件	概述
🟡 建议	paddle/phi/infermeta/multiary.cc:6323	新增 token_dim == 0 拒绝逻辑缺少对应回归测试

📝 PR 规范检查

标题 Tag [PHI] 不在 Paddle 模板枚举内，且本 PR 主要修改 Phi InferMeta / Kernel，建议使用 [Operator Mechanism]。

标题建议（可直接复制）：

• [Operator Mechanism] Reject token_dim == 0 in moe_permute / moe_unpermute

PR 描述建议（点击展开，可直接复制）

### PR Category
<!-- One of [ User Experience | Execute Infrastructure | Operator Mechanism | CINN | Custom Device | Performance Optimization | Distributed Strategy | Parameter Server | Communication Library | Auto Parallel | Inference | Environment Adaptation ] -->
Operator Mechanism

### PR Types
<!-- One of [ New features | Bug fixes | Improvements | Performance | BC Breaking | Deprecations | Docs | Devs | Not User Facing | Security | Others ] -->
Bug fixes

### Description
<!-- Describe what you’ve done -->
`paddle.nn.functional.moe_permute` / `moe_unpermute` 不再接受 `token_dim == 0` 输入。该输入在 MoE 场景中没有有效特征维度，且此前会触发 0-size 早退导致 `zipped_expertwise_rowmap` / `zipped_probs_topk` 等输出未初始化或出现 NaN。

修改内容：
- `MoePermuteInferMeta` / `MoeUnpermuteInferMeta` 将第 2 维校验从 non-negative 收紧为 positive。
- GPU `MoePermuteKernel` / `MoeUnpermuteKernel` 在运行时增加 `cols > 0` 校验，保证下游早退和 `static_cast<int>` 前拒绝 0 列输入。

验证方式：
- N/A

### 是否引起精度变化
<!-- one of the following [ 是 | 否 ]-->
否

总体评价

新增约束放在 GPU 下游 static_cast<int> 和 0-size 早退之前，InferMeta 与 GPU kernel 的核心逻辑方向一致。建议补上 [N, 0] / [rows, 0] 的异常回归测试，并按 Paddle 官方 Tag 修正 PR 标题。

[wanghuancoder]

LGTM

[A-nnonymous]

LGTM in shrink col-wise acceptance

[cangtianhuang]

/re-run all-failed

[codecov-commenter]

Codecov Report

❌ Patch coverage is 0% with 2 lines in your changes missing coverage. Please review.
⚠️ Please upload report for BASE (develop@b6fb7ea). Learn more about missing BASE report.

Files with missing lines	Patch %	Lines
paddle/phi/infermeta/multiary.cc	0.00%	2 Missing ⚠️

❌ Your patch status has failed because the patch coverage (0.00%) is below the target coverage (90.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             develop   #79470   +/-   ##
==========================================
  Coverage           ?    0.00%           
==========================================
  Files              ?        1           
  Lines              ?        2           
  Branches           ?        0           
==========================================
  Hits               ?        0           
  Misses             ?        2           
  Partials           ?        0           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[paddle-bot]

你的PR提交成功，感谢你对开源项目的贡献!
请关注后续CI自动化测试结果，详情请参考Paddle-CI手册。
Your PR has been submitted. Thanks for your contribution!
Please wait for the result of CI firstly. See Paddle CI Manual for details.

[sneaxiy]

LGTM