diff --git a/tests/kernels/moe/test_routed_experts_capture_monolithic.py b/tests/kernels/moe/test_routed_experts_capture_monolithic.py index d4a95750ae..895d6e1e63 100644 --- a/tests/kernels/moe/test_routed_experts_capture_monolithic.py +++ b/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, diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index b80b9039e0..f43312a372 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/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 diff --git a/tests/v1/executor/test_ray_utils.py b/tests/v1/executor/test_ray_utils.py index 8da9d5459e..0fa2c129f1 100644 --- a/tests/v1/executor/test_ray_utils.py +++ b/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) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index 9d39621f4f..89a8f6878a 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/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() diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index df4a4a4384..70f111db6e 100644 --- a/vllm/config/vllm.py +++ b/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 diff --git a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py index c71ee30fb0..df6805eb1f 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py +++ b/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 diff --git a/vllm/v1/executor/ray_utils.py b/vllm/v1/executor/ray_utils.py index cc17c39e35..6cfc0c5169 100644 --- a/vllm/v1/executor/ray_utils.py +++ b/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): diff --git a/vllm/v1/worker/gpu/async_utils.py b/vllm/v1/worker/gpu/async_utils.py index a5d5e87ce1..ab91d9dded 100644 --- a/vllm/v1/worker/gpu/async_utils.py +++ b/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() diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index cea735c0e8..92b5790ce2 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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( diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 5f6ce2bc76..4291d72b56 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/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.