From a96c609074928e689fab4d1473cdab423995d82b Mon Sep 17 00:00:00 2001 From: Id545 Date: Sun, 20 Sep 2026 21:01:25 +0200 Subject: [PATCH 1/2] [Bugfix][MoE] Respect padded row strides in top-k routing Propagate the router-logit row stride through the native top-k kernels and use the stride-aware unfused path for padded rows. Add regression coverage for softmax and sigmoid routing with padded power-of-two and non-power-of-two expert counts. Assisted-by: OpenAI Signed-off-by: Id545 --- .../moe/topk_softmax_kernels.cu | 77 ++++++++++++++----- tests/kernels/moe/test_fused_topk.py | 54 +++++++++++++ 2 files changed, 110 insertions(+), 21 deletions(-) diff --git a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu index 098644a1b901..af0f7df62879 100644 --- a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu +++ b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu @@ -78,7 +78,8 @@ enum ScoringFunc { // in the softmax kernel when we extend this module to support expert-choice routing. template __launch_bounds__(TPB) __global__ - void moeSoftmax(const InputType* input, const bool* finished, float* output, const int num_cols) + void moeSoftmax(const InputType* input, const bool* finished, float* output, const int num_cols, + const int64_t input_row_stride) { using BlockReduce = cub::BlockReduce; __shared__ typename BlockReduce::TempStorage tmpStorage; @@ -86,7 +87,8 @@ __launch_bounds__(TPB) __global__ __shared__ float normalizing_factor; __shared__ float float_max; - const int thread_row_offset = blockIdx.x * num_cols; + const int64_t input_row_offset = blockIdx.x * input_row_stride; + const int output_row_offset = blockIdx.x * num_cols; float threadData(-FLT_MAX); @@ -98,8 +100,7 @@ __launch_bounds__(TPB) __global__ for (int ii = threadIdx.x; ii < num_cols; ii += TPB) { - const int idx = thread_row_offset + ii; - const float val = toFloat(input[idx]); + const float val = toFloat(input[input_row_offset + ii]); threadData = max(val, threadData); } @@ -114,8 +115,7 @@ __launch_bounds__(TPB) __global__ for (int ii = threadIdx.x; ii < num_cols; ii += TPB) { - const int idx = thread_row_offset + ii; - const float val = toFloat(input[idx]); + const float val = toFloat(input[input_row_offset + ii]); threadData += expf(val - float_max); } @@ -129,20 +129,21 @@ __launch_bounds__(TPB) __global__ for (int ii = threadIdx.x; ii < num_cols; ii += TPB) { - const int idx = thread_row_offset + ii; - const float val = toFloat(input[idx]); + const float val = toFloat(input[input_row_offset + ii]); float softmax_val = expf(val - float_max) * normalizing_factor; // Clamp NaN/Inf to 0 to prevent duplicate expert IDs downstream. if (isnan(softmax_val) || isinf(softmax_val)) softmax_val = 0.f; - output[idx] = softmax_val; + output[output_row_offset + ii] = softmax_val; } } template __launch_bounds__(TPB) __global__ - void moeSigmoid(const InputType* input, const bool* finished, float* output, const int num_cols) + void moeSigmoid(const InputType* input, const bool* finished, float* output, const int num_cols, + const int64_t input_row_stride) { - const int thread_row_offset = blockIdx.x * num_cols; + const int64_t input_row_offset = blockIdx.x * input_row_stride; + const int output_row_offset = blockIdx.x * num_cols; // Don't touch finished rows. if ((finished != nullptr) && finished[blockIdx.x]) @@ -152,12 +153,11 @@ __launch_bounds__(TPB) __global__ for (int ii = threadIdx.x; ii < num_cols; ii += TPB) { - const int idx = thread_row_offset + ii; - const float val = toFloat(input[idx]); + const float val = toFloat(input[input_row_offset + ii]); float sigmoid_val = 1.0f / (1.0f + __expf(-val)); // Clamp NaN/Inf to 0 to prevent duplicate expert IDs downstream. if (isnan(sigmoid_val) || isinf(sigmoid_val)) sigmoid_val = 0.f; - output[idx] = sigmoid_val; + output[output_row_offset + ii] = sigmoid_val; } } @@ -662,6 +662,7 @@ void topkGatingKernelLauncher( float* workspace, const int num_tokens, const int num_experts, + const int64_t input_row_stride, const int topk, const bool renormalize, const float* bias, @@ -676,6 +677,28 @@ void topkGatingKernelLauncher( static constexpr int BYTES_PER_LDG_MULTIPLE_64 = (std::is_same_v || std::is_same_v) ? 4 : 8; #endif + // Inductor can pad GEMM outputs (for example, a [M, 60] view with a row + // stride of 64). The fused kernels assume packed rows, so use the + // stride-aware unfused path for row-padded inputs. + if (input_row_stride != num_experts) { + STD_TORCH_CHECK(workspace != nullptr, + "workspace must be provided for row-padded gating output."); + static constexpr int TPB = 256; + if constexpr (SF == SCORING_SOFTMAX) { + moeSoftmax<<>>( + gating_output, nullptr, workspace, num_experts, input_row_stride); + } else if constexpr (SF == SCORING_SIGMOID) { + moeSigmoid<<>>( + gating_output, nullptr, workspace, num_experts, input_row_stride); + } else { + STD_TORCH_CHECK(false, "Unsupported scoring func"); + } + moeTopK<<>>( + workspace, nullptr, topk_weights, topk_indices, token_expert_indices, + num_experts, topk, 0, num_experts, renormalize, bias, + routed_scaling_factor, is_padding); + return; + } switch (num_experts) { case 1: LAUNCH_TOPK(1, WARPS_PER_TB, BYTES_PER_LDG_POWER_OF_2); @@ -733,10 +756,12 @@ void topkGatingKernelLauncher( static constexpr int TPB = 256; if constexpr (SF == SCORING_SOFTMAX) { moeSoftmax<<>>( - gating_output, nullptr, workspace, num_experts); + gating_output, nullptr, workspace, num_experts, + input_row_stride); } else if constexpr (SF == SCORING_SIGMOID) { moeSigmoid<<>>( - gating_output, nullptr, workspace, num_experts); + gating_output, nullptr, workspace, num_experts, + input_row_stride); } else { STD_TORCH_CHECK(false, "Unsupported scoring func"); } @@ -795,7 +820,7 @@ void dispatch_topk_launch( topk_indices.mutable_data_ptr(), token_expert_indices.mutable_data_ptr(), softmax_workspace.mutable_data_ptr(), - num_tokens, num_experts, topk, renormalize, + num_tokens, num_experts, gating_output.stride(0), topk, renormalize, bias_ptr, routed_scaling_factor, stream, is_padding_ptr); } else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) { vllm::moe::topkGatingKernelLauncher( @@ -804,7 +829,7 @@ void dispatch_topk_launch( topk_indices.mutable_data_ptr(), token_expert_indices.mutable_data_ptr(), softmax_workspace.mutable_data_ptr(), - num_tokens, num_experts, topk, renormalize, + num_tokens, num_experts, gating_output.stride(0), topk, renormalize, bias_ptr, routed_scaling_factor, stream, is_padding_ptr); } else { STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long); @@ -814,7 +839,7 @@ void dispatch_topk_launch( topk_indices.mutable_data_ptr(), token_expert_indices.mutable_data_ptr(), softmax_workspace.mutable_data_ptr(), - num_tokens, num_experts, topk, renormalize, + num_tokens, num_experts, gating_output.stride(0), topk, renormalize, bias_ptr, routed_scaling_factor, stream, is_padding_ptr); } } @@ -828,12 +853,17 @@ void topk_softmax( std::optional bias, std::optional is_padding) { + STD_TORCH_CHECK(gating_output.dim() == 2, + "gating_output must be a 2D tensor"); + STD_TORCH_CHECK(gating_output.stride(1) == 1, + "gating_output's last dimension must be contiguous"); const int num_experts = gating_output.size(-1); const auto num_tokens = gating_output.numel() / num_experts; const int topk = topk_weights.size(-1); const bool is_pow_2 = (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0); - const bool needs_workspace = !is_pow_2 || num_experts > 256; + const bool needs_workspace = !is_pow_2 || num_experts > 256 || + gating_output.stride(0) != num_experts; const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0; torch::stable::accelerator::DeviceGuard guard(gating_output.get_device_index()); @@ -869,12 +899,17 @@ void topk_sigmoid( double routed_scaling_factor, std::optional is_padding) { + STD_TORCH_CHECK(gating_output.dim() == 2, + "gating_output must be a 2D tensor"); + STD_TORCH_CHECK(gating_output.stride(1) == 1, + "gating_output's last dimension must be contiguous"); const int num_experts = gating_output.size(-1); const auto num_tokens = gating_output.numel() / num_experts; const int topk = topk_weights.size(-1); const bool is_pow_2 = (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0); - const bool needs_workspace = !is_pow_2 || num_experts > 256; + const bool needs_workspace = !is_pow_2 || num_experts > 256 || + gating_output.stride(0) != num_experts; const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0; torch::stable::accelerator::DeviceGuard guard(gating_output.get_device_index()); diff --git a/tests/kernels/moe/test_fused_topk.py b/tests/kernels/moe/test_fused_topk.py index 18d1d9b4dfb0..7c27c7111591 100644 --- a/tests/kernels/moe/test_fused_topk.py +++ b/tests/kernels/moe/test_fused_topk.py @@ -89,6 +89,60 @@ def test_fused_topk( torch.testing.assert_close(topk_ids_ref.to(torch.int32), topk_ids, atol=0, rtol=0) +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="This test is skipped on non-CUDA platform.", +) +@pytest.mark.parametrize("num_experts,padded_num_experts", [(60, 64), (64, 72)]) +@pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"]) +def test_fused_topk_padded_row_stride( + num_experts: int, + padded_num_experts: int, + scoring_func: str, +): + """Regression test for row-padded router logits produced by Inductor.""" + torch.manual_seed(0) + num_tokens = 33 + hidden_size = 128 + topk = 4 + dtype = torch.bfloat16 + + hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda") + storage = torch.full( + (num_tokens, padded_num_experts), + 1e4, + dtype=dtype, + device="cuda", + ) + gating_output = storage[:, :num_experts] + gating_output.copy_( + torch.linspace(-4, 4, num_experts, dtype=dtype, device="cuda").expand( + num_tokens, -1 + ) + ) + assert gating_output.stride() == (padded_num_experts, 1) + assert not gating_output.is_contiguous() + + topk_weights_ref, topk_ids_ref = torch_topk( + gating_output=gating_output, + topk=topk, + renormalize=True, + scoring_func=scoring_func, + ) + topk_weights, topk_ids, _ = fused_topk( + hidden_states=hidden_states, + gating_output=gating_output, + topk=topk, + renormalize=True, + scoring_func=scoring_func, + ) + + torch.testing.assert_close( + topk_weights_ref.to(torch.float32), topk_weights, atol=1e-2, rtol=1e-2 + ) + torch.testing.assert_close(topk_ids_ref.to(torch.int32), topk_ids, atol=0, rtol=0) + + @pytest.mark.skipif( not current_platform.is_cuda_alike(), reason="This test is skipped on non-CUDA platform.", From a1b399738d06f1341dc255914c6f731c27ba1156 Mon Sep 17 00:00:00 2001 From: Id545 Date: Tue, 27 Oct 2026 06:37:17 +0100 Subject: [PATCH 2/2] [Bugfix][MoE] Accept a single expert column with any stride A (M, 1) gating output has a contiguous last dimension whatever its stride(1) reports, so the check now only requires stride(1) == 1 when there is more than one expert, as suggested in review. Signed-off-by: Id545 --- csrc/moe/topk_softmax_kernels.cu | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/csrc/moe/topk_softmax_kernels.cu b/csrc/moe/topk_softmax_kernels.cu index c7239ad36c1e..0483583c96ac 100644 --- a/csrc/moe/topk_softmax_kernels.cu +++ b/csrc/moe/topk_softmax_kernels.cu @@ -887,9 +887,9 @@ void topk_softmax( { STD_TORCH_CHECK(gating_output.dim() == 2, "gating_output must be a 2D tensor"); - STD_TORCH_CHECK(gating_output.stride(1) == 1, - "gating_output's last dimension must be contiguous"); const int num_experts = gating_output.size(-1); + STD_TORCH_CHECK(num_experts == 1 || gating_output.stride(1) == 1, + "gating_output's last dimension must be contiguous"); const auto num_tokens = gating_output.numel() / num_experts; const int topk = topk_weights.size(-1); @@ -933,9 +933,9 @@ void topk_sigmoid( { STD_TORCH_CHECK(gating_output.dim() == 2, "gating_output must be a 2D tensor"); - STD_TORCH_CHECK(gating_output.stride(1) == 1, - "gating_output's last dimension must be contiguous"); const int num_experts = gating_output.size(-1); + STD_TORCH_CHECK(num_experts == 1 || gating_output.stride(1) == 1, + "gating_output's last dimension must be contiguous"); const auto num_tokens = gating_output.numel() / num_experts; const int topk = topk_weights.size(-1);