[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>

tests/kernels/moe/test_routed_experts_capture_monolithic.py

@@ -385,11 +385,9 @@ def test_routed_experts_capturer_e2e_via_monolithic_experts() -> None:
     on the monolithic experts and verify the captured rows land in the
     capturer's device buffer at the correct layer slot.
 
-    Mirrors the wiring done in ``GPUModelRunner._bind_routed_experts_capturer``
-    for the monolithic path: a single closure is installed on the monolithic
-    ``fused_experts`` (in addition to ``router.set_capture_fn`` on the
-    non-monolithic path) and the capturer routes per-layer based on the
-    closed-over ``layer_id``.
+    Mirrors the monolithic path in ``bind_routed_experts_capturer``: a single
+    closure is installed on ``fused_experts`` and routes per-layer based on
+    the closed-over ``layer_id``.
     """
     from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
         RoutedExpertsCapturer,

tests/model_executor/test_routed_experts_capture.py

@@ -2,15 +2,19 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import types
 from types import SimpleNamespace
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import pytest
 import torch
 
+from vllm.config import VllmConfig
+from vllm.config.compilation import CompilationMode
 from vllm.distributed.eplb.eplb_state import EplbLayerState
 from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
 from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
     RoutedExpertsCapturer,
+    bind_routed_experts_capturer,
+    get_routed_experts_attn_gid,
 )
 from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
 
@@ -118,97 +122,12 @@ def test_base_router_capture_with_eplb_enabled():
     assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]]))
 
 
-def test_gpu_model_runner_binds_router_capture(monkeypatch):
-    from vllm.v1.worker import gpu_model_runner as gmr
-
-    class _DummyRouter:
-        _routing_replay_out: torch.Tensor | None = None
-
-    class DummyFusedMoE:
-        def __init__(self):
-            self.layer_id = 7
-            self.router = _make_router()
-            self.routed_experts = _make_modular_routed_experts()
-            self._quant_method = self.routed_experts.quant_method
-
-    class DummyCapturer:
-        def __init__(self):
-            self.calls = []
-
-        def capture(self, layer_id, topk_ids):
-            self.calls.append((layer_id, topk_ids))
-
-    dummy_module = DummyFusedMoE()
-
-    # Patch the runtime import inside _bind_routed_experts_capturer.
-    import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
-
-    monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
-
-    dummy_self = types.SimpleNamespace(
-        model=types.SimpleNamespace(modules=lambda: [dummy_module])
-    )
-
-    capturer = DummyCapturer()
-    gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, capturer)
-
-    assert dummy_module.router.capture_fn is not None
-    dummy_module.router.capture_fn(torch.tensor([[5, 6]]))
-
-    assert len(capturer.calls) == 1
-    layer_id, topk_ids = capturer.calls[0]
-    assert layer_id == 7
-    assert torch.equal(topk_ids, torch.tensor([[5, 6]]))
-
-
-def test_gpu_model_runner_binding_stage(monkeypatch):
-    from vllm.v1.worker import gpu_model_runner as gmr
-
-    class DummyFusedMoE:
-        def __init__(self):
-            self.layer_id = 11
-            self.router = _make_router()
-            self.routed_experts = _make_modular_routed_experts()
-            self._quant_method = self.routed_experts.quant_method
-
-    class DummyCapturer:
-        def __init__(self):
-            self.calls = []
-
-        def capture(self, layer_id, topk_ids):
-            self.calls.append((layer_id, topk_ids))
-
-    dummy_module = DummyFusedMoE()
-
-    import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
-
-    monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
-
-    dummy_self = types.SimpleNamespace(
-        model=types.SimpleNamespace(modules=lambda: [dummy_module])
-    )
-
-    # Before binding, no capture hook.
-    assert dummy_module.router.capture_fn is None
-
-    capturer = DummyCapturer()
-    gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, capturer)
-
-    # After binding, hook should exist and be callable.
-    assert callable(dummy_module.router.capture_fn)
-    dummy_module.router.capture_fn(torch.tensor([[9, 10]]))
-    assert len(capturer.calls) == 1
-
-
-def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
-    from vllm.v1.worker import gpu_model_runner as gmr
-
+def test_public_binding_only_visits_target_model(monkeypatch):
     class DummyFusedMoE:
         def __init__(self, layer_id):
             self.layer_id = layer_id
             self.router = _make_router()
-            self.routed_experts = _make_modular_routed_experts()
-            self._quant_method = self.routed_experts.quant_method
+            self._quant_method = _make_modular_routed_experts().quant_method
 
     target_module = DummyFusedMoE(layer_id=7)
     draft_module = DummyFusedMoE(layer_id=0)
@@ -216,27 +135,21 @@ def test_gpu_model_runner_does_not_bind_draft_router_capture(monkeypatch):
     import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
 
     monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
+    calls = []
+    capturer = types.SimpleNamespace(capture=lambda *args: calls.append(args))
 
-    dummy_self = types.SimpleNamespace(
-        model=types.SimpleNamespace(modules=lambda: [target_module]),
-        compilation_config=types.SimpleNamespace(
-            static_forward_context={
-                "model.layers.7.mlp.experts": target_module,
-                "mtp.layers.0.mlp.experts": draft_module,
-            }
-        ),
+    bind_routed_experts_capturer(
+        types.SimpleNamespace(modules=lambda: [target_module]), capturer
     )
 
-    capturer = types.SimpleNamespace(capture=lambda *_: None)
-    gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, capturer)
-
     assert target_module.router.capture_fn is not None
     assert draft_module.router.capture_fn is None
+    topk_ids = torch.tensor([[5, 6]])
+    target_module.router.capture_fn(topk_ids)
+    assert calls == [(7, topk_ids)]
 
 
-def test_gpu_model_runner_rejects_monolithic_without_replay_support(monkeypatch):
-    from vllm.v1.worker import gpu_model_runner as gmr
-
+def test_public_binding_rejects_monolithic_without_replay_support(monkeypatch):
     class DummyFusedMoE:
         def __init__(self):
             self.layer_id = 3
@@ -269,12 +182,10 @@ def test_gpu_model_runner_rejects_monolithic_without_replay_support(monkeypatch)
 
     monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
 
-    dummy_self = types.SimpleNamespace(
-        model=types.SimpleNamespace(modules=lambda: [dummy_module])
-    )
-
     with pytest.raises(ValueError, match="monolithic MoE kernel"):
-        gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, DummyCapturer())
+        bind_routed_experts_capturer(
+            types.SimpleNamespace(modules=lambda: [dummy_module]), DummyCapturer()
+        )
 
 
 def test_routed_experts_capturer_single_dp_no_metadata():
@@ -333,3 +244,111 @@ def test_routed_experts_capturer_dp_unexpected_batch_raises():
     ):
         capturer.capture(layer_id=0, topk_ids=topk)
     assert capturer.device_buffer[0, 0, 0].item() == -1
+
+
+def test_routed_experts_attention_group_is_shared_and_fail_closed(monkeypatch):
+    class FullAttentionSpec:
+        pass
+
+    monkeypatch.setattr(f"{_REC_MODULE}.FullAttentionSpec", FullAttentionSpec)
+    config = SimpleNamespace(
+        kv_cache_groups=[
+            SimpleNamespace(kv_cache_spec=object()),
+            SimpleNamespace(kv_cache_spec=FullAttentionSpec()),
+        ]
+    )
+    assert get_routed_experts_attn_gid(config) == 1
+
+    with pytest.raises(ValueError, match="requires a full-attention KV cache group"):
+        get_routed_experts_attn_gid(SimpleNamespace(kv_cache_groups=[]))
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_mrv2_async_output_returns_existing_routed_experts_field():
+    from vllm.v1.outputs import ModelRunnerOutput, RoutedExpertsTensors
+    from vllm.v1.worker.gpu.async_utils import AsyncOutput
+    from vllm.v1.worker.gpu.sample.output import SamplerOutput
+
+    routed_experts = RoutedExpertsTensors(
+        routing_data=torch.arange(6, dtype=torch.int32, device="cuda").reshape(3, 1, 2),
+        slot_mapping=torch.tensor([11, 12, 13], device="cuda"),
+    )
+    num_sampled = torch.tensor([1], dtype=torch.int32, device="cuda")
+    sampler_output = SamplerOutput(
+        sampled_token_ids=torch.tensor([[1]], device="cuda"),
+        logprobs_tensors=None,
+        num_nans=None,
+        num_sampled=num_sampled,
+        num_rejected=torch.tensor([0], dtype=torch.int32, device="cuda"),
+    )
+    output = AsyncOutput(
+        model_runner_output=ModelRunnerOutput(req_ids=["req"], req_id_to_index={}),
+        sampler_output=sampler_output,
+        num_sampled_tokens=num_sampled,
+        main_stream=torch.cuda.current_stream(),
+        copy_stream=torch.cuda.Stream(),
+        routed_experts=routed_experts,
+    ).get_output()
+
+    assert output.routed_experts is not None
+    assert output.routed_experts.routing_data[:, 0, 0].tolist() == [0, 2, 4]
+    assert output.routed_experts.slot_mapping.tolist() == [11, 12, 13]
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_all_tp_ranks_initialize_capture(monkeypatch, rank):
+    pytest.importorskip("vllm.vllm_flash_attn", exc_type=ImportError)
+    import vllm.v1.worker.gpu.model_runner as model_runner
+
+    capturer = Mock()
+    constructor = Mock(return_value=capturer)
+    bind = Mock()
+    monkeypatch.setattr(model_runner, "RoutedExpertsCapturer", constructor)
+    monkeypatch.setattr(model_runner, "bind_routed_experts_capturer", bind)
+
+    runner = model_runner.GPUModelRunner.__new__(model_runner.GPUModelRunner)
+    runner.max_num_tokens = 32
+    runner.vllm_config = SimpleNamespace(parallel_config=SimpleNamespace(rank=rank))
+    runner.kv_cache_config = SimpleNamespace()
+    runner.model = Mock()
+
+    runner.init_routed_experts_capturer()
+
+    constructor.assert_called_once_with(
+        max_num_batched_tokens=32,
+        vllm_config=runner.vllm_config,
+        kv_cache_config=runner.kv_cache_config,
+    )
+    bind.assert_called_once_with(runner.model, capturer)
+    assert runner.routed_experts_capturer is capturer
+
+
+def test_v2_model_runner_accepts_routed_experts(monkeypatch):
+    monkeypatch.setattr("importlib.metadata.entry_points", lambda **_: ())
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            enable_return_routed_experts=True,
+            use_mla=False,
+            logits_processors=None,
+            enable_prompt_embeds=False,
+        ),
+        speculative_config=None,
+        parallel_config=SimpleNamespace(
+            prefill_context_parallel_size=1,
+            tensor_parallel_size=1,
+            distributed_executor_backend=None,
+            pipeline_parallel_size=1,
+            enable_dbo=False,
+            enable_elastic_ep=False,
+        ),
+        compilation_config=SimpleNamespace(
+            mode=CompilationMode.NONE,
+            pass_config=SimpleNamespace(enable_sp=False),
+        ),
+        cache_config=SimpleNamespace(kv_sharing_fast_prefill=False),
+        ec_transfer_config=None,
+    )
+
+    unsupported = VllmConfig._get_v2_model_runner_unsupported_features(config)
+
+    assert "routed experts capture" not in unsupported

tests/v1/executor/test_ray_utils.py

@@ -4,7 +4,12 @@
 import numpy as np
 
 from vllm.v1.executor.ray_utils import detach_zero_copy_from_model_runner_output
-from vllm.v1.outputs import LogprobsLists, LogprobsTensors, ModelRunnerOutput
+from vllm.v1.outputs import (
+    LogprobsLists,
+    LogprobsTensors,
+    ModelRunnerOutput,
+    RoutedExpertsLists,
+)
 
 
 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():
     assert detached_logprobs.sampled_token_ranks.flags.writeable
     assert detached_logprobs.cu_num_generated_tokens is cu_num_generated_tokens
     assert output.prompt_logprobs_dict["req-0"] is prompt_logprobs
+
+
+def test_detach_zero_copy_routed_experts_without_logprobs():
+    output = ModelRunnerOutput(
+        req_ids=["req-0"],
+        req_id_to_index={"req-0": 0},
+        routed_experts=RoutedExpertsLists(
+            routing_data=_make_readonly(np.arange(12, dtype=np.int32).reshape(2, 3, 2)),
+            slot_mapping=_make_readonly(np.array([7, 8], dtype=np.int64)),
+        ),
+    )
+    original = output.routed_experts
+    assert output.logprobs is None
+
+    detach_zero_copy_from_model_runner_output(output)
+
+    detached = output.routed_experts
+    assert detached is not None
+    assert detached is not original
+    assert detached.routing_data is not original.routing_data
+    assert detached.slot_mapping is not original.slot_mapping
+    assert detached.routing_data.flags.writeable
+    assert detached.slot_mapping.flags.writeable
+    np.testing.assert_array_equal(detached.routing_data, original.routing_data)
+    np.testing.assert_array_equal(detached.slot_mapping, original.slot_mapping)

tests/v1/worker/test_gpu_model_runner_v2_eplb.py

@@ -184,6 +184,7 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):
         hidden_states=None,
         aux_hidden_states=None,
         finished_req_ids=set(),
+        routed_experts=None,
         num_tokens_across_dp=None,
     )
     runner.req_states = SimpleNamespace()

vllm/config/vllm.py

@@ -2228,10 +2228,6 @@ class VllmConfig:
         if self.parallel_config.enable_elastic_ep:
             unsupported.append("elastic expert parallelism")
 
-        if model_config is not None and model_config.enable_return_routed_experts:
-            # Will be added by https://github.com/vllm-project/vllm/pull/38163
-            unsupported.append("routed experts capture")
-
         has_logitsproc_plugins = False
         if model_config is not None:
             from importlib.metadata import entry_points

vllm/model_executor/layers/fused_moe/routed_experts_capturer.py

@@ -15,6 +15,7 @@ from vllm.distributed.parallel_state import get_tp_group
 from vllm.forward_context import get_forward_context
 from vllm.platforms import current_platform
 from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig
+from vllm.v1.outputs import RoutedExpertsTensors
 
 logger = logging.getLogger(__name__)
 
@@ -78,8 +79,7 @@ class RoutedExpertsCapturer:
     Invariants:
         - One instance per worker; shape is fixed at init and covers the
           worst-case step (``max_num_batched_tokens`` tokens).
-        - :meth:`clear_buffer` is called at the start of every step, so
-          unused slots stay zero.
+        - Every routed layer overwrites the current step's token rows.
         - ``device_buffer.dtype`` is ``torch.int32``.
     """
 
@@ -87,6 +87,7 @@ class RoutedExpertsCapturer:
         self,
         max_num_batched_tokens: int,
         vllm_config: VllmConfig,
+        kv_cache_config: KVCacheConfig,
     ) -> None:
         hf_config = vllm_config.model_config.hf_text_config
         num_experts_per_tok = _get_num_experts_per_tok(hf_config)
