diff --git a/paddle/phi/infermeta/multiary.cc b/paddle/phi/infermeta/multiary.cc index 76b58fda898562..290de597609cfd 100644 --- a/paddle/phi/infermeta/multiary.cc +++ b/paddle/phi/infermeta/multiary.cc @@ -6298,6 +6298,18 @@ void MoePermuteInferMeta(const MetaTensor& X, expert_prob_topk.dims(), expert_routemap_topk.dims())); } + if (cols != -1) { + PADDLE_ENFORCE_GT( + cols, + 0, + common::errors::InvalidArgument( + "X.dims()[1] should be positive, but got %ld.", cols)); + PADDLE_ENFORCE_LE( + cols, + static_cast(std::numeric_limits::max()), + common::errors::InvalidArgument( + "X.dims()[1] should be <= INT_MAX, but got %ld.", cols)); + } const bool check_input_shape = !common::contain_unknown_dim(X.dims()) && !common::contain_unknown_dim(expert_routemap_topk.dims()); @@ -6320,16 +6332,6 @@ void MoePermuteInferMeta(const MetaTensor& X, static_cast(std::numeric_limits::max()) - 32, common::errors::InvalidArgument( "X.dims()[0] should be <= INT_MAX - 32, but got %ld.", rows)); - PADDLE_ENFORCE_GE( - cols, - 0, - common::errors::InvalidArgument( - "X.dims()[1] should be non-negative, but got %ld.", cols)); - PADDLE_ENFORCE_LE( - cols, - static_cast(std::numeric_limits::max()), - common::errors::InvalidArgument( - "X.dims()[1] should be <= INT_MAX, but got %ld.", cols)); PADDLE_ENFORCE_GE(topk, 1, common::errors::InvalidArgument( @@ -6577,13 +6579,13 @@ 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)); + if (cols != -1) { + PADDLE_ENFORCE_GT( + cols, + 0, + common::errors::InvalidArgument( + "unzipped_tokens.dims()[1] should be positive, but got %ld.", + cols)); } if (!common::contain_unknown_dim(unzipped_tokens.dims()) && !common::contain_unknown_dim(unzipped_token_probs.dims())) { diff --git a/paddle/phi/kernels/gpu/moe_permute_kernel.cu b/paddle/phi/kernels/gpu/moe_permute_kernel.cu index 06dd8d0d52abf2..180deaa611429d 100644 --- a/paddle/phi/kernels/gpu/moe_permute_kernel.cu +++ b/paddle/phi/kernels/gpu/moe_permute_kernel.cu @@ -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)); PADDLE_ENFORCE_LE( cols, std::numeric_limits::max(), diff --git a/paddle/phi/kernels/gpu/moe_unpermute_kernel.cu b/paddle/phi/kernels/gpu/moe_unpermute_kernel.cu index c23360b006a19f..9c69f680eb6769 100644 --- a/paddle/phi/kernels/gpu/moe_unpermute_kernel.cu +++ b/paddle/phi/kernels/gpu/moe_unpermute_kernel.cu @@ -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( + cols, + 0, + common::errors::InvalidArgument( + "unzipped_tokens.dims()[1] should be positive, but got %ld.", cols)); PADDLE_ENFORCE_LE(cols, std::numeric_limits::max(), common::errors::InvalidArgument( @@ -339,7 +344,23 @@ void MoeUnpermuteKernel(const Context &dev_ctx, total_zipped_tokens_num)); dev_ctx.template Alloc(zipped_tokens); dev_ctx.template Alloc(zipped_probs_topk); - if (unzipped_tokens.numel() == 0 || total_zipped_tokens_num == 0) return; + if (unzipped_tokens.numel() == 0 || total_zipped_tokens_num == 0) { + if (zipped_tokens->numel() > 0) { + PADDLE_ENFORCE_GPU_SUCCESS( + cudaMemsetAsync(zipped_tokens->data(), + 0, + zipped_tokens->numel() * sizeof(T), + dev_ctx.stream())); + } + if (zipped_probs_topk->numel() > 0) { + PADDLE_ENFORCE_GPU_SUCCESS( + cudaMemsetAsync(zipped_probs_topk->data(), + 0, + zipped_probs_topk->numel() * sizeof(float), + dev_ctx.stream())); + } + return; + } void *zipped_probs_topk_ptr = reinterpret_cast(zipped_probs_topk->data()); const int64_t probs_numel = diff --git a/paddle/phi/kernels/xpu/moe_permute_kernel.cc b/paddle/phi/kernels/xpu/moe_permute_kernel.cc index f31d91472d92ec..e703df88933748 100644 --- a/paddle/phi/kernels/xpu/moe_permute_kernel.cc +++ b/paddle/phi/kernels/xpu/moe_permute_kernel.cc @@ -152,6 +152,11 @@ void MoePermuteKernel(const Context &dev_ctx, "moe_permute on XPU does not support override_buffer_size yet.")); const int64_t rows = X.dims()[0]; const int64_t cols = X.dims()[1]; + PADDLE_ENFORCE_GT( + cols, + 0, + common::errors::InvalidArgument( + "X.dims()[1] should be positive, received X.dims()[1]: (%ld)", cols)); PADDLE_ENFORCE_LE( rows, std::numeric_limits::max(), diff --git a/paddle/phi/kernels/xpu/moe_unpermute_kernel.cc b/paddle/phi/kernels/xpu/moe_unpermute_kernel.cc index f95409ad44bcfe..f0b09aee16be3b 100644 --- a/paddle/phi/kernels/xpu/moe_unpermute_kernel.cc +++ b/paddle/phi/kernels/xpu/moe_unpermute_kernel.cc @@ -74,6 +74,11 @@ void MoeUnpermuteKernel(const Context &dev_ctx, common::errors::Unimplemented("moe_unpermute on XPU does not support " "using_weighted_combine=true yet.")); const int64_t cols = unzipped_tokens.dims()[1]; + PADDLE_ENFORCE_GT( + cols, + 0, + common::errors::InvalidArgument( + "unzipped_tokens.dims()[1] should be positive, but got %ld.", cols)); PADDLE_ENFORCE_LE(cols, std::numeric_limits::max(), common::errors::InvalidArgument( @@ -97,7 +102,23 @@ void MoeUnpermuteKernel(const Context &dev_ctx, "topk should be less than INT_MAX, received topk: (%ld)", topk)); dev_ctx.template Alloc(zipped_tokens); dev_ctx.template Alloc(zipped_probs_topk); - if (unzipped_tokens.numel() == 0) return; // 0-size tensor + if (unzipped_tokens.numel() == 0 || total_zipped_tokens_num == 0) { + if (zipped_tokens->numel() > 0) { + PADDLE_ENFORCE_XPU_SUCCESS( + cudaMemsetAsync(zipped_tokens->data(), + 0, + zipped_tokens->numel() * sizeof(T), + reinterpret_cast(dev_ctx.stream()))); + } + if (zipped_probs_topk->numel() > 0) { + PADDLE_ENFORCE_XPU_SUCCESS( + cudaMemsetAsync(zipped_probs_topk->data(), + 0, + zipped_probs_topk->numel() * sizeof(float), + reinterpret_cast(dev_ctx.stream()))); + } + return; + } void *zipped_probs_topk_ptr = reinterpret_cast(zipped_probs_topk->data()); PADDLE_ENFORCE_XPU_SUCCESS( diff --git a/test/legacy_test/test_moe_permute_unpermute.py b/test/legacy_test/test_moe_permute_unpermute.py index 8f6dad2c8551c8..1f2204f9e130fd 100644 --- a/test/legacy_test/test_moe_permute_unpermute.py +++ b/test/legacy_test/test_moe_permute_unpermute.py @@ -666,6 +666,126 @@ def test_permute_invalid_scale_rows(self): 16, ) + def test_unpermute_all_tokens_unassigned(self): + """Test moe_unpermute when all tokens are unassigned (routemap all -1). + + This covers the case where moe_permute outputs 0-row + hidden_states_unzipped but total_zipped_tokens > 0. The unpermute + kernel must zero-initialize outputs instead of leaving them + uninitialized. + """ + seq_len = 4 + token_dim = 1024 + topk = 8 + num_experts = 32 + + hidden_states = paddle.randn([seq_len, token_dim], dtype="bfloat16") + expert_routemap_topk = paddle.full([seq_len, topk], -1, dtype="int32") + expert_prob_topk = paddle.zeros([seq_len, topk], dtype="float32") + tokens_per_expert = [0] * num_experts + + # Permute: should produce 0-row output + ( + hidden_unzipped, + compact_routemap, + prob_unzipped, + _, + ) = moe_permute( + hidden_states, + None, + expert_routemap_topk, + expert_prob_topk, + num_experts, + tokens_per_expert, + padding_alignment=128, + do_gather=True, + ) + self.assertEqual( + hidden_unzipped.shape[0], + 0, + "Expected 0 rows when all tokens are unassigned", + ) + + # Unpermute: should not produce NaN / uninitialized values + zipped_tokens, zipped_probs = moe_unpermute( + hidden_unzipped, + compact_routemap, + expert_routemap_topk, + prob_unzipped, + total_zipped_tokens=seq_len, + num_experts=num_experts, + ) + + self.assertEqual(zipped_tokens.shape, [seq_len, token_dim]) + self.assertEqual(zipped_probs.shape, [seq_len, topk]) + + # All outputs should be zero (no token was routed) + tokens_f32 = zipped_tokens.astype("float32").numpy() + probs_f32 = zipped_probs.numpy() + self.assertFalse( + np.isnan(tokens_f32).any(), + "zipped_tokens contains NaN (uninitialized memory)", + ) + self.assertFalse( + np.isnan(probs_f32).any(), + "zipped_probs contains NaN (uninitialized memory)", + ) + np.testing.assert_array_equal( + tokens_f32, + np.zeros_like(tokens_f32), + err_msg="zipped_tokens should be all zeros when no token is routed", + ) + np.testing.assert_array_equal( + probs_f32, + np.zeros_like(probs_f32), + err_msg="zipped_probs should be all zeros when no token is routed", + ) + + def test_permute_reject_zero_token_dim(self): + """Test that moe_permute rejects token_dim == 0 input.""" + seq_len = 4 + topk = 8 + num_experts = 3 + + hidden_states = paddle.empty([seq_len, 0], dtype="bfloat16") + expert_routemap_topk = paddle.full([seq_len, topk], -1, dtype="int32") + expert_prob_topk = paddle.zeros([seq_len, topk], dtype="float32") + tokens_per_expert = [0] * num_experts + + with self.assertRaisesRegex(ValueError, "positive"): + moe_permute( + hidden_states, + None, + expert_routemap_topk, + expert_prob_topk, + num_experts, + tokens_per_expert, + padding_alignment=128, + ) + + def test_unpermute_reject_zero_token_dim(self): + """Test that moe_unpermute rejects token_dim == 0 input.""" + seq_len = 3 + num_experts = 3 + topk = 8 + + unzipped_tokens = paddle.empty([4, 0], dtype="bfloat16") + zipped_expertwise_rowmap = paddle.full( + [seq_len, num_experts], -1, dtype="int32" + ) + expert_routemap_topk = paddle.full([seq_len, topk], -1, dtype="int32") + token_prob_unzipped = paddle.ones([4], dtype="float32") + + with self.assertRaisesRegex(ValueError, "positive"): + moe_unpermute( + unzipped_tokens, + zipped_expertwise_rowmap, + expert_routemap_topk, + token_prob_unzipped, + seq_len, + num_experts, + ) + if __name__ == "__main__": unittest.main()