[MoE][Bugfix] Skip SP padded rows in grouped MoE routing

tests/kernels/moe/test_grouped_topk.py

@@ -15,8 +15,10 @@
     get_cached_compilation_config,
     set_current_vllm_config,
 )
+from vllm.forward_context import set_forward_context
 from vllm.model_executor.layers.fused_moe.router.grouped_topk_router import (
     GroupedTopk,
+    GroupedTopKRouter,
     fused_grouped_topk,
 )
 from vllm.platforms import current_platform
@@ -70,6 +72,36 @@ def _single_group_reference(
     return values, indices.to(torch.int32)
 
 
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
+)
+@pytest.mark.parametrize("skip_padding", [False, True])
+def test_kimi_grouped_topk_masks_padding_when_enabled(skip_padding: bool):
+    hidden_states = torch.randn((2, 16), device="cuda")
+    gating_output = torch.randn((2, 896), device="cuda")
+    is_padding = torch.tensor([False, True], device="cuda")
+    router = GroupedTopKRouter(
+        top_k=16,
+        global_num_experts=896,
+        num_expert_group=1,
+        topk_group=1,
+        skip_padding=skip_padding,
+    )
+
+    with set_forward_context(None, VllmConfig(), is_padding=is_padding):
+        topk_weights, topk_ids = router._compute_routing(
+            hidden_states, gating_output, torch.int32
+        )
+
+    assert torch.all(topk_ids[0] >= 0)
+    if skip_padding:
+        assert torch.all(topk_ids[1] == -1)
+        assert torch.all(topk_weights[1] == 0)
+    else:
+        assert torch.all(topk_ids[1] >= 0)
+        assert torch.all(topk_weights[1] > 0)
+
+
 @pytest.mark.skipif(
     not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
 )

tests/kernels/moe/test_routing.py

@@ -87,6 +87,21 @@ def test_multiple_expert_groups_use_grouped_topk() -> None:
     )
 
     assert isinstance(router, GroupedTopKRouter)
+    assert not router.skip_padding
+
+
+def test_grouped_topk_padding_skip_must_be_enabled() -> None:
+    router = create_fused_moe_router(
+        top_k=4,
+        global_num_experts=128,
+        use_grouped_topk=True,
+        num_expert_group=8,
+        topk_group=4,
+        skip_padding=True,
+    )
+
+    assert isinstance(router, GroupedTopKRouter)
+    assert router.skip_padding
 
 
 def test_degenerate_grouped_config_with_bias_uses_topk_bias() -> None:

tests/v1/worker/test_gpu_model_runner.py