@@ -116,6 +117,8 @@ class RoutedExpertsCapturer:
         )
         self.dp_rank = vllm_config.parallel_config.data_parallel_rank
         self.tp_size = vllm_config.parallel_config.tensor_parallel_size
+        # KV cache group whose slot layout the routing data is keyed by.
+        self.attn_gid = get_routed_experts_attn_gid(kv_cache_config)
 
     def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None:
         """Capture expert routing decisions for a specific layer.
@@ -205,30 +208,102 @@ class RoutedExpertsCapturer:
                     f"tp_size={self.tp_size})"
                 )
 
-        # Defensive: model may expose more layers than the capture buffer
-        # was sized for (unusual, but guards against miss-config).
         if layer_id >= self.device_buffer.shape[1]:
-            return
+            raise IndexError(
+                f"routed-experts layer {layer_id} exceeds capture buffer "
+                f"layer count {self.device_buffer.shape[1]}"
+            )
 
         self.device_buffer[:token_num_per_dp, layer_id, :] = topk_ids[
             start_loc:end_loc, :
         ]
 
-    def clear_buffer(self) -> None:
-        """Zero the device buffer. Called at the start of every step so
-        slots belonging to finished / preempted tokens don't leak into
-        the next step.
-        """
-        self.device_buffer.zero_()
-
     def get_device_buffer(self) -> torch.Tensor:
         """Return the underlying device buffer so the model runner can
         issue the D2H copy. The tensor is shared; callers must either
