[Bugfix][R3] Exclude draft routers from expert capture (#48622)
Signed-off-by: aoshen02 <aoshen@inferact.ai>
aoshen02 committed
Jul 15, 2026 at 08:45 UTC
3ca242d1b6282084b2a41e247810e632450c639a
2 files changed
+35
-7
tests/model_executor/test_routed_experts_capture.py
+34
-6
@@ -138,9 +138,7 @@ def test_gpu_model_runner_binds_router_capture(monkeypatch):
138
monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
139
140
dummy_self = types.SimpleNamespace(
141
- compilation_config=types.SimpleNamespace(
142
- static_forward_context={"dummy": dummy_module}
143
- )
141
+ model=types.SimpleNamespace(modules=lambda: [dummy_module])
142
)
143
144
capturer = DummyCapturer()
@@ -177,9 +175,7 @@ def test_gpu_model_runner_binding_stage(monkeypatch):
175
monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
176
177
dummy_self = types.SimpleNamespace(
180
- compilation_config=types.SimpleNamespace(
181
- static_forward_context={"dummy": dummy_module}
182
- )
178
+ model=types.SimpleNamespace(modules=lambda: [dummy_module])
179
)
180
181
# Before binding, no capture hook.
@@ -194,6 +190,38 @@ def test_gpu_model_runner_binding_stage(monkeypatch):
190
assert len(capturer.calls) == 1
191
192
193
+def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
194
+ from vllm.v1.worker import gpu_model_runner as gmr
195
+
196
+ class DummyFusedMoE:
197
+ def __init__(self, layer_id):
198
+ self.layer_id = layer_id
199
+ self.router = _make_router()
200
+
201
+ target_module = DummyFusedMoE(layer_id=7)
202
+ draft_module = DummyFusedMoE(layer_id=0)
203
+
204
+ import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
205
+
206
+ monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
207
+
208
+ dummy_self = types.SimpleNamespace(
209
+ model=types.SimpleNamespace(modules=lambda: [target_module]),
210
+ compilation_config=types.SimpleNamespace(
211
+ static_forward_context={
212
+ "model.layers.7.mlp.experts": target_module,
213
+ "mtp.layers.0.mlp.experts": draft_module,
214
+ }
215
+ ),
216
+ )
217
+
218
+ capturer = types.SimpleNamespace(capture=lambda *_: None)
219
+ gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, capturer)
220
+
221
+ assert target_module.router.capture_fn is not None
222
+ assert draft_module.router.capture_fn is None
223
+
224
+
225
def test_routed_experts_capturer_single_dp_no_metadata():
226
"""dp_metadata is None: capture writes the full topk_ids rows."""
227
capturer = _capturer_with_buffer(dp_rank=0)
vllm/v1/worker/gpu_model_runner.py
+1
-1
@@ -7589,7 +7589,7 @@ class GPUModelRunner(
7589
BaseRouter,
7590
)
7591
7592
- for module in self.compilation_config.static_forward_context.values():
7592
+ for module in self.model.modules():
7593
if isinstance(module, MoERunner) and isinstance(module.router, BaseRouter):
7594
layer_id = module.layer_id
7595