Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wanghuancoder merged 6 commits into from
Jul 16, 2026
Merged

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

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
f295ec9
🐛 Reject token_dim == 0 in moe_permute/moe_unpermute
cangtianhuang Jul 14, 2026
ee9a558
🎨 Pre commit
cangtianhuang Jul 14, 2026
1c55748
Merge remote-tracking branch 'upstream/develop' into fix/moe-permute-…
cangtianhuang Jul 15, 2026
432f1c1
🎨 Move code to avoid contain_unknown_dim jump
cangtianhuang Jul 15, 2026
281fde8
🎨 Fix early return, add tests
cangtianhuang Jul 15, 2026
b3e74ff
🐛 Add check in xpu
cangtianhuang Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 8 additions & 8 deletions paddle/phi/infermeta/multiary.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6320,11 +6320,11 @@ void MoePermuteInferMeta(const MetaTensor& X,
static_cast<int64_t>(std::numeric_limits<int32_t>::max()) - 32,
common::errors::InvalidArgument(
"X.dims()[0] should be <= INT_MAX - 32, but got %ld.", rows));
PADDLE_ENFORCE_GE(
PADDLE_ENFORCE_GT(
Comment thread
cangtianhuang marked this conversation as resolved.
Outdated
Comment thread
cangtianhuang marked this conversation as resolved.
Outdated
cols,
0,
common::errors::InvalidArgument(
"X.dims()[1] should be non-negative, but got %ld.", cols));
"X.dims()[1] should be positive, but got %ld.", cols));
Comment thread
cangtianhuang marked this conversation as resolved.
Outdated
PADDLE_ENFORCE_LE(
cols,
static_cast<int64_t>(std::numeric_limits<int32_t>::max()),
Expand Down Expand Up @@ -6578,12 +6578,12 @@ void MoeUnpermuteInferMeta(const MetaTensor& unzipped_tokens,
const int64_t cols = unzipped_tokens.dims()[1];
const int64_t topk = expert_routemap_topk.dims()[1];
if (!common::contain_unknown_dim(unzipped_tokens.dims())) {
PADDLE_ENFORCE_GE(cols,
0,
common::errors::InvalidArgument(
"unzipped_tokens.dims()[1] should be non-negative, "
"but got %ld.",
cols));
PADDLE_ENFORCE_GT(
Comment thread
cangtianhuang marked this conversation as resolved.
cols,
0,
common::errors::InvalidArgument(
"unzipped_tokens.dims()[1] should be positive, but got %ld.",
cols));
Comment thread
cangtianhuang marked this conversation as resolved.
}
if (!common::contain_unknown_dim(unzipped_tokens.dims()) &&
!common::contain_unknown_dim(unzipped_token_probs.dims())) {
Expand Down
5 changes: 5 additions & 0 deletions paddle/phi/kernels/gpu/moe_permute_kernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -742,6 +742,11 @@ void MoePermuteKernel(const Context &dev_ctx,
"X.dims()[0] should be <= INT_MAX - %d, received: (%ld)",
kPermuteBlockSize,
rows));
PADDLE_ENFORCE_GT(
cols,
0,
common::errors::InvalidArgument(
"X.dims()[1] should be positive, received: (%ld)", cols));
Comment thread
cangtianhuang marked this conversation as resolved.
PADDLE_ENFORCE_LE(
cols,
std::numeric_limits<int32_t>::max(),
Expand Down
5 changes: 5 additions & 0 deletions paddle/phi/kernels/gpu/moe_unpermute_kernel.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -315,6 +315,11 @@ void MoeUnpermuteKernel(const Context &dev_ctx,
unzipped_token_probs.numel(),
unzipped_tokens.dims()[0]));
const int64_t cols = unzipped_tokens.dims()[1];
PADDLE_ENFORCE_GT(
Comment thread
cangtianhuang marked this conversation as resolved.
cols,
0,
common::errors::InvalidArgument(
"unzipped_tokens.dims()[1] should be positive, but got %ld.", cols));
Comment thread
cangtianhuang marked this conversation as resolved.
PADDLE_ENFORCE_LE(cols,
std::numeric_limits<int32_t>::max(),
common::errors::InvalidArgument(
Expand Down
Loading