-        clone or fully drain it before the next forward pass runs
-        :meth:`clear_buffer`.
+        clone or fully drain it before the next forward pass overwrites it.
         """
         return self.device_buffer
 
+    def get_routed_experts(
+        self, slot_mappings: torch.Tensor, num_tokens: int
+    ) -> RoutedExpertsTensors:
+        """Snapshot this step's routing data and its attention slot mapping.
+
+        Both tensors are cloned since the capture buffer and the shared
+        ``slot_mappings`` are overwritten by the next step, which may race
+        with an in-flight D2H copy.
+
+        Args:
+            slot_mappings: Per-KV-cache-group slot mappings for this step,
+                shape ``(num_kv_cache_groups, max_num_batched_tokens)``.
+            num_tokens: Total number of tokens scheduled in this step.
+        """
+        return RoutedExpertsTensors(
+            routing_data=self.device_buffer[:num_tokens].clone(),
+            slot_mapping=slot_mappings[self.attn_gid, :num_tokens].clone(),
+        )
+
+
+def bind_routed_experts_capturer(
+    model: torch.nn.Module,
+    capturer: RoutedExpertsCapturer,
+) -> None:
+    """Attach capture callbacks to the target model's MoE routers."""
+    from vllm.model_executor.layers.fused_moe.layer import MoERunner
+    from vllm.model_executor.layers.fused_moe.modular_kernel import (
+        FusedMoEExpertsMonolithic,
+    )
+    from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
+
+    num_bound = 0
+    for module in model.modules():
+        if not isinstance(module, MoERunner):
+            continue
+        layer_id = module.layer_id
+
+        def capture_fn(
+            topk_ids: torch.Tensor,
+            layer_id: int = layer_id,
+            capturer: RoutedExpertsCapturer = capturer,
+        ) -> None:
+            capturer.capture(layer_id, topk_ids)
+
+        quant_method = module._quant_method
+        moe_kernel = getattr(quant_method, "moe_kernel", None)
+        impl = getattr(moe_kernel, "impl", None)
+        fused_experts = getattr(impl, "fused_experts", None)
+        if quant_method.is_monolithic:
+            if not (
+                isinstance(fused_experts, FusedMoEExpertsMonolithic)
+                and fused_experts.supports_routing_replay_capture()
+            ):
+                raise ValueError(
+                    "Routed-experts capture is not supported with monolithic "
+                    f"MoE kernel {type(fused_experts).__name__}."
+                )
+            fused_experts.set_capture_fn(capture_fn)
+            num_bound += 1
+        elif isinstance(module.router, BaseRouter):
+            module.router.set_capture_fn(capture_fn)
+            num_bound += 1
+        else:
+            raise ValueError(
+                "Routed-experts capture is not supported with router "
+                f"{type(module.router).__name__}."
+            )
+
+    if num_bound == 0:
+        raise ValueError("No supported MoE router found for routed-experts capture.")
+
+
+def get_routed_experts_attn_gid(kv_cache_config: KVCacheConfig) -> int:
+    """Return the full-attention KV cache group used for routed experts."""
+    for gid, group in enumerate(kv_cache_config.kv_cache_groups):
+        if isinstance(group.kv_cache_spec, FullAttentionSpec):
+            return gid
+    raise ValueError("Routed-experts capture requires a full-attention KV cache group.")
+
 
 class RoutedExpertsManager:
     """Scheduler-side slot-indexed buffer for routed experts.
@@ -260,17 +335,8 @@ class RoutedExpertsManager:
         vllm_config: VllmConfig,
         kv_cache_config: KVCacheConfig,
     ) -> None:
-        # Pick the attention group for block/slot mapping. We require
-        # a FullAttentionSpec group rather than any AttentionSpec to
-        # stay consistent with the worker-side lookup in
-        # ``GPUModelRunner._get_attention_kv_cache_gid``; hybrid models
-        # (Mamba / linear attention) also have other AttentionSpec
-        # groups whose slot layout differs.
-        self.attn_gid = next(
-            gid
-            for gid, g in enumerate(kv_cache_config.kv_cache_groups)
-            if isinstance(g.kv_cache_spec, FullAttentionSpec)
-        )
+        # Hybrid models also have KV groups whose slot layout differs.
+        self.attn_gid = get_routed_experts_attn_gid(kv_cache_config)
         attn_group = kv_cache_config.kv_cache_groups[self.attn_gid]
         self.block_size = attn_group.kv_cache_spec.block_size
 

vllm/v1/executor/ray_utils.py

@@ -202,10 +202,10 @@ def detach_zero_copy_from_model_runner_output(output: "ModelRunnerOutput") -> No
     backed by Ray's shared-memory object store. Ray's channel docs explicitly
     warn that subsequent reads may block if such an object is still in scope.
 
-    vLLM can return numpy-backed logprobs in `ModelRunnerOutput.logprobs`. If
-    those arrays are backed by Ray SHM (commonly read-only), retaining them in
-    scope across scheduler iterations can stall the channel and eventually hit
-    `RAY_CGRAPH_get_timeout`.
+    vLLM can return numpy-backed logprobs and routed experts in
+    `ModelRunnerOutput`. If those arrays are backed by Ray SHM (commonly
+    read-only), retaining them in scope across scheduler iterations can stall
+    the channel and eventually hit `RAY_CGRAPH_get_timeout`.
 
     Copy read-only numpy arrays so the returned output no longer retains
     references to Ray's shared-memory buffers.
@@ -214,27 +214,37 @@ def detach_zero_copy_from_model_runner_output(output: "ModelRunnerOutput") -> No
     `LogprobsTensors` backed by PyTorch-owned CPU tensors (`to_cpu_nonblocking`
     or `empty_cpu`), not NumPy views decoded from Ray channels.
     """