@@ -174,6 +174,20 @@ def test_freeze_gc_disables_and_restores_automatic_gc(
             gc.disable()
 
 
+def test_prepare_padding_mask_marks_sequence_parallel_padding():
+    runner = GPUModelRunner.__new__(GPUModelRunner)
+    runner.is_padding = torch.empty(8, dtype=torch.bool)
+
+    mask = runner._prepare_padding_mask(1, 8)
+
+    assert mask.tolist() == [False, True, True, True, True, True, True, True]
+    assert mask.data_ptr() == runner.is_padding.data_ptr()
+
+    mask = runner._prepare_padding_mask(0, 8)
+
+    assert mask.all()
+
+
 @pytest.fixture
 def model_runner():
     vllm_config = get_vllm_config()

tests/v1/worker/test_gpu_model_runner_v2.py

@@ -21,6 +21,20 @@
 from vllm.v1.worker.gpu.model_runner import GPUModelRunner
 
 
+def test_prepare_padding_mask_marks_sequence_parallel_padding():
+    runner = GPUModelRunner.__new__(GPUModelRunner)
+    runner.input_buffers = SimpleNamespace(is_padding=torch.empty(8, dtype=torch.bool))
+
+    mask = runner._prepare_padding_mask(1, 8)
+
+    assert mask.tolist() == [False, True, True, True, True, True, True, True]
+    assert mask.data_ptr() == runner.input_buffers.is_padding.data_ptr()
+
+    mask = runner._prepare_padding_mask(0, 8)
+
+    assert mask.all()
+
+
 def test_qsa_circular_group_uses_custom_slot_mapping(monkeypatch):
     runner = GPUModelRunner.__new__(GPUModelRunner)
     runner.max_model_len = 262144

tests/v1/worker/test_gpu_ubatch_slicing.py

@@ -10,6 +10,7 @@
 """
 
 import threading
+from contextlib import nullcontext
 from dataclasses import replace
 from types import SimpleNamespace
 from typing import Any, cast
@@ -30,6 +31,7 @@
 from vllm.forward_context import create_forward_context
 from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
 from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.worker import gpu_ubatch_wrapper as legacy_gpu_ubatch_wrapper
 from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils
 from vllm.v1.worker.gpu import cudagraph_utils, dp_utils
 from vllm.v1.worker.gpu.cudagraph_utils import (
@@ -426,6 +428,94 @@ def test_slice_model_inputs_handles_mrope_positions():
     assert sliced["intermediate_tensors"] is None
 
 
+@pytest.mark.parametrize(
+    ("cudagraph_runtime_mode", "execution_method"),
+    [
+        (CUDAGraphMode.FULL, "_capture_ubatches"),
+        (CUDAGraphMode.NONE, "_run_ubatches"),
+    ],
+)
+def test_legacy_ubatch_contexts_slice_padding_mask(
+    cudagraph_runtime_mode: CUDAGraphMode, execution_method: str
+):
+    """Replacement contexts preserve each microbatch's padding rows."""
+    ubatch_slices = [
+        UBatchSlice(slice(0, 2), slice(0, 3)),
+        UBatchSlice(slice(2, 4), slice(3, 6)),
+    ]
+    is_padding = torch.tensor([False, True, False, True, False, True])
+    forward_context = SimpleNamespace(
+        batch_descriptor=object(),
+        ubatch_slices=ubatch_slices,
+        cudagraph_runtime_mode=cudagraph_runtime_mode,
+        attn_metadata=[{}, {}],
+        slot_mapping={},
+        dp_metadata=object(),
+        is_padding=is_padding,
+    )
+
+    wrapper = object.__new__(legacy_gpu_ubatch_wrapper.UBatchWrapper)
+    wrapper.vllm_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(data_parallel_size=2)
+    )
+    wrapper.comm_stream = object()
+    wrapper.ready_barrier = object()
+    wrapper.cudagraphs = {}
+    wrapper.sm_control = nullcontext()
+    wrapper.runnable = object()
+
+    def make_contexts(*, forward_contexts, **kwargs):
+        return [
+            SimpleNamespace(id=i, forward_context=context)
+            for i, context in enumerate(forward_contexts)
+        ]
+
+    def create_context(attn_metadata, vllm_config, **kwargs):
+        return SimpleNamespace(attn_metadata=attn_metadata, **kwargs)
+
+    model_inputs = {
+        "input_ids": torch.arange(6),
+        "positions": torch.arange(6),
+        "inputs_embeds": None,
+        "intermediate_tensors": None,
+    }
+    expected = object()
+    with (
+        patch.object(
+            legacy_gpu_ubatch_wrapper,
+            "get_forward_context",
+            return_value=forward_context,
+        ),
+        patch.object(
+            legacy_gpu_ubatch_wrapper.torch.cuda,
+            "current_stream",
+            return_value=object(),
+        ),
+        patch.object(
+            legacy_gpu_ubatch_wrapper.DPMetadata, "make", return_value=object()
+        ),
+        patch.object(
+            legacy_gpu_ubatch_wrapper,
+            "create_forward_context",
+            side_effect=create_context,
+        ),
+        patch.object(
+            legacy_gpu_ubatch_wrapper,
+            "make_ubatch_contexts",
+            side_effect=make_contexts,
+        ),
+        patch.object(wrapper, execution_method, return_value=expected) as execute,
+    ):
+        result = wrapper(**model_inputs)
+
+    assert result is expected
+    contexts = [
+        metadata.context.forward_context for metadata in execute.call_args.args[0]
+    ]
+    torch.testing.assert_close(contexts[0].is_padding, is_padding[:3])
+    torch.testing.assert_close(contexts[1].is_padding, is_padding[3:])
+
+
 def _make_dbo_config() -> VllmConfig:
     return VllmConfig(
         model_config=ModelConfig(model="facebook/opt-125m", dtype="float16", seed=0),

vllm/model_executor/layers/fused_moe/layer.py

@@ -137,6 +137,7 @@ def FusedMoEFactory(
     runner_args: dict[str, Any] | None = None,
     routed_experts_cls: type[RoutedExperts] | None = None,
     routed_experts_args: dict[str, Any] | None = None,
+    skip_padding: bool = False,
 ) -> MoERunner:
     """Factory function for creating MoE execution pipeline.
 
@@ -213,6 +214,7 @@ def FusedMoEFactory(
         runner_args: Additional arguments for runner constructor
         routed_experts_cls: Custom RoutedExperts class (None = use default)
         routed_experts_args: Additional arguments for routed_experts constructor
+        skip_padding: Whether grouped routing should invalidate padding rows.
 
     Returns:
         MoERunner: Configured MoE execution pipeline ready for forward passes
@@ -308,6 +310,7 @@ def FusedMoEFactory(
             else 1.0,
             e_score_correction_bias=e_score_correction_bias,
             num_fused_shared_experts=num_fused_shared_experts,
+            skip_padding=(skip_padding and moe_parallel_config.use_deepep_v2_kernels),
             # Fused shared-expert slot weight. With apply_routed_scale_to_output
             # the runner scales the combined output by routed_scaling_factor, so
             # the shared slot weight must be 1/routed_scaling_factor for its net

vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py

@@ -8,6 +8,7 @@
 from vllm import envs as envs
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.distributed.eplb.eplb_state import EplbLayerState
+from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.model_executor.custom_op import CustomOp
 from vllm.model_executor.layers.fused_moe.config import (
     RoutingMethodType,
@@ -258,6 +259,7 @@ def __init__(
         e_score_correction_bias: torch.Tensor | None = None,
         num_fused_shared_experts: int = 0,
         eplb_state: EplbLayerState | None = None,
+        skip_padding: bool = False,
     ):
         super().__init__(
             top_k=top_k,
@@ -271,6 +273,7 @@ def __init__(
         self.routed_scaling_factor = routed_scaling_factor
         self.e_score_correction_bias = e_score_correction_bias
         self.num_fused_shared_experts = num_fused_shared_experts
+        self.skip_padding = skip_padding
 
     @property
     def routing_method_type(self) -> RoutingMethodType:
@@ -346,4 +349,15 @@ def valid_grouping() -> bool:
             e_score_correction_bias=self.e_score_correction_bias,
         )
 
+        if (
+            self.skip_padding
+            and envs.VLLM_MOE_SKIP_PADDING
+            and is_forward_context_available()
+        ):
+            is_padding = get_forward_context().is_padding
+            if is_padding is not None:
+                is_padding = is_padding[: topk_ids.shape[0]].unsqueeze(1)
+                topk_weights = topk_weights.masked_fill(is_padding, 0)
+                topk_ids = topk_ids.masked_fill(is_padding, -1)
+
         return topk_weights, topk_ids

vllm/model_executor/layers/fused_moe/router/router_factory.py

@@ -65,6 +65,7 @@ def create_fused_moe_router(
     # Deepseek V4 vision routing bias parameters
     bias_vl: torch.Tensor | None = None,
     image_sentinel_lo: int = 0,
+    skip_padding: bool = False,
 ) -> FusedMoERouter:
     """Factory function to create the appropriate FusedMoERouter subclass based on
     the provided parameters.
@@ -91,6 +92,7 @@ def create_fused_moe_router(
         topk_group: Top-k within each group (for grouped routing)
         scoring_func: Scoring function to use ("softmax" or "sigmoid")
         num_fused_shared_experts: Number of fused shared experts (for ROCm AITER)
+        skip_padding: Whether to invalidate routes for padding rows
 
     Grouped topk and fused topk bias arguments:
         routed_scaling_factor: Scaling factor for routed weights
@@ -209,6 +211,7 @@ def advertised_routing_method(groups: int | None) -> RoutingMethodType:
                 routed_scaling_factor=routed_scaling_factor,
                 e_score_correction_bias=e_score_correction_bias,
                 num_fused_shared_experts=num_fused_shared_experts,
+                skip_padding=skip_padding,
             )
         # Otherwise fall through to the non-grouped chain below.
 

vllm/models/kimi_k3/nvidia/model.py

@@ -761,6 +761,7 @@ def __init__(
                 routed_input_transform=None,
                 routed_output_transform=self.routed_output_transform,
                 is_sequence_parallel=use_sequence_parallel,
+                skip_padding=True,
                 runner_cls=LatentMoERunner if self.use_latent_moe else None,
             )
         if self.padded_moe_intermediate_size != moe_intermediate_size:

vllm/v1/worker/gpu/model_runner.py

@@ -1258,6 +1258,14 @@ def gather_batch_req_state(
             num_reqs, num_toks, max_query_len, batch_state.has_prefill
         )
 
+    def _prepare_padding_mask(
+        self, num_tokens: int, num_tokens_after_padding: int
+    ) -> torch.Tensor:
+        is_padding = self.input_buffers.is_padding[:num_tokens_after_padding]
+        is_padding[:num_tokens].fill_(False)
+        is_padding[num_tokens:].fill_(True)
+        return is_padding
+
     def prepare_inputs(
         self,
         scheduler_output: SchedulerOutput,
@@ -1267,12 +1275,11 @@ def prepare_inputs(
         num_tokens = batch_req_state.num_tokens
         num_tokens_after_padding = max(num_tokens, batch_desc.num_tokens)
         assert num_tokens > 0
+        is_padding = self.input_buffers.is_padding[:num_tokens_after_padding]
         if envs.VLLM_MOE_SKIP_PADDING:
-            # Mark trailing cudagraph-padding rows so kernels can skip work for
-            # them when supported.
-            is_padding = self.input_buffers.is_padding
-            is_padding[:num_tokens].fill_(False)
-            is_padding[num_tokens:num_tokens_after_padding].fill_(True)
+            is_padding = self._prepare_padding_mask(
+                num_tokens, num_tokens_after_padding
+            )
 
         req_ids = batch_req_state.req_ids
         num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens
@@ -1440,7 +1447,7 @@ def prepare_inputs(
             has_prefill=batch_req_state.has_prefill,
             input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],
             positions=self.input_buffers.positions[:num_tokens_after_padding],
-            is_padding=self.input_buffers.is_padding[:num_tokens_after_padding],
+            is_padding=is_padding,
             logits_indices=logits_indices,
             cu_num_logits=cu_num_logits,
             cu_num_logits_np=cu_num_logits_np,