From 46c8a161f40cb8fbae0b9fd00552c9d2589fb8dc Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Tue, 1 Sep 2026 19:08:17 -0700 Subject: [PATCH] [Bugfix] Handle padded routes in CUTLASS MoE permutations (#54747) Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com> Signed-off-by: Yongye Zhu Co-authored-by: Codex Co-authored-by: Yongye Zhu --- .../moe/moe_permute_unpermute_op.cu | 21 ++- .../quantization/fp4/mxfp4_experts_quant.cu | 9 +- .../quantization/fp4/nvfp4_experts_quant.cu | 9 +- .../quantization/w8a8/cutlass/moe/moe_data.cu | 20 +-- tests/kernels/moe/test_cutlass_moe.py | 163 ++++++++++++++++++ tests/kernels/test_shuffle_rows.py | 15 ++ 6 files changed, 216 insertions(+), 21 deletions(-) diff --git a/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu b/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu index 52cf68442f..251097456f 100644 --- a/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu +++ b/csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu @@ -212,23 +212,30 @@ __global__ void shuffleInputRowsKernel(const T* input, int64_t num_src_rows, int64_t num_dst_rows, int64_t num_cols) { int64_t dest_row_idx = blockIdx.x; - int64_t const source_row_idx = dst2src_map[dest_row_idx]; - if (blockIdx.x < num_dst_rows) { + int64_t const source_row_idx = dst2src_map[dest_row_idx]; // Load 128-bits per thread constexpr int64_t ELEM_PER_THREAD = 128 / sizeof(T) / 8; using DataElem = cutlass::Array; - // Duplicate and permute rows - auto const* source_row_ptr = - reinterpret_cast(input + source_row_idx * num_cols); auto* dest_row_ptr = reinterpret_cast(output + dest_row_idx * num_cols); - int64_t const start_offset = threadIdx.x; int64_t const stride = blockDim.x; int64_t const num_elems_in_col = num_cols / ELEM_PER_THREAD; + if (source_row_idx < 0 || source_row_idx >= num_src_rows) { + DataElem const zero{}; + for (int elem_index = start_offset; elem_index < num_elems_in_col; + elem_index += stride) { + dest_row_ptr[elem_index] = zero; + } + return; + } + + // Duplicate and permute rows + auto const* source_row_ptr = + reinterpret_cast(input + source_row_idx * num_cols); for (int elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) { dest_row_ptr[elem_index] = source_row_ptr[elem_index]; @@ -322,4 +329,4 @@ STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) { m.impl("moe_permute", TORCH_BOX(&moe_permute)); m.impl("moe_permute_with_scratch", TORCH_BOX(&moe_permute_with_scratch)); m.impl("moe_unpermute", TORCH_BOX(&moe_unpermute)); -} \ No newline at end of file +} diff --git a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu index 20f024bcef..086c07acf0 100644 --- a/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu +++ b/csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu @@ -76,9 +76,12 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) int tid = blockIdx.x * blockDim.x + threadIdx.x; int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD; + int32_t const numValidRows = + min(numRows, + static_cast(__ldca(&input_offset_by_experts[n_experts]))); int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow; - for (int globalIdx = tid; globalIdx < numRows * colsPerRow; + for (int globalIdx = tid; globalIdx < numValidRows * colsPerRow; globalIdx += gridDim.x * blockDim.x) { int rowIdx = globalIdx / colsPerRow; int colIdx = globalIdx % colsPerRow; @@ -193,9 +196,11 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024)) int tid = blockIdx.x * blockDim.x + threadIdx.x; int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD; + int32_t const numValidRows = + min(numRows, static_cast(shared_input_offsets[n_experts])); int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow; - for (int globalIdx = tid; globalIdx < numRows * colsPerRow; + for (int globalIdx = tid; globalIdx < numValidRows * colsPerRow; globalIdx += gridDim.x * blockDim.x) { int rowIdx = globalIdx / colsPerRow; int colIdx = globalIdx % colsPerRow; diff --git a/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu b/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu index 92b139b7e4..6296f0d901 100644 --- a/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu +++ b/csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu @@ -53,11 +53,14 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) int tid = blockIdx.x * blockDim.x + threadIdx.x; int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD; + int32_t const numValidRows = + min(numRows, + static_cast(__ldca(&input_offset_by_experts[n_experts]))); // When fusing SiLU+Mul, input has gate || up layout (doubled width) int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow; // Each global thread processes one element - for (int globalIdx = tid; globalIdx < numRows * colsPerRow; + for (int globalIdx = tid; globalIdx < numValidRows * colsPerRow; globalIdx += gridDim.x * blockDim.x) { // Calculate which row and column this global thread should process int rowIdx = globalIdx / colsPerRow; @@ -186,11 +189,13 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024)) int tid = blockIdx.x * blockDim.x + threadIdx.x; int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD; + int32_t const numValidRows = + min(numRows, static_cast(shared_input_offsets[n_experts])); // When fusing SiLU+Mul, input has gate || up layout (doubled width) int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow; // Each global thread processes one element - for (int globalIdx = tid; globalIdx < numRows * colsPerRow; + for (int globalIdx = tid; globalIdx < numValidRows * colsPerRow; globalIdx += gridDim.x * blockDim.x) { // Calculate which row and column this global thread should process int rowIdx = globalIdx / colsPerRow; diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu b/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu index 47e0985b12..8ebc3dd9e8 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu @@ -62,6 +62,7 @@ __global__ void compute_expert_offsets( tot_offset += swap_ab ? problem_sizes1[i * 3 + 1] : problem_sizes1[i * 3]; expert_offsets[i + 1] = tot_offset; } + atomic_buffer[num_experts] = tot_offset; } __global__ void compute_expert_blockscale_offsets( @@ -81,28 +82,28 @@ __global__ void compute_expert_blockscale_offsets( tot_offset_round += (cur_offset + (128 - 1)) / 128 * 128; blockscale_offsets[i + 1] = tot_offset_round; } + atomic_buffer[num_experts] = tot_offset; } __global__ void compute_arg_sorts(const int32_t* __restrict__ topk_ids, - const int32_t* __restrict__ expert_offsets, int32_t* input_permutation, int32_t* output_permutation, int32_t* atomic_buffer, const int topk_length, const int topk) { int const blk_expert_id = blockIdx.x; int const num_experts = gridDim.x; - int32_t const num_tokens = expert_offsets[num_experts]; for (int i = threadIdx.x; i < topk_length; i += THREADS_PER_EXPERT) { int const expert_id = topk_ids[i]; - if (expert_id == -1 && blockIdx.x == 0) { + if ((expert_id < 0 || expert_id >= num_experts) && blockIdx.x == 0) { + int const start = atomicAdd(&atomic_buffer[num_experts], 1); + input_permutation[start] = -1; // output_permutation is used to re-order the moe outputs. It is // used as c2 = c2[c_map], where c2 is a torch.tensor that is the - // output of the cutlass kernels and c_map is the output_permutation. - // c2 is initialized to zeros, therefore by setting the output_permutation - // to num_tokens, we are guaranteed to fill the moe outputs to zero - // for "invalid" topk_ids. - output_permutation[i] = num_tokens; + // output of the cutlass kernels and c_map is the output_permutation. Use + // out-of-range sentinels for invalid routes; shuffle_rows converts them + // to zero rows. + output_permutation[i] = topk_length; } else if (expert_id == blk_expert_id) { int start = atomicAdd(&atomic_buffer[expert_id], 1); input_permutation[start] = i / topk; @@ -246,7 +247,7 @@ void get_cutlass_moe_mm_data_caller( const torch::stable::accelerator::DeviceGuard device_guard(device.index()); auto stream = get_current_cuda_stream(device.index()); torch::stable::Tensor atomic_buffer = torch::stable::new_zeros( - topk_ids, {num_experts}, torch::headeronly::ScalarType::Int); + topk_ids, {num_experts + 1}, torch::headeronly::ScalarType::Int); int num_threads = min(THREADS_PER_EXPERT, topk_ids.numel()); @@ -275,7 +276,6 @@ void get_cutlass_moe_mm_data_caller( } compute_arg_sorts<<>>( static_cast(topk_ids.data_ptr()), - static_cast(expert_offsets.data_ptr()), static_cast(input_permutation.data_ptr()), static_cast(output_permutation.data_ptr()), static_cast(atomic_buffer.data_ptr()), topk_ids.numel(), diff --git a/tests/kernels/moe/test_cutlass_moe.py b/tests/kernels/moe/test_cutlass_moe.py index 274658ec0e..9a5189f592 100644 --- a/tests/kernels/moe/test_cutlass_moe.py +++ b/tests/kernels/moe/test_cutlass_moe.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy import dataclasses +from collections.abc import Callable from math import prod import pytest @@ -33,7 +34,9 @@ from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import ( CutlassExpertsFp8, CutlassExpertsMxfp4, CutlassExpertsW4A8Fp8, + run_cutlass_moe_fp4, run_cutlass_moe_fp8, + run_cutlass_moe_mxfp4, ) from vllm.model_executor.layers.fused_moe.oracle import nvfp4 as nvfp4_oracle from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input @@ -131,6 +134,166 @@ def test_nvfp4_clamp_allows_shared_activation_backends( assert selected == expected +@pytest.mark.skipif( + (lambda x: x is None or not ops.cutlass_group_gemm_supported(x.to_int()))( + current_platform.get_device_capability() + ), + reason="Grouped gemm is not supported on this GPU type.", +) +def test_cutlass_moe_permutation_maps_padding_to_zero(): + """Invalid top-k routes must not leave unsafe permutation-map entries.""" + num_experts = 2 + topk_ids = torch.tensor( + [[0, -1], [1, num_experts]], device="cuda", dtype=torch.int32 + ) + num_routes = topk_ids.numel() + + expert_offsets = torch.empty(num_experts + 1, device="cuda", dtype=torch.int32) + blockscale_offsets = torch.empty_like(expert_offsets) + problem_sizes1 = torch.empty(num_experts, 3, device="cuda", dtype=torch.int32) + problem_sizes2 = torch.empty_like(problem_sizes1) + input_permutation = torch.empty(num_routes, device="cuda", dtype=torch.int32) + output_permutation = torch.empty_like(input_permutation) + + ops.get_cutlass_moe_mm_data( + topk_ids, + expert_offsets, + problem_sizes1, + problem_sizes2, + input_permutation, + output_permutation, + num_experts, + 128, + 128, + blockscale_offsets, + ) + + input_tensor = torch.randn(2, 128, device="cuda", dtype=torch.bfloat16) + permuted = ops.shuffle_rows(input_tensor, input_permutation) + restored = ops.shuffle_rows(permuted, output_permutation) + + torch.testing.assert_close(permuted[:2], input_tensor) + torch.testing.assert_close(permuted[2:], torch.zeros_like(permuted[2:])) + torch.testing.assert_close(restored[0], input_tensor[0]) + torch.testing.assert_close(restored[1], torch.zeros_like(restored[1])) + torch.testing.assert_close(restored[2], input_tensor[1]) + torch.testing.assert_close(restored[3], torch.zeros_like(restored[3])) + + +@pytest.mark.parametrize("quantization", ["nvfp4", "mxfp4"]) +@torch.inference_mode() +def test_cutlass_fp4_moe_padded_routes_do_not_change_valid_output(quantization: str): + """Padded routes must not participate in either activation quantization.""" + experts_cls = CutlassExpertsFp4 if quantization == "nvfp4" else CutlassExpertsMxfp4 + if not experts_cls._supports_current_device(): + pytest.skip(f"CUTLASS {quantization} MoE is not supported on this GPU") + + set_random_seed(7) + num_experts, n, k = 4, 128, 128 + dtype = torch.bfloat16 + device = torch.device("cuda") + w1 = torch.full( + (num_experts, 2 * n, k // 2), 0x11, device=device, dtype=torch.uint8 + ) + w2 = torch.full((num_experts, k, n // 2), 0x11, device=device, dtype=torch.uint8) + + runner: Callable[..., None] + quant_args: tuple[torch.Tensor, ...] + if quantization == "nvfp4": + scale_kwargs = { + "device": device, + "dtype": torch.float8_e4m3fn, + } + w1_scale = torch.full((num_experts, 2 * n, k // 16), 1 / 16, **scale_kwargs) + w2_scale = torch.full((num_experts, k, n // 16), 1 / 16, **scale_kwargs) + gscale = torch.ones(num_experts, device=device, dtype=torch.float32) + runner = run_cutlass_moe_fp4 + quant_args = ( + gscale, + w1, + w1_scale, + gscale, + gscale, + w2, + w2_scale, + gscale, + ) + else: + scale_exponent = 127 - 4 + w1_scale = torch.full( + (num_experts, 2 * n, k // 32), + scale_exponent, + device=device, + dtype=torch.uint8, + ) + w2_scale = torch.full( + (num_experts, k, n // 32), + scale_exponent, + device=device, + dtype=torch.uint8, + ) + runner = run_cutlass_moe_mxfp4 + quant_args = (w1, w1_scale, w2, w2_scale) + + def run_moe( + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + workspace_value: float, + ) -> torch.Tensor: + m, topk = topk_ids.shape + output = torch.empty((m, k), device=device, dtype=dtype) + workspace13 = torch.full( + (m * topk, max(2 * n, k)), + workspace_value, + device=device, + dtype=dtype, + ) + workspace2 = torch.full( + (m * topk, n), workspace_value, device=device, dtype=dtype + ) + + runner( + output, + hidden_states, + *quant_args, + topk_weights, + topk_ids, + MoEActivation.SILU, + workspace13, + workspace2, + m, + n, + k, + num_experts, + device, + ) + return output + + valid_input = torch.randn((1, k), device=device, dtype=dtype) + expected = run_moe( + valid_input, + torch.zeros((1, 1), device=device, dtype=torch.int32), + torch.ones((1, 1), device=device, dtype=torch.float32), + workspace_value=0, + ) + assert expected.abs().max() > 0 + + # Force the quantizers' grid-stride loops to revisit padded tail rows. + num_tokens = 4096 + padded_input = torch.zeros((num_tokens, k), device=device, dtype=dtype) + padded_input[0] = valid_input[0] + padded_ids = torch.full((num_tokens, 2), -1, device=device, dtype=torch.int32) + padded_ids[0, 0] = 0 + padded_weights = torch.zeros((num_tokens, 2), device=device, dtype=torch.float32) + padded_weights[0, 0] = 1 + + for workspace_value in (0, 64): + actual = run_moe(padded_input, padded_ids, padded_weights, workspace_value) + torch.testing.assert_close(actual[0], expected[0], atol=1e-2, rtol=1e-2) + torch.testing.assert_close(actual[1:], torch.zeros_like(actual[1:])) + + @dataclasses.dataclass class MOETensors: a: torch.Tensor diff --git a/tests/kernels/test_shuffle_rows.py b/tests/kernels/test_shuffle_rows.py index c7de64066e..08d5883a1e 100644 --- a/tests/kernels/test_shuffle_rows.py +++ b/tests/kernels/test_shuffle_rows.py @@ -103,6 +103,21 @@ def test_shuffle_rows_expansion(num_tokens: int, hidden_size: int): ) +def test_shuffle_rows_invalid_indices_are_zeroed(): + """Invalid sentinel indices should produce zero rows instead of OOB reads.""" + if not current_platform.is_cuda(): + pytest.skip("shuffle_rows requires CUDA") + + input_tensor = torch.randn(3, 128, device="cuda", dtype=torch.bfloat16) + dst2src_map = torch.tensor([-1, 1, 3], device="cuda", dtype=torch.int32) + + output = shuffle_rows(input_tensor, dst2src_map) + + torch.testing.assert_close(output[0], torch.zeros_like(output[0])) + torch.testing.assert_close(output[1], input_tensor[1]) + torch.testing.assert_close(output[2], torch.zeros_like(output[2])) + + @pytest.mark.parametrize("num_tokens", [16, 64]) @pytest.mark.parametrize("hidden_size", [128, 512]) def test_shuffle_rows_random_permutation(num_tokens: int, hidden_size: int):