[MRV2] Enable routed-experts capture (#50721)
Signed-off-by: aoshen02 <aoshen@inferact.ai> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
aoshen02 committed
Aug 4, 2026 at 05:00 UTC
42ab184ea74ee6cf2966529c77bb51fd825a5d0c
10 files changed
+335
-227
tests/kernels/moe/test_routed_experts_capture_monolithic.py
+3
-5
@@ -385,11 +385,9 @@ def test_routed_experts_capturer_e2e_via_monolithic_experts() -> None:
385
on the monolithic experts and verify the captured rows land in the
386
capturer's device buffer at the correct layer slot.
387
388
- Mirrors the wiring done in ``GPUModelRunner._bind_routed_experts_capturer``
389
- for the monolithic path: a single closure is installed on the monolithic
390
- ``fused_experts`` (in addition to ``router.set_capture_fn`` on the
391
- non-monolithic path) and the capturer routes per-layer based on the
392
- closed-over ``layer_id``.
388
+ Mirrors the monolithic path in ``bind_routed_experts_capturer``: a single
389
+ closure is installed on ``fused_experts`` and routes per-layer based on
390
+ the closed-over ``layer_id``.
391
"""
392
from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
393
RoutedExpertsCapturer,
tests/model_executor/test_routed_experts_capture.py
+126
-107
@@ -2,15 +2,19 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import types
4
from types import SimpleNamespace
5
-from unittest.mock import patch
5
+from unittest.mock import Mock, patch
6
7
import pytest
8
import torch
9
10
+from vllm.config import VllmConfig
11
+from vllm.config.compilation import CompilationMode
12
from vllm.distributed.eplb.eplb_state import EplbLayerState
13
from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
14
from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
15
RoutedExpertsCapturer,
16
+ bind_routed_experts_capturer,
17
+ get_routed_experts_attn_gid,
18
)
19
from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
20
@@ -118,97 +122,12 @@ def test_base_router_capture_with_eplb_enabled():
122
assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]]))
123
124
121
-def test_gpu_model_runner_binds_router_capture(monkeypatch):
122
- from vllm.v1.worker import gpu_model_runner as gmr
123
-
124
- class _DummyRouter:
125
- _routing_replay_out: torch.Tensor | None = None
126
-
127
- class DummyFusedMoE:
128
- def __init__(self):
129
- self.layer_id = 7
130
- self.router = _make_router()
131
- self.routed_experts = _make_modular_routed_experts()
132
- self._quant_method = self.routed_experts.quant_method
133
-
134
- class DummyCapturer:
135
- def __init__(self):
136
- self.calls = []
137
-
138
- def capture(self, layer_id, topk_ids):
139
- self.calls.append((layer_id, topk_ids))
140
-
141
- dummy_module = DummyFusedMoE()
142
-
143
- # Patch the runtime import inside _bind_routed_experts_capturer.
144
- import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
145
-
146
- monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
147
-
148
- dummy_self = types.SimpleNamespace(
149
- model=types.SimpleNamespace(modules=lambda: [dummy_module])
150
- )
151
-
152
- capturer = DummyCapturer()
153
- gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, capturer)
154
-
155
- assert dummy_module.router.capture_fn is not None
156
- dummy_module.router.capture_fn(torch.tensor([[5, 6]]))
157
-
158
- assert len(capturer.calls) == 1
159
- layer_id, topk_ids = capturer.calls[0]
160
- assert layer_id == 7
161
- assert torch.equal(topk_ids, torch.tensor([[5, 6]]))
162
-
163
-
164
-def test_gpu_model_runner_binding_stage(monkeypatch):
165
- from vllm.v1.worker import gpu_model_runner as gmr
166
-
167
- class DummyFusedMoE:
168
- def __init__(self):
169
- self.layer_id = 11
170
- self.router = _make_router()
171
- self.routed_experts = _make_modular_routed_experts()
172
- self._quant_method = self.routed_experts.quant_method
173
-
174
- class DummyCapturer:
175
- def __init__(self):
176
- self.calls = []
177
-
178
- def capture(self, layer_id, topk_ids):
179
- self.calls.append((layer_id, topk_ids))
180
-
181
- dummy_module = DummyFusedMoE()
182
-
183
- import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
184
-
185
- monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
186
-
187
- dummy_self = types.SimpleNamespace(
188
- model=types.SimpleNamespace(modules=lambda: [dummy_module])
189
- )
190
-
191
- # Before binding, no capture hook.
192
- assert dummy_module.router.capture_fn is None
193
-
194
- capturer = DummyCapturer()
195
- gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, capturer)
196
-
197
- # After binding, hook should exist and be callable.
198
- assert callable(dummy_module.router.capture_fn)
199
- dummy_module.router.capture_fn(torch.tensor([[9, 10]]))
200
- assert len(capturer.calls) == 1
201
-
202
-
203
-def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
204
- from vllm.v1.worker import gpu_model_runner as gmr
205
-
125
+def test_public_binding_only_visits_target_model(monkeypatch):
126
class DummyFusedMoE:
127
def __init__(self, layer_id):
128
self.layer_id = layer_id
129
self.router = _make_router()
210
- self.routed_experts = _make_modular_routed_experts()
211
- self._quant_method = self.routed_experts.quant_method
130
+ self._quant_method = _make_modular_routed_experts().quant_method
131
132
target_module = DummyFusedMoE(layer_id=7)
133
draft_module = DummyFusedMoE(layer_id=0)
@@ -216,27 +135,21 @@ def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
135
import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
136
137
monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
138
+ calls = []
139
+ capturer = types.SimpleNamespace(capture=lambda *args: calls.append(args))
140
220
- dummy_self = types.SimpleNamespace(
221
- model=types.SimpleNamespace(modules=lambda: [target_module]),
222
- compilation_config=types.SimpleNamespace(
223
- static_forward_context={
224
- "model.layers.7.mlp.experts": target_module,
225
- "mtp.layers.0.mlp.experts": draft_module,
226
- }
227
- ),
141
+ bind_routed_experts_capturer(
142
+ types.SimpleNamespace(modules=lambda: [target_module]), capturer
143
)
144
230
- capturer = types.SimpleNamespace(capture=lambda *_: None)
231
- gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, capturer)
232
-
145
assert target_module.router.capture_fn is not None
146
assert draft_module.router.capture_fn is None
147
+ topk_ids = torch.tensor([[5, 6]])
148
+ target_module.router.capture_fn(topk_ids)
149
+ assert calls == [(7, topk_ids)]
150
151
237
-def test_gpu_model_runner_rejects_monolithic_without_replay_support(monkeypatch):
238
- from vllm.v1.worker import gpu_model_runner as gmr
239
-
152
+def test_public_binding_rejects_monolithic_without_replay_support(monkeypatch):
153
class DummyFusedMoE:
154
def __init__(self):
155
self.layer_id = 3
@@ -269,12 +182,10 @@ def test_gpu_model_runner_rejects_monolithic_without_replay_support(monkeypatch)
182
183
monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
184
272
- dummy_self = types.SimpleNamespace(
273
- model=types.SimpleNamespace(modules=lambda: [dummy_module])
274
- )
275
-
185
with pytest.raises(ValueError, match="monolithic MoE kernel"):
277
- gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, DummyCapturer())
186
+ bind_routed_experts_capturer(
187
+ types.SimpleNamespace(modules=lambda: [dummy_module]), DummyCapturer()
188
+ )
189
190
191
def test_routed_experts_capturer_single_dp_no_metadata():
@@ -333,3 +244,111 @@ def test_routed_experts_capturer_dp_unexpected_batch_raises():
244
):
245
capturer.capture(layer_id=0, topk_ids=topk)
246
assert capturer.device_buffer[0, 0, 0].item() == -1
247
+
248
+
249
+def test_routed_experts_attention_group_is_shared_and_fail_closed(monkeypatch):
250
+ class FullAttentionSpec:
251
+ pass
252
+
253
+ monkeypatch.setattr(f"{_REC_MODULE}.FullAttentionSpec", FullAttentionSpec)
254
+ config = SimpleNamespace(
255
+ kv_cache_groups=[
256
+ SimpleNamespace(kv_cache_spec=object()),
257
+ SimpleNamespace(kv_cache_spec=FullAttentionSpec()),
258
+ ]
259
+ )
260
+ assert get_routed_experts_attn_gid(config) == 1
261
+
262
+ with pytest.raises(ValueError, match="requires a full-attention KV cache group"):
263
+ get_routed_experts_attn_gid(SimpleNamespace(kv_cache_groups=[]))
264
+
265
+
266
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
267
+def test_mrv2_async_output_returns_existing_routed_experts_field():
268
+ from vllm.v1.outputs import ModelRunnerOutput, RoutedExpertsTensors
269
+ from vllm.v1.worker.gpu.async_utils import AsyncOutput
270
+ from vllm.v1.worker.gpu.sample.output import SamplerOutput
271
+
272
+ routed_experts = RoutedExpertsTensors(
273
+ routing_data=torch.arange(6, dtype=torch.int32, device="cuda").reshape(3, 1, 2),
274
+ slot_mapping=torch.tensor([11, 12, 13], device="cuda"),
275
+ )
276
+ num_sampled = torch.tensor([1], dtype=torch.int32, device="cuda")
277
+ sampler_output = SamplerOutput(
278
+ sampled_token_ids=torch.tensor([[1]], device="cuda"),
279
+ logprobs_tensors=None,
280
+ num_nans=None,
281
+ num_sampled=num_sampled,
282
+ num_rejected=torch.tensor([0], dtype=torch.int32, device="cuda"),
283
+ )
284
+ output = AsyncOutput(
285
+ model_runner_output=ModelRunnerOutput(req_ids=["req"], req_id_to_index={}),
286
+ sampler_output=sampler_output,
287
+ num_sampled_tokens=num_sampled,
288
+ main_stream=torch.cuda.current_stream(),
289
+ copy_stream=torch.cuda.Stream(),
290
+ routed_experts=routed_experts,
291
+ ).get_output()
292
+
293
+ assert output.routed_experts is not None
294
+ assert output.routed_experts.routing_data[:, 0, 0].tolist() == [0, 2, 4]
295
+ assert output.routed_experts.slot_mapping.tolist() == [11, 12, 13]
296
+
297
+
298
+@pytest.mark.parametrize("rank", [0, 1])
299
+def test_all_tp_ranks_initialize_capture(monkeypatch, rank):
300
+ pytest.importorskip("vllm.vllm_flash_attn", exc_type=ImportError)
301
+ import vllm.v1.worker.gpu.model_runner as model_runner
302
+
303
+ capturer = Mock()
304
+ constructor = Mock(return_value=capturer)
305
+ bind = Mock()
306
+ monkeypatch.setattr(model_runner, "RoutedExpertsCapturer", constructor)
307
+ monkeypatch.setattr(model_runner, "bind_routed_experts_capturer", bind)
308
+
309
+ runner = model_runner.GPUModelRunner.__new__(model_runner.GPUModelRunner)
310
+ runner.max_num_tokens = 32
311
+ runner.vllm_config = SimpleNamespace(parallel_config=SimpleNamespace(rank=rank))
312
+ runner.kv_cache_config = SimpleNamespace()
313
+ runner.model = Mock()
314
+
315
+ runner.init_routed_experts_capturer()
316
+
317
+ constructor.assert_called_once_with(
318
+ max_num_batched_tokens=32,
319
+ vllm_config=runner.vllm_config,
320
+ kv_cache_config=runner.kv_cache_config,
321
+ )
322
+ bind.assert_called_once_with(runner.model, capturer)
323
+ assert runner.routed_experts_capturer is capturer
324
+
325
+
326
+def test_v2_model_runner_accepts_routed_experts(monkeypatch):
327
+ monkeypatch.setattr("importlib.metadata.entry_points", lambda **_: ())
328
+ config = SimpleNamespace(
329
+ model_config=SimpleNamespace(
330
+ enable_return_routed_experts=True,
331
+ use_mla=False,
332
+ logits_processors=None,
333
+ enable_prompt_embeds=False,
334
+ ),
335
+ speculative_config=None,
336
+ parallel_config=SimpleNamespace(
337
+ prefill_context_parallel_size=1,
338
+ tensor_parallel_size=1,
339
+ distributed_executor_backend=None,
340
+ pipeline_parallel_size=1,
341
+ enable_dbo=False,
342
+ enable_elastic_ep=False,
343
+ ),
344
+ compilation_config=SimpleNamespace(
345
+ mode=CompilationMode.NONE,
346
+ pass_config=SimpleNamespace(enable_sp=False),
347
+ ),
348
+ cache_config=SimpleNamespace(kv_sharing_fast_prefill=False),
349
+ ec_transfer_config=None,
350
+ )
351
+
352
+ unsupported = VllmConfig._get_v2_model_runner_unsupported_features(config)
353
+
354
+ assert "routed experts capture" not in unsupported
tests/v1/executor/test_ray_utils.py
+31
-1
@@ -4,7 +4,12 @@
4
import numpy as np
5
6
from vllm.v1.executor.ray_utils import detach_zero_copy_from_model_runner_output
7
-from vllm.v1.outputs import LogprobsLists, LogprobsTensors, ModelRunnerOutput
7
+from vllm.v1.outputs import (
8
+ LogprobsLists,
9
+ LogprobsTensors,
10
+ ModelRunnerOutput,
11
+ RoutedExpertsLists,
12
+)
13
14
15
def _make_readonly(arr: np.ndarray) -> np.ndarray:
@@ -52,3 +57,28 @@ def test_detach_zero_copy_from_model_runner_output_copies_only_numpy_views():
57
assert detached_logprobs.sampled_token_ranks.flags.writeable
58
assert detached_logprobs.cu_num_generated_tokens is cu_num_generated_tokens
59
assert output.prompt_logprobs_dict["req-0"] is prompt_logprobs
60
+
61
+
62
+def test_detach_zero_copy_routed_experts_without_logprobs():
63
+ output = ModelRunnerOutput(
64
+ req_ids=["req-0"],
65
+ req_id_to_index={"req-0": 0},
66
+ routed_experts=RoutedExpertsLists(
67
+ routing_data=_make_readonly(np.arange(12, dtype=np.int32).reshape(2, 3, 2)),
68
+ slot_mapping=_make_readonly(np.array([7, 8], dtype=np.int64)),
69
+ ),
70
+ )
71
+ original = output.routed_experts
72
+ assert output.logprobs is None
73
+
74
+ detach_zero_copy_from_model_runner_output(output)
75
+
76
+ detached = output.routed_experts
77
+ assert detached is not None
78
+ assert detached is not original
79
+ assert detached.routing_data is not original.routing_data
80
+ assert detached.slot_mapping is not original.slot_mapping
81
+ assert detached.routing_data.flags.writeable
82
+ assert detached.slot_mapping.flags.writeable
83
+ np.testing.assert_array_equal(detached.routing_data, original.routing_data)
84
+ np.testing.assert_array_equal(detached.slot_mapping, original.slot_mapping)
tests/v1/worker/test_gpu_model_runner_v2_eplb.py
+1
@@ -184,6 +184,7 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):
184
hidden_states=None,
185
aux_hidden_states=None,
186
finished_req_ids=set(),
187
+ routed_experts=None,
188
num_tokens_across_dp=None,
189
)
190
runner.req_states = SimpleNamespace()
vllm/config/vllm.py
-4
@@ -2228,10 +2228,6 @@ class VllmConfig:
2228
if self.parallel_config.enable_elastic_ep:
2229
unsupported.append("elastic expert parallelism")
2230
2231
- if model_config is not None and model_config.enable_return_routed_experts:
2232
- # Will be added by https://github.com/vllm-project/vllm/pull/38163
2233
- unsupported.append("routed experts capture")
2234
-
2231
has_logitsproc_plugins = False
2232
if model_config is not None:
2233
from importlib.metadata import entry_points
vllm/model_executor/layers/fused_moe/routed_experts_capturer.py
+91
-25
@@ -15,6 +15,7 @@ from vllm.distributed.parallel_state import get_tp_group
15
from vllm.forward_context import get_forward_context
16
from vllm.platforms import current_platform
17
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig
18
+from vllm.v1.outputs import RoutedExpertsTensors
19
20
logger = logging.getLogger(__name__)
21
@@ -78,8 +79,7 @@ class RoutedExpertsCapturer:
79
Invariants:
80
- One instance per worker; shape is fixed at init and covers the
81
worst-case step (``max_num_batched_tokens`` tokens).
81
- - :meth:`clear_buffer` is called at the start of every step, so
82
- unused slots stay zero.
82
+ - Every routed layer overwrites the current step's token rows.
83
- ``device_buffer.dtype`` is ``torch.int32``.
84
"""
85
@@ -87,6 +87,7 @@ class RoutedExpertsCapturer:
87
self,
88
max_num_batched_tokens: int,
89
vllm_config: VllmConfig,
90
+ kv_cache_config: KVCacheConfig,
91
) -> None:
92
hf_config = vllm_config.model_config.hf_text_config
93
num_experts_per_tok = _get_num_experts_per_tok(hf_config)
@@ -116,6 +117,8 @@ class RoutedExpertsCapturer:
117
)
118
self.dp_rank = vllm_config.parallel_config.data_parallel_rank
119
self.tp_size = vllm_config.parallel_config.tensor_parallel_size
120
+ # KV cache group whose slot layout the routing data is keyed by.
121
+ self.attn_gid = get_routed_experts_attn_gid(kv_cache_config)
122
123
def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None:
124
"""Capture expert routing decisions for a specific layer.
@@ -205,30 +208,102 @@ class RoutedExpertsCapturer:
208
f"tp_size={self.tp_size})"
209
)
210
208
- # Defensive: model may expose more layers than the capture buffer
209
- # was sized for (unusual, but guards against miss-config).
211
if layer_id >= self.device_buffer.shape[1]:
211
- return
212
+ raise IndexError(
213
+ f"routed-experts layer {layer_id} exceeds capture buffer "
214
+ f"layer count {self.device_buffer.shape[1]}"
215
+ )
216
217
self.device_buffer[:token_num_per_dp, layer_id, :] = topk_ids[
218
start_loc:end_loc, :
219
]
220
217
- def clear_buffer(self) -> None:
218
- """Zero the device buffer. Called at the start of every step so
219
- slots belonging to finished / preempted tokens don't leak into
220
- the next step.
221
- """
222
- self.device_buffer.zero_()
223
-
221
def get_device_buffer(self) -> torch.Tensor:
222
"""Return the underlying device buffer so the model runner can
223
issue the D2H copy. The tensor is shared; callers must either
227
- clone or fully drain it before the next forward pass runs
228
- :meth:`clear_buffer`.
224
+ clone or fully drain it before the next forward pass overwrites it.
225
"""
226
return self.device_buffer
227
228
+ def get_routed_experts(
229
+ self, slot_mappings: torch.Tensor, num_tokens: int
230
+ ) -> RoutedExpertsTensors:
231
+ """Snapshot this step's routing data and its attention slot mapping.
232
+
233
+ Both tensors are cloned since the capture buffer and the shared
234
+ ``slot_mappings`` are overwritten by the next step, which may race
235
+ with an in-flight D2H copy.
236
+
237
+ Args:
238
+ slot_mappings: Per-KV-cache-group slot mappings for this step,
239
+ shape ``(num_kv_cache_groups, max_num_batched_tokens)``.
240
+ num_tokens: Total number of tokens scheduled in this step.
241
+ """
242
+ return RoutedExpertsTensors(
243
+ routing_data=self.device_buffer[:num_tokens].clone(),
244
+ slot_mapping=slot_mappings[self.attn_gid, :num_tokens].clone(),
245
+ )
246
+
247
+
248
+def bind_routed_experts_capturer(
249
+ model: torch.nn.Module,
250
+ capturer: RoutedExpertsCapturer,
251
+) -> None:
252
+ """Attach capture callbacks to the target model's MoE routers."""
253
+ from vllm.model_executor.layers.fused_moe.layer import MoERunner
254
+ from vllm.model_executor.layers.fused_moe.modular_kernel import (
255
+ FusedMoEExpertsMonolithic,
256
+ )
257
+ from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
258
+
259
+ num_bound = 0
260
+ for module in model.modules():
261
+ if not isinstance(module, MoERunner):
262
+ continue
263
+ layer_id = module.layer_id
264
+
265
+ def capture_fn(
266
+ topk_ids: torch.Tensor,
267
+ layer_id: int = layer_id,
268
+ capturer: RoutedExpertsCapturer = capturer,
269
+ ) -> None:
270
+ capturer.capture(layer_id, topk_ids)
271
+
272
+ quant_method = module._quant_method
273
+ moe_kernel = getattr(quant_method, "moe_kernel", None)
274
+ impl = getattr(moe_kernel, "impl", None)
275
+ fused_experts = getattr(impl, "fused_experts", None)
276
+ if quant_method.is_monolithic:
277
+ if not (
278
+ isinstance(fused_experts, FusedMoEExpertsMonolithic)
279
+ and fused_experts.supports_routing_replay_capture()
280
+ ):
281
+ raise ValueError(
282
+ "Routed-experts capture is not supported with monolithic "
283
+ f"MoE kernel {type(fused_experts).__name__}."
284
+ )
285
+ fused_experts.set_capture_fn(capture_fn)
286
+ num_bound += 1
287
+ elif isinstance(module.router, BaseRouter):
288
+ module.router.set_capture_fn(capture_fn)
289
+ num_bound += 1
290
+ else:
291
+ raise ValueError(
292
+ "Routed-experts capture is not supported with router "
293
+ f"{type(module.router).__name__}."
294
+ )
295
+
296
+ if num_bound == 0:
297
+ raise ValueError("No supported MoE router found for routed-experts capture.")
298
+
299
+
300
+def get_routed_experts_attn_gid(kv_cache_config: KVCacheConfig) -> int:
301
+ """Return the full-attention KV cache group used for routed experts."""
302
+ for gid, group in enumerate(kv_cache_config.kv_cache_groups):
303
+ if isinstance(group.kv_cache_spec, FullAttentionSpec):
304
+ return gid
305
+ raise ValueError("Routed-experts capture requires a full-attention KV cache group.")
306
+
307
308
class RoutedExpertsManager:
309
"""Scheduler-side slot-indexed buffer for routed experts.
@@ -260,17 +335,8 @@ class RoutedExpertsManager:
335
vllm_config: VllmConfig,
336
kv_cache_config: KVCacheConfig,
337
) -> None:
263
- # Pick the attention group for block/slot mapping. We require
264
- # a FullAttentionSpec group rather than any AttentionSpec to
265
- # stay consistent with the worker-side lookup in
266
- # ``GPUModelRunner._get_attention_kv_cache_gid``; hybrid models
267
- # (Mamba / linear attention) also have other AttentionSpec
268
- # groups whose slot layout differs.
269
- self.attn_gid = next(
270
- gid
271
- for gid, g in enumerate(kv_cache_config.kv_cache_groups)
272
- if isinstance(g.kv_cache_spec, FullAttentionSpec)
273
- )
338
+ # Hybrid models also have KV groups whose slot layout differs.
339
+ self.attn_gid = get_routed_experts_attn_gid(kv_cache_config)
340
attn_group = kv_cache_config.kv_cache_groups[self.attn_gid]
341
self.block_size = attn_group.kv_cache_spec.block_size
342
vllm/v1/executor/ray_utils.py
+28
-18
@@ -202,10 +202,10 @@ def detach_zero_copy_from_model_runner_output(output: "ModelRunnerOutput") -> No
202
backed by Ray's shared-memory object store. Ray's channel docs explicitly
203
warn that subsequent reads may block if such an object is still in scope.
204
205
- vLLM can return numpy-backed logprobs in `ModelRunnerOutput.logprobs`. If
206
- those arrays are backed by Ray SHM (commonly read-only), retaining them in
207
- scope across scheduler iterations can stall the channel and eventually hit
208
- `RAY_CGRAPH_get_timeout`.
205
+ vLLM can return numpy-backed logprobs and routed experts in
206
+ `ModelRunnerOutput`. If those arrays are backed by Ray SHM (commonly
207
+ read-only), retaining them in scope across scheduler iterations can stall
208
+ the channel and eventually hit `RAY_CGRAPH_get_timeout`.
209
210
Copy read-only numpy arrays so the returned output no longer retains
211
references to Ray's shared-memory buffers.
@@ -214,27 +214,37 @@ def detach_zero_copy_from_model_runner_output(output: "ModelRunnerOutput") -> No
214
`LogprobsTensors` backed by PyTorch-owned CPU tensors (`to_cpu_nonblocking`
215
or `empty_cpu`), not NumPy views decoded from Ray channels.
216
"""
217
- if output.logprobs is None:
218
- return
219
-
220
- token_ids, logprobs, ranks, cu_num_generated_tokens = output.logprobs
217
218
def _copy_if_readonly(arr):
219
if isinstance(arr, np.ndarray) and not arr.flags.writeable:
220
return arr.copy()
221
return arr
222
227
- # `cu_num_generated_tokens` is already a plain Python list (or None), so it
228
- # never aliases Ray SHM buffers and can be reused as-is.
229
- token_ids_c = _copy_if_readonly(token_ids)
230
- logprobs_c = _copy_if_readonly(logprobs)
231
- ranks_c = _copy_if_readonly(ranks)
232
- if token_ids_c is token_ids and logprobs_c is logprobs and ranks_c is ranks:
233
- return
223
+ if output.logprobs is not None:
224
+ token_ids, logprobs, ranks, cu_num_generated_tokens = output.logprobs
225
+
226
+ # `cu_num_generated_tokens` is already a plain Python list (or None),
227
+ # so it never aliases Ray SHM buffers and can be reused as-is.
228
+ token_ids_c = _copy_if_readonly(token_ids)
229
+ logprobs_c = _copy_if_readonly(logprobs)
230
+ ranks_c = _copy_if_readonly(ranks)
231
+ if (
232
+ token_ids_c is not token_ids
233
+ or logprobs_c is not logprobs
234
+ or ranks_c is not ranks
235
+ ):
236
+ output.logprobs = type(output.logprobs)(
237
+ token_ids_c, logprobs_c, ranks_c, cu_num_generated_tokens
238
+ )
239
235
- output.logprobs = type(output.logprobs)(
236
- token_ids_c, logprobs_c, ranks_c, cu_num_generated_tokens
237
- )
240
+ if output.routed_experts is not None:
241
+ routing_data, slot_mapping = output.routed_experts
242
+ routing_data_c = _copy_if_readonly(routing_data)
243
+ slot_mapping_c = _copy_if_readonly(slot_mapping)
244
+ if routing_data_c is not routing_data or slot_mapping_c is not slot_mapping:
245
+ output.routed_experts = type(output.routed_experts)(
246
+ routing_data_c, slot_mapping_c
247
+ )
248
249
250
class FutureWrapper(Future):
vllm/v1/worker/gpu/async_utils.py
+8
@@ -11,6 +11,7 @@ from vllm.v1.outputs import (
11
LogprobsTensors,
12
ModelRunnerOutput,
13
PoolerOutput,
14
+ RoutedExpertsTensors,
15
)
16
from vllm.v1.worker.gpu.sample.output import SamplerOutput
17
@@ -24,6 +25,7 @@ class AsyncOutput(AsyncModelRunnerOutput):
25
main_stream: torch.cuda.Stream,
26
copy_stream: torch.cuda.Stream,
27
check_ep_fault: bool = False,
28
+ routed_experts: RoutedExpertsTensors | None = None,
29
):
30
# NOTE(woosuk): We must retain references to the GPU tensors,
31
# as the copy operations are performed on a different CUDA stream than
@@ -31,6 +33,7 @@ class AsyncOutput(AsyncModelRunnerOutput):
33
self.model_runner_output = model_runner_output
34
self.sampler_output = sampler_output
35
self.num_sampled_tokens = num_sampled_tokens
36
+ self.routed_experts = routed_experts
37
# Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
38
self.copy_event = torch.cuda.Event(blocking=True)
39
self._has_fault: torch.Tensor | None = None
@@ -48,6 +51,9 @@ class AsyncOutput(AsyncModelRunnerOutput):
51
if sampler_output.num_nans is not None:
52
self.num_nans = async_copy_to_np(sampler_output.num_nans)
53
self.num_sampled_tokens_np = async_copy_to_np(num_sampled_tokens)
54
+ self.routed_experts_cpu: RoutedExpertsTensors | None = None
55
+ if routed_experts is not None:
56
+ self.routed_experts_cpu = routed_experts.to_cpu_nonblocking()
57
self.prompt_logprobs_dict = {
58
k: v.to_cpu_nonblocking() if v is not None else None
59
for k, v in self.model_runner_output.prompt_logprobs_dict.items()
@@ -78,6 +84,8 @@ class AsyncOutput(AsyncModelRunnerOutput):
84
if self.logprobs_tensors is not None:
85
self.model_runner_output.logprobs = self.logprobs_tensors.tolists()
86
self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict
87
+ if self.routed_experts_cpu is not None:
88
+ self.model_runner_output.routed_experts = self.routed_experts_cpu.tolists()
89
90
if self._has_fault is not None and self._has_fault.item():
91
mask = get_ep_all2all_manager().query_active_mask()
vllm/v1/worker/gpu/model_runner.py
+28
-1
@@ -39,6 +39,10 @@ from vllm.distributed.parallel_state import (
39
from vllm.forward_context import BatchDescriptor, set_forward_context
40
from vllm.logger import init_logger
41
from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
42
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
43
+ RoutedExpertsCapturer,
44
+ bind_routed_experts_capturer,
45
+)
46
from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
47
initialize_mamba_ssu_backend,
48
)
@@ -55,7 +59,11 @@ from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
59
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
60
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
61
from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
58
-from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
62
+from vllm.v1.outputs import (
63
+ DraftTokenIds,
64
+ ModelRunnerOutput,
65
+ RoutedExpertsTensors,
66
+)
67
from vllm.v1.worker.block_table import get_block_table_width
68
from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
69
from vllm.v1.worker.gpu import pcp_manager as pcp
@@ -279,11 +287,21 @@ class GPUModelRunner(LoRAModelRunnerMixin):
287
288
# Expert parallelism load balancer.
289
self.eplb = EPLBController(self.parallel_config, self.device)
290
+ self.routed_experts_capturer: RoutedExpertsCapturer | None = None
291
292
def update_max_model_len(self, max_model_len: int) -> None:
293
self.max_model_len = max_model_len
294
self.req_states.max_model_len = max_model_len
295
296
+ def init_routed_experts_capturer(self) -> None:
297
+ """Initialize target-model capture on every participating worker."""
298
+ self.routed_experts_capturer = RoutedExpertsCapturer(
299
+ max_num_batched_tokens=self.max_num_tokens,
300
+ vllm_config=self.vllm_config,
301
+ kv_cache_config=self.kv_cache_config,
302
+ )
303
+ bind_routed_experts_capturer(self.model, self.routed_experts_capturer)
304
+
305
def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
306
tasks: list[SupportedTask] = []
307
if self.model_config.runner_type == "generate":
@@ -1437,6 +1455,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1455
aux_hidden_states = None
1456
output_intermediate_tensors = model_output
1457
1458
+ routed_experts = None
1459
+ if not dummy_run and (capturer := self.routed_experts_capturer) is not None:
1460
+ assert slot_mappings is not None
1461
+ routed_experts = capturer.get_routed_experts(slot_mappings, num_toks)
1462
+
1463
finished_req_ids = scheduler_output.finished_req_ids
1464
self.execute_model_state = ExecuteModelState(
1465
input_batch=input_batch,
@@ -1445,6 +1468,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1468
hidden_states=hidden_states,
1469
aux_hidden_states=aux_hidden_states,
1470
finished_req_ids=finished_req_ids,
1471
+ routed_experts=routed_experts,
1472
)
1473
1474
if not self.is_last_pp_rank:
@@ -1467,6 +1491,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1491
hidden_states = self.execute_model_state.hidden_states
1492
aux_hidden_states = self.execute_model_state.aux_hidden_states
1493
finished_req_ids = self.execute_model_state.finished_req_ids
1494
+ routed_experts = self.execute_model_state.routed_experts
1495
self.execute_model_state = None
1496
1497
if not self.is_last_pp_rank:
@@ -1532,6 +1557,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1557
main_stream=self.main_stream,
1558
copy_stream=self.output_copy_stream,
1559
check_ep_fault=self.check_ep_fault,
1560
+ routed_experts=routed_experts,
1561
)
1562
1563
mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
@@ -1712,6 +1738,7 @@ class ExecuteModelState(NamedTuple):
1738
hidden_states: torch.Tensor | None
1739
aux_hidden_states: list[torch.Tensor] | None
1740
finished_req_ids: set[str]
1741
+ routed_experts: RoutedExpertsTensors | None
1742
1743
1744
def sort_batch_req_ids(
vllm/v1/worker/gpu_model_runner.py
+19
-66
@@ -63,6 +63,7 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
63
from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
64
from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
65
RoutedExpertsCapturer,
66
+ bind_routed_experts_capturer,
67
)
68
from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
69
initialize_mamba_ssu_backend,
@@ -2350,8 +2351,7 @@ class GPUModelRunner(
2351
# by the next ``_prepare_inputs``; we need a stable snapshot
2352
# because the async D2H may still be in flight on the copy
2353
# stream when the next step runs.
2353
- attn_gid = self.routed_experts_attn_gid
2354
- slot_mapping_attn = slot_mappings[attn_gid]
2354
+ slot_mapping_attn = slot_mappings[self.routed_experts_capturer.attn_gid]
2355
self.routed_experts_slot_mapping_device[:num_tokens].copy_(
2356
slot_mapping_attn[:num_tokens]
2357
)
@@ -4173,9 +4173,6 @@ class GPUModelRunner(
4173
"after execute_model() returns None."
4174
)
4175
4176
- if self.routed_experts_initialized:
4177
- self.routed_experts_capturer.clear_buffer()
4178
-
4176
# If ngram_gpu is used, we need to copy the scheduler_output to avoid
4177
# the modification has influence on the scheduler_output in engine core process.
4178
# The replace is much faster than deepcopy.
@@ -4789,24 +4786,16 @@ class GPUModelRunner(
4786
# copy stream can D2H later. Both tensors must be private
4787
# clones because:
4788
# - ``routing_data`` source is the shared capturer buffer,
4792
- # which is ``clear_buffer()``-ed at the start of the
4793
- # next step on the default stream.
4789
+ # which the next forward overwrites on the default stream.
4790
# - ``slot_mapping`` source is our own
4791
# ``routed_experts_slot_mapping_device``, which the
4792
# next ``_prepare_inputs`` overwrites on the default
4793
# stream while the D2H is still pending on the copy
4794
# stream.
4795
# Without clones, the copy stream would read torn data.
4800
- routed_experts_snapshot = None
4801
- if self.routed_experts_initialized:
4802
- buf = self.routed_experts_capturer.get_device_buffer()
4803
- total = scheduler_output.total_num_scheduled_tokens
4804
- routed_experts_snapshot = RoutedExpertsTensors(
4805
- routing_data=buf[:total].clone(),
4806
- slot_mapping=self.routed_experts_slot_mapping_device[
4807
- :total
4808
- ].clone(),
4809
- )
4796
+ routed_experts_snapshot = self.get_routed_experts(
4797
+ scheduler_output.total_num_scheduled_tokens
4798
+ )
4799
4800
async_output = AsyncGPUModelRunnerOutput(
4801
model_runner_output=output,
@@ -7663,19 +7652,18 @@ class GPUModelRunner(
7652
kv_transfer_group.register_kv_caches(kv_caches)
7653
kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks)
7654
7666
- def _get_attention_kv_cache_gid(self) -> int:
7667
- """Find the KV cache group index for attention layers.
7655
+ def get_routed_experts(
7656
+ self,
7657
+ num_tokens: int,
7658
+ ) -> RoutedExpertsTensors | None:
7659
+ if not self.routed_experts_initialized:
7660
+ return None
7661
7669
- Must match :attr:`RoutedExpertsManager.attn_gid` in the scheduler:
7670
- both pick the first ``FullAttentionSpec`` group so hybrid models
7671
- (Mamba / linear-attention layers that use other AttentionSpec
7672
- subclasses) end up indexing the same slot layout on both sides.
7673
- Falls back to 0 only for legacy single-group configs.
7674
- """
7675
- for gid, group in enumerate(self.kv_cache_config.kv_cache_groups):
7676
- if isinstance(group.kv_cache_spec, FullAttentionSpec):
7677
- return gid
7678
- return 0
7662
+ device_buffer = self.routed_experts_capturer.get_device_buffer()
7663
+ return RoutedExpertsTensors(
7664
+ routing_data=device_buffer[:num_tokens].clone(),
7665
+ slot_mapping=self.routed_experts_slot_mapping_device[:num_tokens].clone(),
7666
+ )
7667
7668
def init_routed_experts_capturer(self):
7669
logger.info(
@@ -7685,9 +7673,9 @@ class GPUModelRunner(
7673
self.routed_experts_capturer = RoutedExpertsCapturer(
7674
max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens,
7675
vllm_config=self.vllm_config,
7676
+ kv_cache_config=self.kv_cache_config,
7677
)
7689
- self.routed_experts_attn_gid = self._get_attention_kv_cache_gid()
7690
- self._bind_routed_experts_capturer(self.routed_experts_capturer)
7678
+ bind_routed_experts_capturer(self.model, self.routed_experts_capturer)
7679
7680
# Pinned CPU buffer for non-blocking D2H of ``routing_data`` on
7681
# the sync scheduling path. Shape / dtype mirror the device
@@ -7719,41 +7707,6 @@ class GPUModelRunner(
7707
)
7708
self.routed_experts_initialized = True
7709
7722
- def _bind_routed_experts_capturer(self, capturer: RoutedExpertsCapturer) -> None:
7723
- from vllm.model_executor.layers.fused_moe.layer import MoERunner
7724
- from vllm.model_executor.layers.fused_moe.modular_kernel import (
7725
- FusedMoEExpertsMonolithic,
7726
- )
7727
- from vllm.model_executor.layers.fused_moe.router.base_router import (
7728
- BaseRouter,
7729
- )
7730
-
7731
- for module in self.model.modules():
7732
- if not isinstance(module, MoERunner):
7733
- continue
7734
- layer_id = module.layer_id
7735
-
7736
- def _capture_fn(topk_ids, _layer_id=layer_id, _capturer=capturer):
7737
- _capturer.capture(_layer_id, topk_ids)
7738
-
7739
- quant_method = module._quant_method
7740
- moe_kernel = getattr(quant_method, "moe_kernel", None)
7741
- impl = getattr(moe_kernel, "impl", None)
7742
- fused_experts = getattr(impl, "fused_experts", None)
7743
- if quant_method.is_monolithic:
7744
- if not (
7745
- isinstance(fused_experts, FusedMoEExpertsMonolithic)
7746
- and fused_experts.supports_routing_replay_capture()
7747
- ):
7748
- raise ValueError(
7749
- "--enable-return-routed-experts is not supported with "
7750
- f"monolithic MoE kernel {type(fused_experts).__name__}; "
7751
- "routed expert IDs would be silently all-zero."
7752
- )
7753
- fused_experts.set_capture_fn(_capture_fn)
7754
- elif isinstance(module.router, BaseRouter):
7755
- module.router.set_capture_fn(_capture_fn)
7756
-
7710
def may_add_encoder_only_layers_to_kv_cache_config(self) -> None:
7711
"""
7712
Add encoder-only layers to the KV cache config.