-    if output.logprobs is None:
-        return
-
-    token_ids, logprobs, ranks, cu_num_generated_tokens = output.logprobs
 
     def _copy_if_readonly(arr):
         if isinstance(arr, np.ndarray) and not arr.flags.writeable:
             return arr.copy()
         return arr
 
-    # `cu_num_generated_tokens` is already a plain Python list (or None), so it
-    # never aliases Ray SHM buffers and can be reused as-is.
-    token_ids_c = _copy_if_readonly(token_ids)
-    logprobs_c = _copy_if_readonly(logprobs)
-    ranks_c = _copy_if_readonly(ranks)
-    if token_ids_c is token_ids and logprobs_c is logprobs and ranks_c is ranks:
-        return
+    if output.logprobs is not None:
+        token_ids, logprobs, ranks, cu_num_generated_tokens = output.logprobs
 
-    output.logprobs = type(output.logprobs)(
-        token_ids_c, logprobs_c, ranks_c, cu_num_generated_tokens
-    )
+        # `cu_num_generated_tokens` is already a plain Python list (or None),
+        # so it never aliases Ray SHM buffers and can be reused as-is.
+        token_ids_c = _copy_if_readonly(token_ids)
+        logprobs_c = _copy_if_readonly(logprobs)
+        ranks_c = _copy_if_readonly(ranks)
+        if (
+            token_ids_c is not token_ids
+            or logprobs_c is not logprobs
+            or ranks_c is not ranks
+        ):
+            output.logprobs = type(output.logprobs)(
+                token_ids_c, logprobs_c, ranks_c, cu_num_generated_tokens
+            )
+
+    if output.routed_experts is not None:
+        routing_data, slot_mapping = output.routed_experts
+        routing_data_c = _copy_if_readonly(routing_data)
+        slot_mapping_c = _copy_if_readonly(slot_mapping)
+        if routing_data_c is not routing_data or slot_mapping_c is not slot_mapping:
+            output.routed_experts = type(output.routed_experts)(
+                routing_data_c, slot_mapping_c
+            )
 
 
 class FutureWrapper(Future):

