[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>
Robert Shaw committed
Sep 17, 2026 at 10:11 UTC
861a299adc72bbbb1b60853b6173075465a6b425
3 files changed
+30
-1
tests/v1/worker/test_gpu_input_batch_v2.py
+25
@@ -6,6 +6,9 @@ import numpy as np
6
import pytest
7
import torch
8
9
+from vllm.config import VllmConfig
10
+from vllm.forward_context import set_forward_context
11
+from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk
12
from vllm.platforms import current_platform
13
from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
14
from vllm.v1.worker.gpu import cp_utils
@@ -56,6 +59,28 @@ def test_make_dummy_distributes_remainder(num_reqs: int, num_tokens: int):
59
)
60
61
62
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA top-k.")
63
+@pytest.mark.parametrize("is_padding", [True, False])
64
+def test_make_dummy_padding_controls_moe_routing(monkeypatch, is_padding: bool):
65
+ """Dummy tokens marked as padding are dropped by the MoE router (top-k id
66
+ -1), which is wanted for idle DP ranks. Profile runs must not mark them, or
67
+ no token reaches the experts and MoE memory is never profiled."""
68
+ monkeypatch.setenv("VLLM_MOE_SKIP_PADDING", "1")
69
+ num_tokens = 16
70
+ buffers = InputBuffers(
71
+ max_num_reqs=4, max_num_tokens=num_tokens, device=torch.device(DEVICE)
72
+ )
73
+ batch = InputBatch.make_dummy(4, num_tokens, buffers, is_padding=is_padding)
74
+ hidden_states = torch.randn(num_tokens, 4, device=DEVICE)
75
+ router_logits = torch.randn(num_tokens, 8, device=DEVICE)
76
+
77
+ with set_forward_context(None, VllmConfig(), is_padding=batch.is_padding):
78
+ _, topk_ids, _ = fused_topk(hidden_states, router_logits, 2, False)
79
+
80
+ assert bool((topk_ids == -1).all()) is is_padding
81
+ assert bool((topk_ids >= 0).all()) is not is_padding
82
+
83
+
84
def test_maybe_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch):
85
"""The batch must view the caller-owned buffer, sliced to padded length.
86
vllm/v1/worker/gpu/input_batch.py
+2
-1
@@ -125,6 +125,7 @@ class InputBatch:
125
num_tokens: int,
126
input_buffers: InputBuffers,
127
max_query_len: int | None = None,
128
+ is_padding: bool = True,
129
) -> "InputBatch":
130
assert 0 < num_reqs <= num_tokens
131
device = input_buffers.device
@@ -168,7 +169,7 @@ class InputBatch:
169
input_ids = input_buffers.input_ids[:num_tokens].zero_()
170
positions = input_buffers.positions[:num_tokens].zero_()
171
171
- input_buffers.is_padding[:num_tokens].fill_(True)
172
+ input_buffers.is_padding[:num_tokens].fill_(is_padding)
173
is_padding = input_buffers.is_padding[:num_tokens]
174
175
logits_indices = query_start_loc[1:] - 1
vllm/v1/worker/gpu/model_runner.py
+3
@@ -1733,6 +1733,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1733
batch_desc.num_tokens,
1734
self.input_buffers,
1735
max_query_len=batch_desc.max_query_len,
1736
+ # Profiling and warmup must route the dummy tokens to experts
1737
+ # so MoE memory is measured and MoE kernels are exercised.
1738
+ is_padding=not is_profile,
1739
)
1740
if self.pcp_manager is not None:
1741
input_batch = self.pcp_manager.prepare_inputs_to_capture(input_batch)