From 861a299adc72bbbb1b60853b6173075465a6b425 Mon Sep 17 00:00:00 2001 From: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Date: Thu, 17 Sep 2026 10:11:14 -0400 Subject: [PATCH] [Bugfix][Model Runner V2] Route dummy tokens to MoE experts during profiling (#57270) Signed-off-by: Robert Shaw Co-authored-by: Claude Opus 5 (1M context) --- tests/v1/worker/test_gpu_input_batch_v2.py | 25 ++++++++++++++++++++++ vllm/v1/worker/gpu/input_batch.py | 3 ++- vllm/v1/worker/gpu/model_runner.py | 3 +++ 3 files changed, 30 insertions(+), 1 deletion(-) diff --git a/tests/v1/worker/test_gpu_input_batch_v2.py b/tests/v1/worker/test_gpu_input_batch_v2.py index 8937f2dc18..13a92d2c0f 100644 --- a/tests/v1/worker/test_gpu_input_batch_v2.py +++ b/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. diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index d3c65bd06f..32c083e459 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/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 diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 97571df8a9..631e6b04d4 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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)