vllm/v1/worker/gpu/async_utils.py

@@ -11,6 +11,7 @@ from vllm.v1.outputs import (
     LogprobsTensors,
     ModelRunnerOutput,
     PoolerOutput,
+    RoutedExpertsTensors,
 )
 from vllm.v1.worker.gpu.sample.output import SamplerOutput
 
@@ -24,6 +25,7 @@ class AsyncOutput(AsyncModelRunnerOutput):
         main_stream: torch.cuda.Stream,
         copy_stream: torch.cuda.Stream,
         check_ep_fault: bool = False,
+        routed_experts: RoutedExpertsTensors | None = None,
     ):
         # NOTE(woosuk): We must retain references to the GPU tensors,
         # as the copy operations are performed on a different CUDA stream than
@@ -31,6 +33,7 @@ class AsyncOutput(AsyncModelRunnerOutput):
         self.model_runner_output = model_runner_output
         self.sampler_output = sampler_output
         self.num_sampled_tokens = num_sampled_tokens
+        self.routed_experts = routed_experts
         # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
         self.copy_event = torch.cuda.Event(blocking=True)
         self._has_fault: torch.Tensor | None = None
@@ -48,6 +51,9 @@ class AsyncOutput(AsyncModelRunnerOutput):
             if sampler_output.num_nans is not None:
                 self.num_nans = async_copy_to_np(sampler_output.num_nans)
             self.num_sampled_tokens_np = async_copy_to_np(num_sampled_tokens)
