[MoE][Bugfix] Skip SP padded rows in grouped MoE routing (#56079)
Signed-off-by: Elvir Crncevic <elvircrn@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Elvir Crnčević committed
Sep 17, 2026 at 21:36 UTC
ac2f0ea82c0d3005eb1de081cda715f0ee524c16
12 files changed
+224
-6
tests/kernels/moe/test_grouped_topk.py
+32
@@ -15,8 +15,10 @@ from vllm.config import (
15
get_cached_compilation_config,
16
set_current_vllm_config,
17
)
18
+from vllm.forward_context import set_forward_context
19
from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import (
20
GroupedTopk,
21
+ GroupedTopKRouter,
22
fused_grouped_topk,
23
)
24
from vllm.platforms import current_platform
@@ -70,6 +72,36 @@ def _single_group_reference(
72
return values, indices.to(torch.int32)
73
74
75
+@pytest.mark.skipif(
76
+ not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
77
+)
78
+@pytest.mark.parametrize("skip_padding", [False, True])
79
+def test_kimi_grouped_topk_masks_padding_when_enabled(skip_padding: bool):
80
+ hidden_states = torch.randn((2, 16), device="cuda")
81
+ gating_output = torch.randn((2, 896), device="cuda")
82
+ is_padding = torch.tensor([False, True], device="cuda")
83
+ router = GroupedTopKRouter(
84
+ top_k=16,
85
+ global_num_experts=896,
86
+ num_expert_group=1,
87
+ topk_group=1,
88
+ skip_padding=skip_padding,
89
+ )
90
+
91
+ with set_forward_context(None, VllmConfig(), is_padding=is_padding):
92
+ topk_weights, topk_ids = router._compute_routing(
93
+ hidden_states, gating_output, torch.int32
94
+ )
95
+
96
+ assert torch.all(topk_ids[0] >= 0)
97
+ if skip_padding:
98
+ assert torch.all(topk_ids[1] == -1)
99
+ assert torch.all(topk_weights[1] == 0)
100
+ else:
101
+ assert torch.all(topk_ids[1] >= 0)
102
+ assert torch.all(topk_weights[1] > 0)
103
+
104
+
105
@pytest.mark.skipif(
106
not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
107
)
tests/kernels/moe/test_routing.py
+15
@@ -87,6 +87,21 @@ def test_multiple_expert_groups_use_grouped_topk() -> None:
87
)
88
89
assert isinstance(router, GroupedTopKRouter)
90
+ assert not router.skip_padding
91
+
92
+
93
+def test_grouped_topk_padding_skip_must_be_enabled() -> None:
94
+ router = create_fused_moe_router(
95
+ top_k=4,
96
+ global_num_experts=128,
97
+ use_grouped_topk=True,
98
+ num_expert_group=8,
99
+ topk_group=4,
100
+ skip_padding=True,
101
+ )
102
+
103
+ assert isinstance(router, GroupedTopKRouter)
104
+ assert router.skip_padding
105
106
107
def test_degenerate_grouped_config_with_bias_uses_topk_bias() -> None:
tests/v1/worker/test_gpu_model_runner.py
+14
@@ -174,6 +174,20 @@ def test_freeze_gc_disables_and_restores_automatic_gc(
174
gc.disable()
175
176
177
+def test_prepare_padding_mask_marks_sequence_parallel_padding():
178
+ runner = GPUModelRunner.__new__(GPUModelRunner)
179
+ runner.is_padding = torch.empty(8, dtype=torch.bool)
180
+
181
+ mask = runner._prepare_padding_mask(1, 8)
182
+
183
+ assert mask.tolist() == [False, True, True, True, True, True, True, True]
184
+ assert mask.data_ptr() == runner.is_padding.data_ptr()
185
+
186
+ mask = runner._prepare_padding_mask(0, 8)
187
+
188
+ assert mask.all()
189
+
190
+
191
@pytest.fixture
192
def model_runner():
193
vllm_config = get_vllm_config()
tests/v1/worker/test_gpu_model_runner_v2.py
+14
@@ -21,6 +21,20 @@ from vllm.v1.worker.gpu.block_table import BlockTables
21
from vllm.v1.worker.gpu.model_runner import GPUModelRunner
22
23
24
+def test_prepare_padding_mask_marks_sequence_parallel_padding():
25
+ runner = GPUModelRunner.__new__(GPUModelRunner)
26
+ runner.input_buffers = SimpleNamespace(is_padding=torch.empty(8, dtype=torch.bool))
27
+
28
+ mask = runner._prepare_padding_mask(1, 8)
29
+
30
+ assert mask.tolist() == [False, True, True, True, True, True, True, True]
31
+ assert mask.data_ptr() == runner.input_buffers.is_padding.data_ptr()
32
+
33
+ mask = runner._prepare_padding_mask(0, 8)
34
+
35
+ assert mask.all()
36
+
37
+
38
def test_qsa_circular_group_uses_custom_slot_mapping(monkeypatch):
39
runner = GPUModelRunner.__new__(GPUModelRunner)
40
runner.max_model_len = 262144
tests/v1/worker/test_gpu_ubatch_slicing.py
+90
@@ -10,6 +10,7 @@ threaded execution of the microbatches.
10
"""
11
12
import threading
13
+from contextlib import nullcontext
14
from dataclasses import replace
15
from types import SimpleNamespace
16
from typing import Any, cast
@@ -30,6 +31,7 @@ from vllm.config import (
31
from vllm.forward_context import create_forward_context
32
from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
33
from vllm.v1.kv_cache_interface import KVCacheConfig
34
+from vllm.v1.worker import gpu_ubatch_wrapper as legacy_gpu_ubatch_wrapper
35
from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils
36
from vllm.v1.worker.gpu import cudagraph_utils, dp_utils
37
from vllm.v1.worker.gpu.cudagraph_utils import (
@@ -426,6 +428,94 @@ def test_slice_model_inputs_handles_mrope_positions():
428
assert sliced["intermediate_tensors"] is None
429
430
431
+@pytest.mark.parametrize(
432
+ ("cudagraph_runtime_mode", "execution_method"),
433
+ [
434
+ (CUDAGraphMode.FULL, "_capture_ubatches"),
435
+ (CUDAGraphMode.NONE, "_run_ubatches"),
436
+ ],
437
+)
438
+def test_legacy_ubatch_contexts_slice_padding_mask(
439
+ cudagraph_runtime_mode: CUDAGraphMode, execution_method: str
440
+):
441
+ """Replacement contexts preserve each microbatch's padding rows."""
442
+ ubatch_slices = [
443
+ UBatchSlice(slice(0, 2), slice(0, 3)),
444
+ UBatchSlice(slice(2, 4), slice(3, 6)),
445
+ ]
446
+ is_padding = torch.tensor([False, True, False, True, False, True])
447
+ forward_context = SimpleNamespace(
448
+ batch_descriptor=object(),
449
+ ubatch_slices=ubatch_slices,
450
+ cudagraph_runtime_mode=cudagraph_runtime_mode,
451
+ attn_metadata=[{}, {}],
452
+ slot_mapping={},
453
+ dp_metadata=object(),
454
+ is_padding=is_padding,
455
+ )
456
+
457
+ wrapper = object.__new__(legacy_gpu_ubatch_wrapper.UBatchWrapper)
458
+ wrapper.vllm_config = SimpleNamespace(
459
+ parallel_config=SimpleNamespace(data_parallel_size=2)
460
+ )
461
+ wrapper.comm_stream = object()
462
+ wrapper.ready_barrier = object()
463
+ wrapper.cudagraphs = {}
464
+ wrapper.sm_control = nullcontext()
465
+ wrapper.runnable = object()
466
+
467
+ def make_contexts(*, forward_contexts, **kwargs):
468
+ return [
469
+ SimpleNamespace(id=i, forward_context=context)
470
+ for i, context in enumerate(forward_contexts)
471
+ ]
472
+
473
+ def create_context(attn_metadata, vllm_config, **kwargs):
474
+ return SimpleNamespace(attn_metadata=attn_metadata, **kwargs)
475
+
476
+ model_inputs = {
477
+ "input_ids": torch.arange(6),
478
+ "positions": torch.arange(6),
479
+ "inputs_embeds": None,
480
+ "intermediate_tensors": None,
481
+ }
482
+ expected = object()
483
+ with (
484
+ patch.object(
485
+ legacy_gpu_ubatch_wrapper,
486
+ "get_forward_context",
487
+ return_value=forward_context,
488
+ ),
489
+ patch.object(
490
+ legacy_gpu_ubatch_wrapper.torch.cuda,
491
+ "current_stream",
492
+ return_value=object(),
493
+ ),
494
+ patch.object(
495
+ legacy_gpu_ubatch_wrapper.DPMetadata, "make", return_value=object()
496
+ ),
497
+ patch.object(
498
+ legacy_gpu_ubatch_wrapper,
499
+ "create_forward_context",
500
+ side_effect=create_context,
501
+ ),
502
+ patch.object(
503
+ legacy_gpu_ubatch_wrapper,
504
+ "make_ubatch_contexts",
505
+ side_effect=make_contexts,
506
+ ),
507
+ patch.object(wrapper, execution_method, return_value=expected) as execute,
508
+ ):
509
+ result = wrapper(**model_inputs)
510
+
511
+ assert result is expected
512
+ contexts = [
513
+ metadata.context.forward_context for metadata in execute.call_args.args[0]
514
+ ]
515
+ torch.testing.assert_close(contexts[0].is_padding, is_padding[:3])
516
+ torch.testing.assert_close(contexts[1].is_padding, is_padding[3:])
517
+
518
+
519
def _make_dbo_config() -> VllmConfig:
520
return VllmConfig(
521
model_config=ModelConfig(model="facebook/opt-125m", dtype="float16", seed=0),
vllm/model_executor/layers/fused_moe/layer.py
+3
@@ -137,6 +137,7 @@ def FusedMoEFactory(
137
runner_args: dict[str, Any] | None = None,
138
routed_experts_cls: type[RoutedExperts] | None = None,
139
routed_experts_args: dict[str, Any] | None = None,
140
+ skip_padding: bool = False,
141
) -> MoERunner:
142
"""Factory function for creating MoE execution pipeline.
143
@@ -213,6 +214,7 @@ def FusedMoEFactory(
214
runner_args: Additional arguments for runner constructor
215
routed_experts_cls: Custom RoutedExperts class (None = use default)
216
routed_experts_args: Additional arguments for routed_experts constructor
217
+ skip_padding: Whether grouped routing should invalidate padding rows.
218
219
Returns:
220
MoERunner: Configured MoE execution pipeline ready for forward passes
@@ -308,6 +310,7 @@ def FusedMoEFactory(
310
else 1.0,
311
e_score_correction_bias=e_score_correction_bias,
312
num_fused_shared_experts=num_fused_shared_experts,
313
+ skip_padding=(skip_padding and moe_parallel_config.use_deepep_v2_kernels),
314
# Fused shared-expert slot weight. With apply_routed_scale_to_output
315
# the runner scales the combined output by routed_scaling_factor, so
316
# the shared slot weight must be 1/routed_scaling_factor for its net
vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py
+14
@@ -8,6 +8,7 @@ from vllm import _custom_ops as ops
8
from vllm import envs as envs
9
from vllm._aiter_ops import rocm_aiter_ops
10
from vllm.distributed.eplb.eplb_state import EplbLayerState
11
+from vllm.forward_context import get_forward_context, is_forward_context_available
12
from vllm.model_executor.custom_op import CustomOp
13
from vllm.model_executor.layers.fused_moe.config import (
14
RoutingMethodType,
@@ -258,6 +259,7 @@ class GroupedTopKRouter(BaseRouter):
259
e_score_correction_bias: torch.Tensor | None = None,
260
num_fused_shared_experts: int = 0,
261
eplb_state: EplbLayerState | None = None,
262
+ skip_padding: bool = False,
263
):
264
super().__init__(
265
top_k=top_k,
@@ -271,6 +273,7 @@ class GroupedTopKRouter(BaseRouter):
273
self.routed_scaling_factor = routed_scaling_factor
274
self.e_score_correction_bias = e_score_correction_bias
275
self.num_fused_shared_experts = num_fused_shared_experts
276
+ self.skip_padding = skip_padding
277
278
@property
279
def routing_method_type(self) -> RoutingMethodType:
@@ -346,4 +349,15 @@ class GroupedTopKRouter(BaseRouter):
349
e_score_correction_bias=self.e_score_correction_bias,
350
)
351
352
+ if (
353
+ self.skip_padding
354
+ and envs.VLLM_MOE_SKIP_PADDING
355
+ and is_forward_context_available()
356
+ ):
357
+ is_padding = get_forward_context().is_padding
358
+ if is_padding is not None:
359
+ is_padding = is_padding[: topk_ids.shape[0]].unsqueeze(1)
360
+ topk_weights = topk_weights.masked_fill(is_padding, 0)
361
+ topk_ids = topk_ids.masked_fill(is_padding, -1)
362
+
363
return topk_weights, topk_ids
vllm/model_executor/layers/fused_moe/router/router_factory.py
+3
@@ -65,6 +65,7 @@ def create_fused_moe_router(
65
# Deepseek V4 vision routing bias parameters
66
bias_vl: torch.Tensor | None = None,
67
image_sentinel_lo: int = 0,
68
+ skip_padding: bool = False,
69
) -> FusedMoERouter:
70
"""Factory function to create the appropriate FusedMoERouter subclass based on
71
the provided parameters.
@@ -91,6 +92,7 @@ def create_fused_moe_router(
92
topk_group: Top-k within each group (for grouped routing)
93
scoring_func: Scoring function to use ("softmax" or "sigmoid")
94
num_fused_shared_experts: Number of fused shared experts (for ROCm AITER)
95
+ skip_padding: Whether to invalidate routes for padding rows
96
97
Grouped topk and fused topk bias arguments:
98
routed_scaling_factor: Scaling factor for routed weights
@@ -209,6 +211,7 @@ def create_fused_moe_router(
211
routed_scaling_factor=routed_scaling_factor,
212
e_score_correction_bias=e_score_correction_bias,
213
num_fused_shared_experts=num_fused_shared_experts,
214
+ skip_padding=skip_padding,
215
)
216
# Otherwise fall through to the non-grouped chain below.
217
vllm/models/kimi_k3/nvidia/model.py
+1
@@ -761,6 +761,7 @@ class KimiMoE(nn.Module):
761
routed_input_transform=None,
762
routed_output_transform=self.routed_output_transform,
763
is_sequence_parallel=use_sequence_parallel,
764
+ skip_padding=True,
765
runner_cls=LatentMoERunner if self.use_latent_moe else None,
766
)
767
if self.padded_moe_intermediate_size != moe_intermediate_size:
vllm/v1/worker/gpu/model_runner.py
+13
-6
@@ -1258,6 +1258,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1258
num_reqs, num_toks, max_query_len, batch_state.has_prefill
1259
)
1260
1261
+ def _prepare_padding_mask(
1262
+ self, num_tokens: int, num_tokens_after_padding: int
1263
+ ) -> torch.Tensor:
1264
+ is_padding = self.input_buffers.is_padding[:num_tokens_after_padding]
1265
+ is_padding[:num_tokens].fill_(False)
1266
+ is_padding[num_tokens:].fill_(True)
1267
+ return is_padding
1268
+
1269
def prepare_inputs(
1270
self,
1271
scheduler_output: SchedulerOutput,
@@ -1267,12 +1275,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1275
num_tokens = batch_req_state.num_tokens
1276
num_tokens_after_padding = max(num_tokens, batch_desc.num_tokens)
1277
assert num_tokens > 0
1278
+ is_padding = self.input_buffers.is_padding[:num_tokens_after_padding]
1279
if envs.VLLM_MOE_SKIP_PADDING:
1271
- # Mark trailing cudagraph-padding rows so kernels can skip work for
1272
- # them when supported.
1273
- is_padding = self.input_buffers.is_padding
1274
- is_padding[:num_tokens].fill_(False)
1275
- is_padding[num_tokens:num_tokens_after_padding].fill_(True)
1280
+ is_padding = self._prepare_padding_mask(
1281
+ num_tokens, num_tokens_after_padding
1282
+ )
1283
1284
req_ids = batch_req_state.req_ids
1285
num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens
@@ -1440,7 +1447,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1447
has_prefill=batch_req_state.has_prefill,
1448
input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],
1449
positions=self.input_buffers.positions[:num_tokens_after_padding],
1443
- is_padding=self.input_buffers.is_padding[:num_tokens_after_padding],
1450
+ is_padding=is_padding,
1451
logits_indices=logits_indices,
1452
cu_num_logits=cu_num_logits,
1453
cu_num_logits_np=cu_num_logits_np,
vllm/v1/worker/gpu_model_runner.py
+16
@@ -816,6 +816,9 @@ class GPUModelRunner(
816
self.positions = torch.zeros(
817
self.max_num_tokens, dtype=torch.int64, device=self.device
818
)
819
+ self.is_padding = torch.zeros(
820
+ self.max_num_tokens, dtype=torch.bool, device=self.device
821
+ )
822
self.query_start_loc = self._make_buffer(
823
self.max_num_reqs + 1, dtype=torch.int32
824
)
@@ -3500,6 +3503,14 @@ class GPUModelRunner(
3503
return round_up(num_scheduled_tokens, tp_size)
3504
return num_scheduled_tokens
3505
3506
+ def _prepare_padding_mask(
3507
+ self, num_tokens_unpadded: int, num_tokens_padded: int
3508
+ ) -> torch.Tensor:
3509
+ padding_mask = self.is_padding[:num_tokens_padded]
3510
+ padding_mask[:num_tokens_unpadded].fill_(False)
3511
+ padding_mask[num_tokens_unpadded:].fill_(True)
3512
+ return padding_mask
3513
+
3514
def _prepare_mm_inputs(
3515
self, num_tokens: int
3516
) -> tuple[torch.Tensor | None, torch.Tensor]:
@@ -4441,6 +4452,7 @@ class GPUModelRunner(
4452
num_tokens_unpadded,
4453
ubatch_slices_padded,
4454
)
4455
+ is_padding = self._prepare_padding_mask(num_tokens_unpadded, num_tokens_padded)
4456
with (
4457
set_forward_context(
4458
attn_metadata,
@@ -4452,6 +4464,7 @@ class GPUModelRunner(
4464
ubatch_slices=ubatch_slices_padded,
4465
slot_mapping=slot_mappings,
4466
skip_compiled=has_encoder_input,
4467
+ is_padding=is_padding,
4468
),
4469
record_function_or_nullcontext("gpu_model_runner: forward"),
4470
self.maybe_get_kv_connector_output(
@@ -6159,6 +6172,8 @@ class GPUModelRunner(
6172
if num_tokens_across_dp is not None:
6173
num_tokens_across_dp[:] = num_tokens_padded
6174
6175
+ is_padding = self._prepare_padding_mask(0, num_tokens_padded)
6176
+
6177
with (
6178
self.maybe_randomize_inputs(
6179
input_ids, inputs_embeds, randomize_inputs=randomize_inputs
@@ -6172,6 +6187,7 @@ class GPUModelRunner(
6187
batch_descriptor=batch_desc,
6188
ubatch_slices=ubatch_slices_padded,
6189
slot_mapping=slot_mappings,
6190
+ is_padding=is_padding,
6191
),
6192
):
6193
outputs = self.model(
vllm/v1/worker/gpu_ubatch_wrapper.py
+9
@@ -261,6 +261,7 @@ class UBatchWrapper:
261
dp_metadata,
262
batch_descriptor,
263
cudagraph_runtime_mode,
264
+ is_padding,
265
) -> list[UbatchMetadata]:
266
# Create one forward context per ubatch
267
forward_contexts = []
@@ -276,6 +277,11 @@ class UBatchWrapper:
277
batch_descriptor=batch_descriptor,
278
cudagraph_runtime_mode=cudagraph_runtime_mode,
279
slot_mapping=slot_mapping[i] if has_slot_mapping else None,
280
+ is_padding=(
281
+ is_padding[ubatch_slice.token_slice]
282
+ if is_padding is not None
283
+ else None
284
+ ),
285
)
286
)
287
@@ -373,6 +379,7 @@ class UBatchWrapper:
379
380
attn_metadata = forward_context.attn_metadata
381
slot_mapping = forward_context.slot_mapping
382
+ is_padding = forward_context.is_padding
383
num_tokens = sum(ubatch_slice.num_tokens for ubatch_slice in ubatch_slices)
384
input_ids = kwargs["input_ids"]
385
positions = kwargs["positions"]
@@ -414,6 +421,7 @@ class UBatchWrapper:
421
dp_metadata=ubatch_dp_metadata,
422
batch_descriptor=batch_descriptor,
423
cudagraph_runtime_mode=CUDAGraphMode.NONE,
424
+ is_padding=is_padding,
425
)
426
with self.sm_control:
427
return self._capture_ubatches(ubatch_metadata, self.runnable)
@@ -440,6 +448,7 @@ class UBatchWrapper:
448
dp_metadata=ubatch_dp_metadata,
449
batch_descriptor=batch_descriptor,
450
cudagraph_runtime_mode=CUDAGraphMode.NONE,
451
+ is_padding=is_padding,
452
)
453
with self.sm_control:
454
return self._run_ubatches(ubatch_metadata, self.runnable)