A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 05:00 UTC
[BugFix][Router Replay] Capture Logical Experts with EPLB (#33013)
Signed-off-by: Hollow Man <hollowman@opensuse.org>
ℍ𝕠𝕝𝕝𝕠𝕨 𝕄𝕒𝕟 committed
Jan 31, 2026 at 17:12 UTC
13b842f2714558abd96f260ab450d192ce687ffb
4 files changed
+185
-21
tests/model_executor/test_routed_experts_capture.py
new
+160
@@ -0,0 +1,160 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import types
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.distributed.eplb.eplb_state import EplbLayerState
9
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
10
+from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
11
+
12
+pytestmark = pytest.mark.cpu_test
13
+
14
+
15
+class DummyRouter(BaseRouter):
16
+ @property
17
+ def routing_method_type(self) -> RoutingMethodType:
18
+ return RoutingMethodType.FUSED_TOPK
19
+
20
+ def _compute_routing(self, hidden_states, router_logits, indices_type):
21
+ topk_ids = torch.tensor([[1, 2], [3, 4]], dtype=torch.int64)
22
+ topk_weights = torch.ones_like(topk_ids, dtype=torch.float32)
23
+ return topk_weights, topk_ids
24
+
25
+ def _apply_eplb_mapping(self, topk_ids: torch.Tensor) -> torch.Tensor:
26
+ # Make mapping observable without requiring CUDA EPLB path.
27
+ return topk_ids + 10
28
+
29
+
30
+def _make_router() -> DummyRouter:
31
+ return DummyRouter(
32
+ top_k=2,
33
+ global_num_experts=16,
34
+ eplb_state=EplbLayerState(),
35
+ enable_eplb=False,
36
+ indices_type_getter=None,
37
+ )
38
+
39
+
40
+def test_base_router_capture_pre_eplb_mapping():
41
+ router = _make_router()
42
+ captured = []
43
+
44
+ def capture_fn(ids):
45
+ captured.append(ids.clone())
46
+
47
+ router.set_capture_fn(capture_fn)
48
+ topk_weights, topk_ids = router.select_experts(
49
+ hidden_states=torch.empty(1),
50
+ router_logits=torch.empty(1),
51
+ )
52
+
53
+ assert topk_weights.shape == topk_ids.shape
54
+ assert len(captured) == 1
55
+ assert torch.equal(captured[0], torch.tensor([[1, 2], [3, 4]]))
56
+ assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]]))
57
+
58
+
59
+def test_base_router_capture_with_eplb_enabled():
60
+ router = _make_router()
61
+ router.enable_eplb = True
62
+ router.eplb_state.expert_load_view = torch.zeros(32, dtype=torch.int64)
63
+ router.eplb_state.logical_to_physical_map = torch.arange(32).view(32, 1)
64
+ router.eplb_state.logical_replica_count = torch.ones(32, dtype=torch.int64)
65
+
66
+ captured = []
67
+
68
+ def capture_fn(ids):
69
+ captured.append(ids.clone())
70
+
71
+ router.set_capture_fn(capture_fn)
72
+ _, topk_ids = router.select_experts(
73
+ hidden_states=torch.empty(1),
74
+ router_logits=torch.empty(1),
75
+ )
76
+
77
+ assert len(captured) == 1
78
+ # Capture should see logical ids pre-EPLB mapping.
79
+ assert torch.equal(captured[0], torch.tensor([[1, 2], [3, 4]]))
80
+ # Our DummyRouter mapping adds +10.
81
+ assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]]))
82
+
83
+
84
+def test_gpu_model_runner_binds_router_capture(monkeypatch):
85
+ from vllm.v1.worker import gpu_model_runner as gmr
86
+
87
+ class DummyFusedMoE:
88
+ def __init__(self):
89
+ self.layer_id = 7
90
+ self.router = _make_router()
91
+
92
+ class DummyCapturer:
93
+ def __init__(self):
94
+ self.calls = []
95
+
96
+ def capture(self, layer_id, topk_ids):
97
+ self.calls.append((layer_id, topk_ids))
98
+
99
+ dummy_module = DummyFusedMoE()
100
+
101
+ # Patch the runtime import inside _bind_routed_experts_capturer.
102
+ import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
103
+
104
+ monkeypatch.setattr(fused_moe_layer, "FusedMoE", DummyFusedMoE)
105
+
106
+ dummy_self = types.SimpleNamespace(
107
+ compilation_config=types.SimpleNamespace(
108
+ static_forward_context={"dummy": dummy_module}
109
+ )
110
+ )
111
+
112
+ capturer = DummyCapturer()
113
+ gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, capturer)
114
+
115
+ assert dummy_module.router.capture_fn is not None
116
+ dummy_module.router.capture_fn(torch.tensor([[5, 6]]))
117
+
118
+ assert len(capturer.calls) == 1
119
+ layer_id, topk_ids = capturer.calls[0]
120
+ assert layer_id == 7
121
+ assert torch.equal(topk_ids, torch.tensor([[5, 6]]))
122
+
123
+
124
+def test_gpu_model_runner_binding_stage(monkeypatch):
125
+ from vllm.v1.worker import gpu_model_runner as gmr
126
+
127
+ class DummyFusedMoE:
128
+ def __init__(self):
129
+ self.layer_id = 11
130
+ self.router = _make_router()
131
+
132
+ class DummyCapturer:
133
+ def __init__(self):
134
+ self.calls = []
135
+
136
+ def capture(self, layer_id, topk_ids):
137
+ self.calls.append((layer_id, topk_ids))
138
+
139
+ dummy_module = DummyFusedMoE()
140
+
141
+ import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
142
+
143
+ monkeypatch.setattr(fused_moe_layer, "FusedMoE", DummyFusedMoE)
144
+
145
+ dummy_self = types.SimpleNamespace(
146
+ compilation_config=types.SimpleNamespace(
147
+ static_forward_context={"dummy": dummy_module}
148
+ )
149
+ )
150
+
151
+ # Before binding, no capture hook.
152
+ assert dummy_module.router.capture_fn is None
153
+
154
+ capturer = DummyCapturer()
155
+ gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, capturer)
156
+
157
+ # After binding, hook should exist and be callable.
158
+ assert callable(dummy_module.router.capture_fn)
159
+ dummy_module.router.capture_fn(torch.tensor([[9, 10]]))
160
+ assert len(capturer.calls) == 1
vllm/model_executor/layers/fused_moe/layer.py
-21
@@ -44,9 +44,6 @@ from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import (
44
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
45
init_aiter_topK_meta_data,
46
)
47
-from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
48
- RoutedExpertsCapturer,
49
-)
47
from vllm.model_executor.layers.fused_moe.router.router_factory import (
48
create_fused_moe_router,
49
)
@@ -523,18 +520,6 @@ class FusedMoE(CustomOp):
520
self.apply_router_weight_on_input = apply_router_weight_on_input
521
self.activation = activation
522
526
- self.capture: Callable[[torch.Tensor], None] | None = None
527
- if (
528
- self.vllm_config.model_config is not None
529
- and self.vllm_config.model_config.enable_return_routed_experts
530
- ):
531
- # In dummy runs, the capturer is not initialized.
532
- capturer = RoutedExpertsCapturer.get_instance()
533
- if capturer is not None:
534
- self.capture = lambda topk_ids: capturer.capture(
535
- self.layer_id, topk_ids
536
- )
537
-
523
self.router = create_fused_moe_router(
524
top_k=top_k,
525
global_num_experts=self.global_num_experts,
@@ -1688,9 +1673,6 @@ class FusedMoE(CustomOp):
1673
router_logits=staged_router_logits,
1674
)
1675
1691
- if self.capture is not None:
1692
- self.capture(topk_ids)
1693
-
1676
final_hidden_states = self.quant_method.apply(
1677
layer=self,
1678
x=staged_hidden_states,
@@ -1883,9 +1865,6 @@ class FusedMoE(CustomOp):
1865
router_logits=router_logits,
1866
)
1867
1886
- if self.capture is not None:
1887
- self.capture(topk_ids)
1888
-
1868
final_hidden_states = self.quant_method.apply(
1869
layer=self,
1870
x=x, # The type signture of this is wrong due to the hack.
vllm/model_executor/layers/fused_moe/router/base_router.py
+9
@@ -127,6 +127,11 @@ class BaseRouter(FusedMoERouter):
127
self.eplb_state = eplb_state
128
self.enable_eplb = enable_eplb
129
self.indices_type_getter = indices_type_getter
130
+ self.capture_fn: Callable[[torch.Tensor], None] | None = None
131
+
132
+ def set_capture_fn(self, capture_fn: Callable[[torch.Tensor], None] | None) -> None:
133
+ """Set a capture callback for logical routed expert IDs."""
134
+ self.capture_fn = capture_fn
135
136
def _validate_eplb_state(self) -> None:
137
"""Validate that EPLB state is properly initialized if EPLB is enabled."""
@@ -231,6 +236,10 @@ class BaseRouter(FusedMoERouter):
236
hidden_states, router_logits, indices_type
237
)
238
239
+ # Capture logical ids before EPLB mapping.
240
+ if self.capture_fn is not None:
241
+ self.capture_fn(topk_ids)
242
+
243
# Step 4: Apply EPLB mapping
244
topk_ids = self._apply_eplb_mapping(topk_ids)
245
vllm/v1/worker/gpu_model_runner.py
+16
@@ -6068,6 +6068,22 @@ class GPUModelRunner(
6068
max_num_kv_tokens=self.max_num_kv_tokens,
6069
vllm_config=self.vllm_config,
6070
)
6071
+ self._bind_routed_experts_capturer(routed_experts_capturer)
6072
+
6073
+ def _bind_routed_experts_capturer(self, capturer: RoutedExpertsCapturer) -> None:
6074
+ from vllm.model_executor.layers.fused_moe.layer import FusedMoE
6075
+ from vllm.model_executor.layers.fused_moe.router.base_router import (
6076
+ BaseRouter,
6077
+ )
6078
+
6079
+ for module in self.compilation_config.static_forward_context.values():
6080
+ if isinstance(module, FusedMoE) and isinstance(module.router, BaseRouter):
6081
+ layer_id = module.layer_id
6082
+
6083
+ def _capture_fn(topk_ids, _layer_id=layer_id, _capturer=capturer):
6084
+ _capturer.capture(_layer_id, topk_ids)
6085
+
6086
+ module.router.set_capture_fn(_capture_fn)
6087
6088
def may_add_encoder_only_layers_to_kv_cache_config(self) -> None:
6089
"""