+            self.routed_experts_cpu: RoutedExpertsTensors | None = None
+            if routed_experts is not None:
+                self.routed_experts_cpu = routed_experts.to_cpu_nonblocking()
             self.prompt_logprobs_dict = {
                 k: v.to_cpu_nonblocking() if v is not None else None
                 for k, v in self.model_runner_output.prompt_logprobs_dict.items()
@@ -78,6 +84,8 @@ class AsyncOutput(AsyncModelRunnerOutput):
         if self.logprobs_tensors is not None:
             self.model_runner_output.logprobs = self.logprobs_tensors.tolists()
         self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict
+        if self.routed_experts_cpu is not None:
+            self.model_runner_output.routed_experts = self.routed_experts_cpu.tolists()
 
         if self._has_fault is not None and self._has_fault.item():
             mask = get_ep_all2all_manager().query_active_mask()

vllm/v1/worker/gpu/model_runner.py

@@ -39,6 +39,10 @@ from vllm.distributed.parallel_state import (
 from vllm.forward_context import BatchDescriptor, set_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
+    RoutedExpertsCapturer,
+    bind_routed_experts_capturer,
+)
 from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
     initialize_mamba_ssu_backend,
 )
@@ -55,7 +59,11 @@ from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
 from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
-from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
+from vllm.v1.outputs import (
+    DraftTokenIds,
+    ModelRunnerOutput,
+    RoutedExpertsTensors,
+)
 from vllm.v1.worker.block_table import get_block_table_width
 from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
 from vllm.v1.worker.gpu import pcp_manager as pcp
