@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 46c8a161f4

[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 <zyy1102000@gmail.com> Co-authored-by: Codex <codex@openai.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>

Kevin H. Luu committed Sep 1, 2026 at 19:08 UTC 46c8a161f40cb8fbae0b9fd00552c9d2589fb8dc
6 files changed +216 -21
csrc/libtorch_stable/moe/moe_permute_unpermute_op.cu
+14 -7
@@ -212,23 +212,30 @@ __global__ void shuffleInputRowsKernel(const T* input,
212 int64_t num_src_rows,
213 int64_t num_dst_rows, int64_t num_cols) {
214 int64_t dest_row_idx = blockIdx.x;
215 - int64_t const source_row_idx = dst2src_map[dest_row_idx];
216 -
215 if (blockIdx.x < num_dst_rows) {
216 + int64_t const source_row_idx = dst2src_map[dest_row_idx];
217 // Load 128-bits per thread
218 constexpr int64_t ELEM_PER_THREAD = 128 / sizeof(T) / 8;
219 using DataElem = cutlass::Array<T, ELEM_PER_THREAD>;
220
222 - // Duplicate and permute rows
223 - auto const* source_row_ptr =
224 - reinterpret_cast<DataElem const*>(input + source_row_idx * num_cols);
221 auto* dest_row_ptr =
222 reinterpret_cast<DataElem*>(output + dest_row_idx * num_cols);
227 -
223 int64_t const start_offset = threadIdx.x;
224 int64_t const stride = blockDim.x;
225 int64_t const num_elems_in_col = num_cols / ELEM_PER_THREAD;
226
227 + if (source_row_idx < 0 || source_row_idx >= num_src_rows) {
228 + DataElem const zero{};
229 + for (int elem_index = start_offset; elem_index < num_elems_in_col;
230 + elem_index += stride) {
231 + dest_row_ptr[elem_index] = zero;
232 + }
233 + return;
234 + }
235 +
236 + // Duplicate and permute rows
237 + auto const* source_row_ptr =
238 + reinterpret_cast<DataElem const*>(input + source_row_idx * num_cols);
239 for (int elem_index = start_offset; elem_index < num_elems_in_col;
240 elem_index += stride) {
241 dest_row_ptr[elem_index] = source_row_ptr[elem_index];
@@ -322,4 +329,4 @@ STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) {
329 m.impl("moe_permute", TORCH_BOX(&moe_permute));
330 m.impl("moe_permute_with_scratch", TORCH_BOX(&moe_permute_with_scratch));
331 m.impl("moe_unpermute", TORCH_BOX(&moe_unpermute));
325 -}
\ No newline at end of file
332 +}
csrc/libtorch_stable/quantization/fp4/mxfp4_experts_quant.cu
+7 -2
@@ -76,9 +76,12 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
76
77 int tid = blockIdx.x * blockDim.x + threadIdx.x;
78 int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
79 + int32_t const numValidRows =
80 + min(numRows,
81 + static_cast<int32_t>(__ldca(&input_offset_by_experts[n_experts])));
82 int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
83
81 - for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
84 + for (int globalIdx = tid; globalIdx < numValidRows * colsPerRow;
85 globalIdx += gridDim.x * blockDim.x) {
86 int rowIdx = globalIdx / colsPerRow;
87 int colIdx = globalIdx % colsPerRow;
@@ -193,9 +196,11 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
196
197 int tid = blockIdx.x * blockDim.x + threadIdx.x;
198 int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
199 + int32_t const numValidRows =
200 + min(numRows, static_cast<int32_t>(shared_input_offsets[n_experts]));
201 int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
202
198 - for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
203 + for (int globalIdx = tid; globalIdx < numValidRows * colsPerRow;
204 globalIdx += gridDim.x * blockDim.x) {
205 int rowIdx = globalIdx / colsPerRow;
206 int colIdx = globalIdx % colsPerRow;
csrc/libtorch_stable/quantization/fp4/nvfp4_experts_quant.cu
+7 -2
@@ -53,11 +53,14 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
53
54 int tid = blockIdx.x * blockDim.x + threadIdx.x;
55 int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
56 + int32_t const numValidRows =
57 + min(numRows,
58 + static_cast<int32_t>(__ldca(&input_offset_by_experts[n_experts])));
59 // When fusing SiLU+Mul, input has gate || up layout (doubled width)
60 int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
61
62 // Each global thread processes one element
60 - for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
63 + for (int globalIdx = tid; globalIdx < numValidRows * colsPerRow;
64 globalIdx += gridDim.x * blockDim.x) {
65 // Calculate which row and column this global thread should process
66 int rowIdx = globalIdx / colsPerRow;
@@ -186,11 +189,13 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
189
190 int tid = blockIdx.x * blockDim.x + threadIdx.x;
191 int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
192 + int32_t const numValidRows =
193 + min(numRows, static_cast<int32_t>(shared_input_offsets[n_experts]));
194 // When fusing SiLU+Mul, input has gate || up layout (doubled width)
195 int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
196
197 // Each global thread processes one element
193 - for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
198 + for (int globalIdx = tid; globalIdx < numValidRows * colsPerRow;
199 globalIdx += gridDim.x * blockDim.x) {
200 // Calculate which row and column this global thread should process
201 int rowIdx = globalIdx / colsPerRow;
csrc/libtorch_stable/quantization/w8a8/cutlass/moe/moe_data.cu
+10 -10
@@ -62,6 +62,7 @@ __global__ void compute_expert_offsets(
62 tot_offset += swap_ab ? problem_sizes1[i * 3 + 1] : problem_sizes1[i * 3];
63 expert_offsets[i + 1] = tot_offset;
64 }
65 + atomic_buffer[num_experts] = tot_offset;
66 }
67
68 __global__ void compute_expert_blockscale_offsets(
@@ -81,28 +82,28 @@ __global__ void compute_expert_blockscale_offsets(
82 tot_offset_round += (cur_offset + (128 - 1)) / 128 * 128;
83 blockscale_offsets[i + 1] = tot_offset_round;
84 }
85 + atomic_buffer[num_experts] = tot_offset;
86 }
87
88 __global__ void compute_arg_sorts(const int32_t* __restrict__ topk_ids,
87 - const int32_t* __restrict__ expert_offsets,
89 int32_t* input_permutation,
90 int32_t* output_permutation,
91 int32_t* atomic_buffer, const int topk_length,
92 const int topk) {
93 int const blk_expert_id = blockIdx.x;
94 int const num_experts = gridDim.x;
94 - int32_t const num_tokens = expert_offsets[num_experts];
95
96 for (int i = threadIdx.x; i < topk_length; i += THREADS_PER_EXPERT) {
97 int const expert_id = topk_ids[i];
98 - if (expert_id == -1 && blockIdx.x == 0) {
98 + if ((expert_id < 0 || expert_id >= num_experts) && blockIdx.x == 0) {
99 + int const start = atomicAdd(&atomic_buffer[num_experts], 1);
100 + input_permutation[start] = -1;
101 // output_permutation is used to re-order the moe outputs. It is
102 // used as c2 = c2[c_map], where c2 is a torch.tensor that is the
101 - // output of the cutlass kernels and c_map is the output_permutation.
102 - // c2 is initialized to zeros, therefore by setting the output_permutation
103 - // to num_tokens, we are guaranteed to fill the moe outputs to zero
104 - // for "invalid" topk_ids.
105 - output_permutation[i] = num_tokens;
103 + // output of the cutlass kernels and c_map is the output_permutation. Use
104 + // out-of-range sentinels for invalid routes; shuffle_rows converts them
105 + // to zero rows.
106 + output_permutation[i] = topk_length;
107 } else if (expert_id == blk_expert_id) {
108 int start = atomicAdd(&atomic_buffer[expert_id], 1);
109 input_permutation[start] = i / topk;
@@ -246,7 +247,7 @@ void get_cutlass_moe_mm_data_caller(
247 const torch::stable::accelerator::DeviceGuard device_guard(device.index());
248 auto stream = get_current_cuda_stream(device.index());
249 torch::stable::Tensor atomic_buffer = torch::stable::new_zeros(
249 - topk_ids, {num_experts}, torch::headeronly::ScalarType::Int);
250 + topk_ids, {num_experts + 1}, torch::headeronly::ScalarType::Int);
251
252 int num_threads = min(THREADS_PER_EXPERT, topk_ids.numel());
253
@@ -275,7 +276,6 @@ void get_cutlass_moe_mm_data_caller(
276 }
277 compute_arg_sorts<<<num_experts, num_threads, 0, stream>>>(
278 static_cast<const int32_t*>(topk_ids.data_ptr()),
278 - static_cast<const int32_t*>(expert_offsets.data_ptr()),
279 static_cast<int32_t*>(input_permutation.data_ptr()),
280 static_cast<int32_t*>(output_permutation.data_ptr()),
281 static_cast<int32_t*>(atomic_buffer.data_ptr()), topk_ids.numel(),
tests/kernels/moe/test_cutlass_moe.py
+163
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import copy
4 import dataclasses
5 +from collections.abc import Callable
6 from math import prod
7
8 import pytest
@@ -33,7 +34,9 @@ from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
34 CutlassExpertsFp8,
35 CutlassExpertsMxfp4,
36 CutlassExpertsW4A8Fp8,
37 + run_cutlass_moe_fp4,
38 run_cutlass_moe_fp8,
39 + run_cutlass_moe_mxfp4,
40 )
41 from vllm.model_executor.layers.fused_moe.oracle import nvfp4 as nvfp4_oracle
42 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(
134 assert selected == expected
135
136
137 +@pytest.mark.skipif(
138 + (lambda x: x is None or not ops.cutlass_group_gemm_supported(x.to_int()))(
139 + current_platform.get_device_capability()
140 + ),
141 + reason="Grouped gemm is not supported on this GPU type.",
142 +)
143 +def test_cutlass_moe_permutation_maps_padding_to_zero():
144 + """Invalid top-k routes must not leave unsafe permutation-map entries."""
145 + num_experts = 2
146 + topk_ids = torch.tensor(
147 + [[0, -1], [1, num_experts]], device="cuda", dtype=torch.int32
148 + )
149 + num_routes = topk_ids.numel()
150 +
151 + expert_offsets = torch.empty(num_experts + 1, device="cuda", dtype=torch.int32)
152 + blockscale_offsets = torch.empty_like(expert_offsets)
153 + problem_sizes1 = torch.empty(num_experts, 3, device="cuda", dtype=torch.int32)
154 + problem_sizes2 = torch.empty_like(problem_sizes1)
155 + input_permutation = torch.empty(num_routes, device="cuda", dtype=torch.int32)
156 + output_permutation = torch.empty_like(input_permutation)
157 +
158 + ops.get_cutlass_moe_mm_data(
159 + topk_ids,
160 + expert_offsets,
161 + problem_sizes1,
162 + problem_sizes2,
163 + input_permutation,
164 + output_permutation,
165 + num_experts,
166 + 128,
167 + 128,
168 + blockscale_offsets,
169 + )
170 +
171 + input_tensor = torch.randn(2, 128, device="cuda", dtype=torch.bfloat16)
172 + permuted = ops.shuffle_rows(input_tensor, input_permutation)
173 + restored = ops.shuffle_rows(permuted, output_permutation)
174 +
175 + torch.testing.assert_close(permuted[:2], input_tensor)
176 + torch.testing.assert_close(permuted[2:], torch.zeros_like(permuted[2:]))
177 + torch.testing.assert_close(restored[0], input_tensor[0])
178 + torch.testing.assert_close(restored[1], torch.zeros_like(restored[1]))
179 + torch.testing.assert_close(restored[2], input_tensor[1])
180 + torch.testing.assert_close(restored[3], torch.zeros_like(restored[3]))
181 +
182 +
183 +@pytest.mark.parametrize("quantization", ["nvfp4", "mxfp4"])
184 +@torch.inference_mode()
185 +def test_cutlass_fp4_moe_padded_routes_do_not_change_valid_output(quantization: str):
186 + """Padded routes must not participate in either activation quantization."""
187 + experts_cls = CutlassExpertsFp4 if quantization == "nvfp4" else CutlassExpertsMxfp4
188 + if not experts_cls._supports_current_device():
189 + pytest.skip(f"CUTLASS {quantization} MoE is not supported on this GPU")
190 +
191 + set_random_seed(7)
192 + num_experts, n, k = 4, 128, 128
193 + dtype = torch.bfloat16
194 + device = torch.device("cuda")
195 + w1 = torch.full(
196 + (num_experts, 2 * n, k // 2), 0x11, device=device, dtype=torch.uint8
197 + )
198 + w2 = torch.full((num_experts, k, n // 2), 0x11, device=device, dtype=torch.uint8)
199 +
200 + runner: Callable[..., None]
201 + quant_args: tuple[torch.Tensor, ...]
202 + if quantization == "nvfp4":
203 + scale_kwargs = {
204 + "device": device,
205 + "dtype": torch.float8_e4m3fn,
206 + }
207 + w1_scale = torch.full((num_experts, 2 * n, k // 16), 1 / 16, **scale_kwargs)
208 + w2_scale = torch.full((num_experts, k, n // 16), 1 / 16, **scale_kwargs)
209 + gscale = torch.ones(num_experts, device=device, dtype=torch.float32)
210 + runner = run_cutlass_moe_fp4
211 + quant_args = (
212 + gscale,
213 + w1,
214 + w1_scale,
215 + gscale,
216 + gscale,
217 + w2,
218 + w2_scale,
219 + gscale,
220 + )
221 + else:
222 + scale_exponent = 127 - 4
223 + w1_scale = torch.full(
224 + (num_experts, 2 * n, k // 32),
225 + scale_exponent,
226 + device=device,
227 + dtype=torch.uint8,
228 + )
229 + w2_scale = torch.full(
230 + (num_experts, k, n // 32),
231 + scale_exponent,
232 + device=device,
233 + dtype=torch.uint8,
234 + )
235 + runner = run_cutlass_moe_mxfp4
236 + quant_args = (w1, w1_scale, w2, w2_scale)
237 +
238 + def run_moe(
239 + hidden_states: torch.Tensor,
240 + topk_ids: torch.Tensor,
241 + topk_weights: torch.Tensor,
242 + workspace_value: float,
243 + ) -> torch.Tensor:
244 + m, topk = topk_ids.shape
245 + output = torch.empty((m, k), device=device, dtype=dtype)
246 + workspace13 = torch.full(
247 + (m * topk, max(2 * n, k)),
248 + workspace_value,
249 + device=device,
250 + dtype=dtype,
251 + )
252 + workspace2 = torch.full(
253 + (m * topk, n), workspace_value, device=device, dtype=dtype
254 + )
255 +
256 + runner(
257 + output,
258 + hidden_states,
259 + *quant_args,
260 + topk_weights,
261 + topk_ids,
262 + MoEActivation.SILU,
263 + workspace13,
264 + workspace2,
265 + m,
266 + n,
267 + k,
268 + num_experts,
269 + device,
270 + )
271 + return output
272 +
273 + valid_input = torch.randn((1, k), device=device, dtype=dtype)
274 + expected = run_moe(
275 + valid_input,
276 + torch.zeros((1, 1), device=device, dtype=torch.int32),
277 + torch.ones((1, 1), device=device, dtype=torch.float32),
278 + workspace_value=0,
279 + )
280 + assert expected.abs().max() > 0
281 +
282 + # Force the quantizers' grid-stride loops to revisit padded tail rows.
283 + num_tokens = 4096
284 + padded_input = torch.zeros((num_tokens, k), device=device, dtype=dtype)
285 + padded_input[0] = valid_input[0]
286 + padded_ids = torch.full((num_tokens, 2), -1, device=device, dtype=torch.int32)
287 + padded_ids[0, 0] = 0
288 + padded_weights = torch.zeros((num_tokens, 2), device=device, dtype=torch.float32)
289 + padded_weights[0, 0] = 1
290 +
291 + for workspace_value in (0, 64):
292 + actual = run_moe(padded_input, padded_ids, padded_weights, workspace_value)
293 + torch.testing.assert_close(actual[0], expected[0], atol=1e-2, rtol=1e-2)
294 + torch.testing.assert_close(actual[1:], torch.zeros_like(actual[1:]))
295 +
296 +
297 @dataclasses.dataclass
298 class MOETensors:
299 a: torch.Tensor
tests/kernels/test_shuffle_rows.py
+15
@@ -103,6 +103,21 @@ def test_shuffle_rows_expansion(num_tokens: int, hidden_size: int):
103 )
104
105
106 +def test_shuffle_rows_invalid_indices_are_zeroed():
107 + """Invalid sentinel indices should produce zero rows instead of OOB reads."""
108 + if not current_platform.is_cuda():
109 + pytest.skip("shuffle_rows requires CUDA")
110 +
111 + input_tensor = torch.randn(3, 128, device="cuda", dtype=torch.bfloat16)
112 + dst2src_map = torch.tensor([-1, 1, 3], device="cuda", dtype=torch.int32)
113 +
114 + output = shuffle_rows(input_tensor, dst2src_map)
115 +
116 + torch.testing.assert_close(output[0], torch.zeros_like(output[0]))
117 + torch.testing.assert_close(output[1], input_tensor[1])
118 + torch.testing.assert_close(output[2], torch.zeros_like(output[2]))
119 +
120 +
121 @pytest.mark.parametrize("num_tokens", [16, 64])
122 @pytest.mark.parametrize("hidden_size", [128, 512])
123 def test_shuffle_rows_random_permutation(num_tokens: int, hidden_size: int):