[Bugfix][Model Runner V2] Route dummy tokens to MoE experts during profiling (#57270)

Signed-off-by: Robert Shaw <robertgshaw2@gmail.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

tests/v1/worker/test_gpu_input_batch_v2.py

@@ -6,6 +6,9 @@ import numpy as np
 import pytest
 import torch
 
+from vllm.config import VllmConfig
+from vllm.forward_context import set_forward_context
+from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk
 from vllm.platforms import current_platform
 from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
 from vllm.v1.worker.gpu import cp_utils
@@ -56,6 +59,28 @@ def test_make_dummy_distributes_remainder(num_reqs: int, num_tokens: int):
     )
 
 
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA top-k.")
+@pytest.mark.parametrize("is_padding", [True, False])
+def test_make_dummy_padding_controls_moe_routing(monkeypatch, is_padding: bool):
+    """Dummy tokens marked as padding are dropped by the MoE router (top-k id
+    -1), which is wanted for idle DP ranks. Profile runs must not mark them, or
+    no token reaches the experts and MoE memory is never profiled."""
+    monkeypatch.setenv("VLLM_MOE_SKIP_PADDING", "1")
+    num_tokens = 16
+    buffers = InputBuffers(
+        max_num_reqs=4, max_num_tokens=num_tokens, device=torch.device(DEVICE)
+    )
+    batch = InputBatch.make_dummy(4, num_tokens, buffers, is_padding=is_padding)
+    hidden_states = torch.randn(num_tokens, 4, device=DEVICE)
+    router_logits = torch.randn(num_tokens, 8, device=DEVICE)
+
+    with set_forward_context(None, VllmConfig(), is_padding=batch.is_padding):
+        _, topk_ids, _ = fused_topk(hidden_states, router_logits, 2, False)
+
+    assert bool((topk_ids == -1).all()) is is_padding
+    assert bool((topk_ids >= 0).all()) is not is_padding
+
+
 def test_maybe_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch):
     """The batch must view the caller-owned buffer, sliced to padded length.
 

vllm/v1/worker/gpu/input_batch.py

@@ -125,6 +125,7 @@ class InputBatch:
         num_tokens: int,
         input_buffers: InputBuffers,
         max_query_len: int | None = None,
+        is_padding: bool = True,
     ) -> "InputBatch":
         assert 0 < num_reqs <= num_tokens
         device = input_buffers.device
@@ -168,7 +169,7 @@ class InputBatch:
         input_ids = input_buffers.input_ids[:num_tokens].zero_()
         positions = input_buffers.positions[:num_tokens].zero_()
 
-        input_buffers.is_padding[:num_tokens].fill_(True)
+        input_buffers.is_padding[:num_tokens].fill_(is_padding)
         is_padding = input_buffers.is_padding[:num_tokens]
 
         logits_indices = query_start_loc[1:] - 1

vllm/v1/worker/gpu/model_runner.py

@@ -1733,6 +1733,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 batch_desc.num_tokens,
                 self.input_buffers,
                 max_query_len=batch_desc.max_query_len,
+                # Profiling and warmup must route the dummy tokens to experts
+                # so MoE memory is measured and MoE kernels are exercised.
+                is_padding=not is_profile,
             )
             if self.pcp_manager is not None:
                 input_batch = self.pcp_manager.prepare_inputs_to_capture(input_batch)