@@ -279,11 +287,21 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
         # Expert parallelism load balancer.
         self.eplb = EPLBController(self.parallel_config, self.device)
+        self.routed_experts_capturer: RoutedExpertsCapturer | None = None
 
     def update_max_model_len(self, max_model_len: int) -> None:
         self.max_model_len = max_model_len
         self.req_states.max_model_len = max_model_len
 
+    def init_routed_experts_capturer(self) -> None:
+        """Initialize target-model capture on every participating worker."""
+        self.routed_experts_capturer = RoutedExpertsCapturer(
+            max_num_batched_tokens=self.max_num_tokens,
+            vllm_config=self.vllm_config,
+            kv_cache_config=self.kv_cache_config,
+        )
+        bind_routed_experts_capturer(self.model, self.routed_experts_capturer)
+
     def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
         tasks: list[SupportedTask] = []
         if self.model_config.runner_type == "generate":
@@ -1437,6 +1455,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             aux_hidden_states = None
             output_intermediate_tensors = model_output
 
+        routed_experts = None
+        if not dummy_run and (capturer := self.routed_experts_capturer) is not None:
+            assert slot_mappings is not None
+            routed_experts = capturer.get_routed_experts(slot_mappings, num_toks)
+
         finished_req_ids = scheduler_output.finished_req_ids
         self.execute_model_state = ExecuteModelState(
             input_batch=input_batch,
@@ -1445,6 +1468,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             hidden_states=hidden_states,
             aux_hidden_states=aux_hidden_states,
             finished_req_ids=finished_req_ids,
+            routed_experts=routed_experts,
         )
 
         if not self.is_last_pp_rank:
@@ -1467,6 +1491,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         hidden_states = self.execute_model_state.hidden_states
         aux_hidden_states = self.execute_model_state.aux_hidden_states
         finished_req_ids = self.execute_model_state.finished_req_ids
+        routed_experts = self.execute_model_state.routed_experts
         self.execute_model_state = None
 
         if not self.is_last_pp_rank:
@@ -1532,6 +1557,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             main_stream=self.main_stream,
             copy_stream=self.output_copy_stream,
             check_ep_fault=self.check_ep_fault,
+            routed_experts=routed_experts,
         )
 
         mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
@@ -1712,6 +1738,7 @@ class ExecuteModelState(NamedTuple):
     hidden_states: torch.Tensor | None
     aux_hidden_states: list[torch.Tensor] | None
     finished_req_ids: set[str]
+    routed_experts: RoutedExpertsTensors | None
 
 
 def sort_batch_req_ids(

vllm/v1/worker/gpu_model_runner.py

@@ -63,6 +63,7 @@ from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
 from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
     RoutedExpertsCapturer,
+    bind_routed_experts_capturer,
 )
 from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
     initialize_mamba_ssu_backend,
@@ -2350,8 +2351,7 @@ class GPUModelRunner(
             # by the next ``_prepare_inputs``; we need a stable snapshot
             # because the async D2H may still be in flight on the copy
             # stream when the next step runs.
-            attn_gid = self.routed_experts_attn_gid
-            slot_mapping_attn = slot_mappings[attn_gid]
+            slot_mapping_attn = slot_mappings[self.routed_experts_capturer.attn_gid]
             self.routed_experts_slot_mapping_device[:num_tokens].copy_(
                 slot_mapping_attn[:num_tokens]
             )
@@ -4173,9 +4173,6 @@ class GPUModelRunner(
                 "after execute_model() returns None."
             )
 
-        if self.routed_experts_initialized:
-            self.routed_experts_capturer.clear_buffer()
-
         # If ngram_gpu is used, we need to copy the scheduler_output to avoid
         # the modification has influence on the scheduler_output in engine core process.
         # The replace is much faster than deepcopy.
@@ -4789,24 +4786,16 @@ class GPUModelRunner(
             # copy stream can D2H later. Both tensors must be private
             # clones because:
             #   - ``routing_data`` source is the shared capturer buffer,
-            #     which is ``clear_buffer()``-ed at the start of the
-            #     next step on the default stream.
+            #     which the next forward overwrites on the default stream.
             #   - ``slot_mapping`` source is our own
             #     ``routed_experts_slot_mapping_device``, which the
             #     next ``_prepare_inputs`` overwrites on the default
             #     stream while the D2H is still pending on the copy
             #     stream.
             # Without clones, the copy stream would read torn data.
-            routed_experts_snapshot = None
-            if self.routed_experts_initialized:
-                buf = self.routed_experts_capturer.get_device_buffer()
-                total = scheduler_output.total_num_scheduled_tokens
-                routed_experts_snapshot = RoutedExpertsTensors(
-                    routing_data=buf[:total].clone(),
-                    slot_mapping=self.routed_experts_slot_mapping_device[
-                        :total
-                    ].clone(),
-                )
+            routed_experts_snapshot = self.get_routed_experts(
+                scheduler_output.total_num_scheduled_tokens
+            )
 
             async_output = AsyncGPUModelRunnerOutput(
                 model_runner_output=output,
@@ -7663,19 +7652,18 @@ class GPUModelRunner(
                 kv_transfer_group.register_kv_caches(kv_caches)
             kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks)
 
-    def _get_attention_kv_cache_gid(self) -> int:
-        """Find the KV cache group index for attention layers.
+    def get_routed_experts(
+        self,
+        num_tokens: int,
+    ) -> RoutedExpertsTensors | None:
+        if not self.routed_experts_initialized:
+            return None
 
-        Must match :attr:`RoutedExpertsManager.attn_gid` in the scheduler:
-        both pick the first ``FullAttentionSpec`` group so hybrid models
-        (Mamba / linear-attention layers that use other AttentionSpec
-        subclasses) end up indexing the same slot layout on both sides.
-        Falls back to 0 only for legacy single-group configs.
-        """
-        for gid, group in enumerate(self.kv_cache_config.kv_cache_groups):
-            if isinstance(group.kv_cache_spec, FullAttentionSpec):
-                return gid
-        return 0
+        device_buffer = self.routed_experts_capturer.get_device_buffer()
+        return RoutedExpertsTensors(
+            routing_data=device_buffer[:num_tokens].clone(),
+            slot_mapping=self.routed_experts_slot_mapping_device[:num_tokens].clone(),
+        )
 
     def init_routed_experts_capturer(self):
         logger.info(
@@ -7685,9 +7673,9 @@ class GPUModelRunner(
         self.routed_experts_capturer = RoutedExpertsCapturer(
             max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens,
             vllm_config=self.vllm_config,
+            kv_cache_config=self.kv_cache_config,
         )
-        self.routed_experts_attn_gid = self._get_attention_kv_cache_gid()
-        self._bind_routed_experts_capturer(self.routed_experts_capturer)
+        bind_routed_experts_capturer(self.model, self.routed_experts_capturer)
 
         # Pinned CPU buffer for non-blocking D2H of ``routing_data`` on
         # the sync scheduling path. Shape / dtype mirror the device
@@ -7719,41 +7707,6 @@ class GPUModelRunner(
         )
         self.routed_experts_initialized = True
 
-    def _bind_routed_experts_capturer(self, capturer: RoutedExpertsCapturer) -> None:
-        from vllm.model_executor.layers.fused_moe.layer import MoERunner
-        from vllm.model_executor.layers.fused_moe.modular_kernel import (
-            FusedMoEExpertsMonolithic,
-        )
-        from vllm.model_executor.layers.fused_moe.router.base_router import (
-            BaseRouter,
-        )
-
-        for module in self.model.modules():
-            if not isinstance(module, MoERunner):
-                continue
-            layer_id = module.layer_id
-
-            def _capture_fn(topk_ids, _layer_id=layer_id, _capturer=capturer):
-                _capturer.capture(_layer_id, topk_ids)
-
-            quant_method = module._quant_method
-            moe_kernel = getattr(quant_method, "moe_kernel", None)
-            impl = getattr(moe_kernel, "impl", None)
-            fused_experts = getattr(impl, "fused_experts", None)
-            if quant_method.is_monolithic:
-                if not (
-                    isinstance(fused_experts, FusedMoEExpertsMonolithic)
-                    and fused_experts.supports_routing_replay_capture()
-                ):
-                    raise ValueError(
-                        "--enable-return-routed-experts is not supported with "
-                        f"monolithic MoE kernel {type(fused_experts).__name__}; "
-                        "routed expert IDs would be silently all-zero."
-                    )
-                fused_experts.set_capture_fn(_capture_fn)
-            elif isinstance(module.router, BaseRouter):
-                module.router.set_capture_fn(_capture_fn)
-
     def may_add_encoder_only_layers_to_kv_cache_config(self) -> None:
         """
         Add encoder-only layers to the KV cache config.