From c7560af42487b1570c4e6f4cea5df1605a4d59fc Mon Sep 17 00:00:00 2001 From: Hongxin Xu <70438206+xhx1022@users.noreply.github.com> Date: Thu, 14 May 2026 22:12:30 +0800 Subject: [PATCH] [RFC] Replace shared-memory routed experts with ModelRunnerOutput transfer and HTTP support (#39568) Signed-off-by: xhx1022 <1737006628@qq.com> Signed-off-by: arlenxu Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: arlenxu Co-authored-by: Junjie Zhang --- .../test_routed_experts_capture.py | 19 +- tests/v1/core/test_async_scheduler.py | 1 + tests/v1/core/test_scheduler.py | 1 + vllm/config/vllm.py | 24 + vllm/entrypoints/serve/disagg/protocol.py | 10 + vllm/entrypoints/serve/disagg/serving.py | 16 + .../fused_moe/routed_experts_capturer.py | 496 +++++++++--------- vllm/sampling_params.py | 7 + vllm/v1/core/sched/scheduler.py | 152 +++--- vllm/v1/engine/output_processor.py | 20 +- vllm/v1/outputs.py | 78 +++ vllm/v1/worker/gpu_model_runner.py | 150 ++++-- 12 files changed, 605 insertions(+), 369 deletions(-) diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index 726c419419..152feac9e3 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/tests/model_executor/test_routed_experts_capture.py @@ -25,10 +25,15 @@ def _capturer_with_buffer( num_layers: int = 4, num_experts_per_tok: int = 2, dp_rank: int = 0, + tp_size: int = 1, ) -> RoutedExpertsCapturer: - c = RoutedExpertsCapturer() + # Bypass __init__ so the test can use a CPU buffer and skip the + # VllmConfig dependency. The CUDA device-tensor allocation in the + # real constructor is not what we are exercising here. + c = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer) c.dp_rank = dp_rank - c._device_buffer = torch.full( + c.tp_size = tp_size + c.device_buffer = torch.full( (max_tokens, num_layers, num_experts_per_tok), -1, dtype=torch.int32, @@ -193,8 +198,8 @@ def test_routed_experts_capturer_single_dp_no_metadata(): ctx = SimpleNamespace(dp_metadata=None) with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): capturer.capture(layer_id=0, topk_ids=topk) - assert torch.equal(capturer._device_buffer[:3, 0, :], topk) - assert capturer._device_buffer[3, 0, 0].item() == -1 + assert torch.equal(capturer.device_buffer[:3, 0, :], topk) + assert capturer.device_buffer[3, 0, 0].item() == -1 def test_routed_experts_capturer_dp_naive_concatenated_all_ranks(): @@ -211,7 +216,7 @@ def test_routed_experts_capturer_dp_naive_concatenated_all_ranks(): with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): capturer.capture(layer_id=0, topk_ids=topk) want = topk[2:5] - assert torch.equal(capturer._device_buffer[:3, 0, :], want) + assert torch.equal(capturer.device_buffer[:3, 0, :], want) def test_routed_experts_capturer_dp_modular_local_tokens(): @@ -224,7 +229,7 @@ def test_routed_experts_capturer_dp_modular_local_tokens(): topk = torch.tensor([[10, 11], [12, 13], [14, 15]], dtype=torch.int32) with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): capturer.capture(layer_id=0, topk_ids=topk) - assert torch.equal(capturer._device_buffer[:3, 0, :], topk) + assert torch.equal(capturer.device_buffer[:3, 0, :], topk) def test_routed_experts_capturer_dp_unexpected_batch_raises(): @@ -241,4 +246,4 @@ def test_routed_experts_capturer_dp_unexpected_batch_raises(): pytest.raises(AssertionError, match="unexpected topk_ids batch dim"), ): capturer.capture(layer_id=0, topk_ids=topk) - assert capturer._device_buffer[0, 0, 0].item() == -1 + assert capturer.device_buffer[0, 0, 0].item() == -1 diff --git a/tests/v1/core/test_async_scheduler.py b/tests/v1/core/test_async_scheduler.py index c05e979a9f..a77a50173f 100644 --- a/tests/v1/core/test_async_scheduler.py +++ b/tests/v1/core/test_async_scheduler.py @@ -282,6 +282,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.finished_req_ids_dict = None scheduler.vllm_config = Mock() scheduler.vllm_config.model_config.enable_return_routed_experts = False + scheduler.enable_return_routed_experts = False scheduler.recompute_kv_load_failures = False scheduler.make_stats = Mock(return_value=None) scheduler.max_model_len = 128 diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 42f4825e2b..1c5e1ac4ba 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -2545,6 +2545,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance(): scheduler.finished_req_ids_dict = None scheduler.vllm_config = Mock() scheduler.vllm_config.model_config.enable_return_routed_experts = False + scheduler.enable_return_routed_experts = False scheduler.recompute_kv_load_failures = False scheduler.make_stats = Mock(return_value=None) scheduler.max_model_len = 128 diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 05c3e6bc9e..9f93ab7944 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -785,6 +785,30 @@ class VllmConfig: self.parallel_config.is_moe_model = self.model_config.is_moe + if ( + self.model_config is not None + and self.model_config.enable_return_routed_experts + ): + if self.parallel_config.pipeline_parallel_size > 1: + raise ValueError( + "--enable-return-routed-experts is incompatible with " + "pipeline parallelism (PP > 1)." + ) + + # Incompatible with any KV connector — covers both PD disaggregation + # (kv_producer/kv_consumer: routing captured on P can't reach D) and + # single-instance KV offload/sharing (kv_both: slot_mapping semantics + # change when KV blocks live outside local GPU memory, breaking the + # slot-indexed routed_experts buffer). + if ( + self.kv_transfer_config is not None + and self.kv_transfer_config.is_kv_transfer_instance + ): + raise ValueError( + "--enable-return-routed-experts is incompatible with KV " + "connectors (PD disaggregation, KV cache offload)." + ) + if self.lora_config is not None: self.lora_config.verify_with_model_config(self.model_config) diff --git a/vllm/entrypoints/serve/disagg/protocol.py b/vllm/entrypoints/serve/disagg/protocol.py index 13d0e0d21d..7d468d5d92 100644 --- a/vllm/entrypoints/serve/disagg/protocol.py +++ b/vllm/entrypoints/serve/disagg/protocol.py @@ -160,6 +160,16 @@ class GenerateResponseChoice(BaseModel): # per OpenAI spec this is the default finish_reason: str | None = "stop" token_ids: list[int] | None = None + # Per-token expert routing decisions, base64-encoded ``.npy`` bytes + # (numpy serialization). Shape after decode: + # (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16 + # ``num_tokens - 1`` because the last sampled token has not been + # forwarded yet and therefore has no routing data. + # Decode: + # np.load(io.BytesIO(base64.b64decode(s))) + # ``None`` if (a) the request was aborted before any forward pass, + # or (b) ``enable_return_routed_experts`` is off server-side. + routed_experts: str | None = None class GenerateResponseStreamChoice(BaseModel): diff --git a/vllm/entrypoints/serve/disagg/serving.py b/vllm/entrypoints/serve/disagg/serving.py index 9ef3198464..0edfa61487 100644 --- a/vllm/entrypoints/serve/disagg/serving.py +++ b/vllm/entrypoints/serve/disagg/serving.py @@ -3,10 +3,13 @@ import asyncio +import io import time from collections.abc import AsyncGenerator from collections.abc import Sequence as GenericSequence +import numpy as np +import pybase64 as base64 from fastapi import Request from vllm.engine.protocol import EngineClient @@ -255,11 +258,24 @@ class ServingTokens(OpenAIServing): else: logprobs = None + # Encode routed_experts for transport. JSON can't carry raw + # bytes, so we write the ndarray as a ``.npy`` byte stream + # and base64-encode it. ``pybase64`` is ~3x faster than the + # stdlib ``base64`` on large payloads thanks to SIMD. + # This is the only base64 hop in the pipeline -- the + # engine<->API-server link is binary msgpack + zmq. + routed_experts_b64 = None + if output.routed_experts is not None: + buf = io.BytesIO() + np.save(buf, output.routed_experts) + routed_experts_b64 = base64.b64encode(buf.getvalue()).decode("ascii") + choice_data = GenerateResponseChoice( index=output.index, logprobs=logprobs, finish_reason=output.finish_reason if output.finish_reason else "stop", token_ids=as_list(output.token_ids), + routed_experts=routed_experts_b64, ) choices.append(choice_data) 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 bf4d02782a..115f43a58b 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py @@ -5,22 +5,16 @@ from __future__ import annotations -import fcntl import logging -import os -import tempfile -from collections.abc import Generator -from contextlib import contextmanager -from multiprocessing import shared_memory -from unittest.mock import patch import numpy as np import torch from vllm.config import VllmConfig -from vllm.distributed import get_tensor_model_parallel_rank +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 logger = logging.getLogger(__name__) @@ -42,151 +36,104 @@ def _get_num_experts_per_tok(hf_config) -> int: return val -# Constants -_TMP_DIR = tempfile.gettempdir() -_LOCK_FILE_PREFIX = os.path.join(_TMP_DIR, "vllm_routed_experts") -_BUFFER_PREFIX = "vllm_routed_experts_buffer" +def get_num_experts(hf_config) -> int: + """Resolve ``num_experts`` across HuggingFace config naming conventions. -# Global singleton instances -_global_experts_capturer: RoutedExpertsCapturer | None = None -_global_experts_reader: RoutedExpertsReader | None = None - - -@contextmanager -def _file_lock(lock_file: str, mode: str = "wb+") -> Generator[None, None, None]: - """Context manager for file-based locking.""" - with open(lock_file, mode) as fp: - fcntl.flock(fp, fcntl.LOCK_EX) - try: - yield - finally: - fcntl.flock(fp, fcntl.LOCK_UN) - - -def _create_or_attach_shared_memory( - name: str, size: int, lock_file: str -) -> shared_memory.SharedMemory: - """Create or attach to shared memory with proper locking.""" - # Ensure lock file exists before acquiring lock - with open(lock_file, "wb"): - pass - - with _file_lock(lock_file): - try: - shm = shared_memory.SharedMemory(name=name, create=True, size=size) - except FileExistsError: - shm = shared_memory.SharedMemory(name=name, create=False, size=size) - - if shm.size != size: - logger.warning( - "Shared memory %s size mismatch; recreating", - name, - ) - shm.close() - shm.unlink() - try: - shm = shared_memory.SharedMemory(name=name, create=True, size=size) - logger.info("Created shared memory %s", name) - except FileExistsError: - shm = shared_memory.SharedMemory(name=name, create=False, size=size) - logger.info("Linked to existing shared memory %s", name) - - return shm + Different MoE model families expose this under different keys: + - ``num_experts``: Mixtral, Qwen2-MoE, Qwen3-MoE + - ``n_routed_experts``: DeepSeek-V2/V3 + - ``num_local_experts``: Mixtral (older exports) + """ + for key in ("num_experts", "n_routed_experts", "num_local_experts"): + val = getattr(hf_config, key, None) + if val is not None: + return val + raise ValueError( + "Could not resolve num_experts from model config. " + "Expected one of 'num_experts', 'n_routed_experts', " + "or 'num_local_experts'." + ) class RoutedExpertsCapturer: - """ - Capturer for routed experts with device and optional shared memory buffer. + """Worker-side capturer for routed experts, lives on GPU. - This class captures expert routing decisions during model forward passes - and optionally stores them in shared memory for cross-process access. + Layer-level hooks call :meth:`capture` from inside the forward pass + with the per-layer ``topk_ids`` tensor. The tensor is sliced to the + tokens owned by this DP rank and written into a preallocated device + buffer. At the end of the step, :class:`GPUModelRunner` reads the + device buffer, issues a D2H copy into a pinned CPU buffer, and hands + the result to the scheduler via :class:`RoutedExpertsLists`. + + The device / pinned-CPU transit buffers use ``torch.int32`` (not a + narrow ``uint8``/``uint16`` sized by ``num_experts``). This keeps the + SP all-gather path free of dtype casts, matches the router's native + ``topk_ids`` indices dtype more closely, and costs only a few MB per + worker (``max_num_batched_tokens * num_layers * top_k * 4`` bytes). + The scheduler-side slot buffer + (``RoutedExpertsManager.routed_experts_by_slot``) still uses the + narrow dtype -- numpy fancy-index assignment in ``store_batch`` + narrows the data on the way in. + + 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. + - ``device_buffer.dtype`` is ``torch.int32``. """ - _instance: RoutedExpertsCapturer | None = None - - def __init__(self) -> None: - self._device_buffer: torch.Tensor | None = None - self._shm: shared_memory.SharedMemory | None = None - self._host_buffer_view: np.ndarray | None = None - self._lock_file: str | None = None - - @classmethod - def create(cls) -> RoutedExpertsCapturer: - """Create a global singleton instance.""" - global _global_experts_capturer - if _global_experts_capturer is not None: - raise RuntimeError("Experts capturer already created.") - - _global_experts_capturer = cls() - return _global_experts_capturer - - @staticmethod - def get_instance() -> RoutedExpertsCapturer | None: - """Get the global singleton instance.""" - return _global_experts_capturer - - def init_buffer( + def __init__( self, max_num_batched_tokens: int, - max_num_kv_tokens: int, vllm_config: VllmConfig, ) -> None: - """ - Initialize the device buffer and optionally shared memory buffer. - - Args: - max_num_batched_tokens: Maximum number of tokens in a batch. - max_num_kv_tokens: Maximum number of KV tokens for shared memory. - vllm_config: vllm configuration containing layer and expert info. - """ - - if self._device_buffer is not None: - raise RuntimeError("Device buffer has already been initialized") - hf_config = vllm_config.model_config.hf_text_config - num_layers = hf_config.num_hidden_layers num_experts_per_tok = _get_num_experts_per_tok(hf_config) - - # Initialize device buffer - self._device_buffer = torch.zeros( - (max_num_batched_tokens, num_layers, num_experts_per_tok), + self.device_buffer = torch.zeros( + ( + max_num_batched_tokens, + hf_config.num_hidden_layers, + num_experts_per_tok, + ), + # Use int32 for the device / host transit buffers: it + # matches the router's native topk_ids dtype, is universally + # supported by NCCL (uint8/uint16 are version-dependent), + # and the extra bytes are small (few MB per worker). The + # big scheduler-side slot buffer stays narrow. dtype=torch.int32, device=current_platform.device_type, ) self.dp_rank = vllm_config.parallel_config.data_parallel_rank - - if get_tensor_model_parallel_rank() != 0: - return - - # Initialize shared memory - shape = (max_num_kv_tokens, num_layers, num_experts_per_tok) - buffer_size = int(np.prod(shape)) * np.dtype(np.int32).itemsize - instance_id = vllm_config.instance_id - self._lock_file = f"{_LOCK_FILE_PREFIX}_{instance_id}_{self.dp_rank}.lock" - shm_name = f"{_BUFFER_PREFIX}_{instance_id}_{self.dp_rank}" - - self._shm = _create_or_attach_shared_memory( - shm_name, buffer_size, self._lock_file - ) - self._host_buffer_view = np.ndarray(shape, dtype=np.int32, buffer=self._shm.buf) - self._host_buffer_view.fill(0) - - logger.debug( - "Created shared memory buffer '%s' with shape %s", - shm_name, - shape, - ) + self.tp_size = vllm_config.parallel_config.tensor_parallel_size def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None: - """ - Capture expert routing decisions for a specific layer. + """Capture expert routing decisions for a specific layer. + + Under data parallelism, ``topk_ids`` may have three different batch + layouts depending on where the DP combine happens and whether + Sequence Parallelism (SP) is active for the MoE layer: + - ``n == total`` (naive dispatch): all DP ranks' tokens are + concatenated before routing; we slice out this rank's span + using the cumulative per-rank counts. + - ``n == token_num_per_dp`` (modular-kernel path): DP combine + happens inside ``quant_method.apply``; ``select_experts`` only + ever sees this rank's tokens, so we take the whole tensor. + - ``n == ceil(token_num_per_dp / tp_size)`` (SP + modular-kernel + path): tokens were split along dim=0 across the TP group by + ``_sequence_parallel_context`` + (``moe_runner_base.py:_sequence_parallel_context``), so each + TP rank only sees its shard. We all-gather along dim=0 to + reconstruct this DP rank's full routing tensor. SP pads with + ceil-div (see ``_compute_sp_num_tokens`` in + ``forward_context.py``), so the gathered tensor may contain a + few trailing padding rows which are trimmed by the downstream + ``[:token_num_per_dp]`` slice. Args: layer_id: The layer index. topk_ids: Tensor of shape (batch_size, num_routed_experts). """ - if self._device_buffer is None: - raise RuntimeError("Buffer not initialized. Call init_buffer() first.") ctx = get_forward_context() if ctx.dp_metadata is None: # single dp @@ -200,172 +147,203 @@ class RoutedExpertsCapturer: n = topk_ids.shape[0] if n == total: - # Naive dispatch: all DP ranks' tokens concatenated before routing. + # Naive dispatch: all DP ranks' tokens concatenated + # before routing. This rank owns tokens + # [end_loc - token_num_per_dp, end_loc). cumsum = torch.cumsum(num_tokens_dp, dim=0) end_loc = int(cumsum[self.dp_rank].item()) start_loc = end_loc - token_num_per_dp elif n == token_num_per_dp: - # Modular-kernel path: DP combine happens inside quant_method.apply; - # select_experts only sees this rank's tokens. + # Modular-kernel path: DP combine happens inside + # quant_method.apply; select_experts only sees this + # rank's tokens, take the whole tensor. + start_loc = 0 + end_loc = token_num_per_dp + elif ( + self.tp_size > 1 + and n != token_num_per_dp + and n == (token_num_per_dp + self.tp_size - 1) // self.tp_size + ): + # SP + modular-kernel path. All-gather across the TP + # group along dim=0 to reconstruct the full per-DP-rank + # tensor; keep only the first ``token_num_per_dp`` rows + # (trailing rows are SP ceil-div padding). The TP group + # is always initialized on real rollout workers, and + # every rank in the group reaches this branch in + # lockstep (bind is per-FusedMoE layer, SP is a global + # condition), so a bare all_gather here will not + # deadlock -- let it raise if the precondition is + # violated rather than skip silently. + # + # ``topk_ids`` is already whatever the router produced + # (typically int32/int64, both supported by NCCL); the + # downstream ``device_buffer[...] = topk_ids[...]`` + # setitem narrows into int32 automatically. + topk_ids = get_tp_group().all_gather(topk_ids, dim=0) start_loc = 0 end_loc = token_num_per_dp else: + sp_expected = ( + (token_num_per_dp + self.tp_size - 1) // self.tp_size + if self.tp_size > 0 + else -1 + ) raise AssertionError( - "RoutedExpertsCapturer: unexpected topk_ids batch dim " - f"{n} (expected {total} or {token_num_per_dp} " - f"for dp_rank={self.dp_rank})" + "RoutedExpertsCapturer: unexpected topk_ids batch " + f"dim {n} (expected {total}, {token_num_per_dp}, " + f"or {sp_expected} for dp_rank={self.dp_rank}, " + f"tp_size={self.tp_size})" ) - if layer_id >= self._device_buffer.shape[1]: + # 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 - self._device_buffer[:token_num_per_dp, layer_id, :] = topk_ids[ + self.device_buffer[:token_num_per_dp, layer_id, :] = topk_ids[ start_loc:end_loc, : ] def clear_buffer(self) -> None: - """Clear the device buffer.""" - if self._device_buffer is not None: - self._device_buffer.zero_() - - def save_captured_experts(self, indices: np.ndarray) -> 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. """ - Save captured experts from device buffer to shared memory. + self.device_buffer.zero_() - Args: - indices: Array of indices indicating where to store the data. + 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`. """ - if get_tensor_model_parallel_rank() != 0: - return - if self._lock_file is None: - raise RuntimeError("Shared memory not initialized.") - if self._host_buffer_view is None: - return - if self._device_buffer is None: - raise RuntimeError("Device buffer not initialized.") - - num_tokens = len(indices) - data = self._device_buffer[:num_tokens, :, :].cpu().numpy() - - with _file_lock(self._lock_file): - self._host_buffer_view[indices, :, :] = data - - def cleanup(self) -> None: - """Explicitly clean up shared memory resources.""" - if self._shm is not None: - try: - self._shm.close() - self._shm.unlink() - except Exception: - logger.debug("Exception during cleanup for capturer", exc_info=True) - finally: - self._shm = None - - def __del__(self) -> None: - """Clean up shared memory on destruction.""" - self.cleanup() + return self.device_buffer -class RoutedExpertsReader: - """ - Reader for routed experts from shared memory. +class RoutedExpertsManager: + """Scheduler-side slot-indexed buffer for routed experts. - This class attaches to shared memory created by RoutedExpertsCapturer - and reads expert routing decisions. + Lives on CPU in the scheduler process. Each slot corresponds to + ``block_id * block_size + offset_in_block`` where ``block_id`` is + drawn from the physical KV-cache block pool, so routing data is + tied to physical blocks and naturally survives preemption for + prefix-cached blocks (prefix hits re-expose the same slots). + + Data flow per step: + 1. Worker D2Hs its device capture buffer into + :class:`RoutedExpertsLists` and returns it via + :class:`ModelRunnerOutput`. + 2. Scheduler calls :meth:`store_batch` with that step's + ``(routing_data, slot_mapping)`` — a single CPU->CPU + fancy-index assign, ~few MB per step. + 3. On request completion / abort / preemption, the scheduler + calls :meth:`get` with the request's block IDs to recover + the full per-token routing. + + Memory: ``routed_experts_by_slot`` is sized for the whole block + pool (``num_blocks * block_size`` slots). For large block pools + this can reach multiple GB; see the init log for the exact size. """ - _instance: RoutedExpertsReader | None = None - - def __init__(self) -> None: - self._shm: shared_memory.SharedMemory | None = None - self._host_buffer_view: np.ndarray | None = None - self._lock_file: str | None = None - - @classmethod - def create(cls) -> RoutedExpertsReader: - """Create a global singleton instance.""" - global _global_experts_reader - if _global_experts_reader is not None: - raise RuntimeError("Experts reader already created.") - - _global_experts_reader = cls() - return _global_experts_reader - - @staticmethod - def get_instance() -> RoutedExpertsReader | None: - """Get the global singleton instance.""" - if _global_experts_reader is None: - logger.info("Experts reader not initialized.") - return _global_experts_reader - - def attach_buffer( + def __init__( self, - max_num_kv_tokens: int, vllm_config: VllmConfig, + kv_cache_config: KVCacheConfig, ) -> None: - """ - Attach to an existing shared memory buffer. - - Args: - max_num_kv_tokens: Maximum number of KV tokens. - vllm_config: vllm configuration. - """ - if self._shm is not None: - logger.warning("Already attached to shared memory buffer.") - return # Already attached + # 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) + ) + attn_group = kv_cache_config.kv_cache_groups[self.attn_gid] + self.block_size = attn_group.kv_cache_spec.block_size + # All kv_cache_groups share the same physical block pool, so + # block IDs span [0, num_blocks) regardless of how many groups + # exist. Sizing to the full pool avoids index-out-of-range + # when different groups happen to land on the same block. hf_config = vllm_config.model_config.hf_text_config - shape = ( - max_num_kv_tokens, + num_experts = get_num_experts(hf_config) + num_experts_per_tok = _get_num_experts_per_tok(hf_config) + max_num_slots = kv_cache_config.num_blocks * self.block_size + # Expert IDs are 0..num_experts-1; uint8 fits 256 distinct + # values so the boundary is ``<= 256`` (NOT ``< 256``). Keeping + # this narrow matters because the slot buffer is sized for the + # whole block pool and can reach multiple GB. + expert_id_dtype = np.uint8 if num_experts <= 256 else np.uint16 + self.routed_experts_by_slot = np.zeros( + ( + max_num_slots, + hf_config.num_hidden_layers, + num_experts_per_tok, + ), + dtype=expert_id_dtype, + ) + logger.info( + "RoutedExpertsManager CPU buffer: %.2f GB " + "(slots=%d, layers=%d, top_k=%d, dtype=%s)", + self.routed_experts_by_slot.nbytes / 1e9, + max_num_slots, hf_config.num_hidden_layers, - _get_num_experts_per_tok(hf_config), + hf_config.num_experts_per_tok, + self.routed_experts_by_slot.dtype.name, ) - self.dp_rank = vllm_config.parallel_config.data_parallel_rank - instance_id = vllm_config.instance_id - self._lock_file = f"{_LOCK_FILE_PREFIX}_{instance_id}_{self.dp_rank}.lock" - shm_name = f"{_BUFFER_PREFIX}_{instance_id}_{self.dp_rank}" + def store_batch(self, data: np.ndarray, slot_mapping: np.ndarray) -> None: + """Persist one step's routed experts into the slot buffer. - with _file_lock(self._lock_file, mode="rb+"): - # Avoid resource_tracker registering the shared memory - with patch( - "multiprocessing.resource_tracker.register", - lambda *args, **kwargs: None, - ): - self._shm = shared_memory.SharedMemory(name=shm_name) - - self._host_buffer_view = np.ndarray( - shape, dtype=np.int32, buffer=self._shm.buf - ) - - def get_routed_experts(self, indices: np.ndarray) -> np.ndarray: + Equivalent to ``slot_buffer[slot_mapping] = data``; numpy fancy + indexing handles repeated / out-of-order indices. Called once + per scheduler step in ``update_from_output``. """ - Read routed expert data from shared memory. + self.routed_experts_by_slot[slot_mapping] = data + + def get( + self, + block_ids: list[int], + num_tokens: int, + token_start: int = 0, + ) -> np.ndarray: + """Read routed experts data for a completed / preempted request. + + Reconstructs a per-token slot_mapping from the request's block + IDs and returns the routing slice. Because numpy fancy indexing + returns a **copy** (not a view), the returned ndarray is safe + to hold across subsequent :meth:`store_batch` calls — do not + replace the fancy index with a slice without re-verifying. Args: - indices: Array of indices to read. + block_ids: Block IDs from the attention KV-cache group. + num_tokens: Number of tokens that have gone through a forward + pass and therefore have routing data written to their + slots (typically ``request.num_tokens - 1``; the last + sampled token has not been forwarded yet). Slots beyond + ``request.num_computed_tokens`` are zero-initialized. + token_start: Skip the first ``token_start`` tokens from the + result. The slot_mapping is sliced before the fancy-index + read, so only the requested slots are fetched — no large + intermediate array is allocated. Clamped to + ``[0, num_tokens]`` automatically. Returns: - Copy of the expert routing data for the given indices. + Array of shape (num_tokens - token_start, num_layers, + num_experts_per_tok). """ - if self._host_buffer_view is None: - raise RuntimeError("Buffer not attached. Call attach_buffer() first.") - if self._lock_file is None: - raise RuntimeError("Lock file not initialized.") - - with _file_lock(self._lock_file, mode="rb+"): - return self._host_buffer_view[indices, :, :].copy() - - def cleanup(self) -> None: - """Explicitly clean up resources (close without unlink).""" - if self._shm is not None: - try: - self._shm.close() - except Exception: - logger.debug("Exception during cleanup for reader", exc_info=True) - finally: - self._shm = None - - def __del__(self) -> None: - """Close shared memory on destruction (do not unlink).""" - self.cleanup() + bs = self.block_size + block_ids_array = np.array(block_ids, dtype=np.int32) + block_offsets = np.arange(bs) + # slot = block_id * block_size + offset_in_block; flatten the + # (num_blocks, block_size) grid and trim to num_tokens, then + # skip the first token_start entries so only the requested + # range is fetched in a single fancy-index read. + slot_mapping = ( + block_ids_array.reshape(-1, 1) * bs + block_offsets.reshape(1, -1) + ).flatten()[:num_tokens] + slot_mapping = slot_mapping[token_start:] + return self.routed_experts_by_slot[slot_mapping] diff --git a/vllm/sampling_params.py b/vllm/sampling_params.py index 23cdc33341..e24bd4eef0 100644 --- a/vllm/sampling_params.py +++ b/vllm/sampling_params.py @@ -294,6 +294,13 @@ class SamplingParams( """Arbitrary additional args, that can be used by custom sampling implementations, plugins, etc. Not used by any in-tree sampling implementations.""" + routed_experts_prompt_start: int = 0 + """When enable_return_routed_experts is active, skip the first + routed_experts_prompt_start prompt tokens from the returned routing + data. In multi-turn agent scenarios, set this to the length of the + already-returned prefix to avoid duplicating routing for prompt tokens + covered by earlier turns. Default 0 returns routing for all prompt + tokens.""" # Fields used for bad words bad_words: list[str] | None = None diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 4dd3a8dc82..e6edbb232b 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -7,8 +7,6 @@ from collections.abc import Iterable from dataclasses import replace from typing import Any -import numpy as np - from vllm import envs from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.config import VllmConfig @@ -28,7 +26,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadat from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( - RoutedExpertsReader, + RoutedExpertsManager, ) from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry from vllm.multimodal.encoder_budget import MultiModalBudget @@ -52,7 +50,7 @@ from vllm.v1.core.sched.request_queue import ( ) from vllm.v1.core.sched.utils import check_stop, remove_all from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs -from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig +from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.metrics.perf import ModelMetrics, PerfStats from vllm.v1.metrics.stats import PrefixCacheStats, SchedulerStats from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput @@ -258,42 +256,24 @@ class Scheduler(SchedulerInterface): if self.log_stats and vllm_config.observability_config.enable_mfu_metrics: self.perf_metrics = ModelMetrics(vllm_config) - if self.vllm_config.model_config.enable_return_routed_experts: + self.enable_return_routed_experts = ( + vllm_config.model_config.enable_return_routed_experts + ) + + if self.enable_return_routed_experts: assert self.dcp_world_size == 1 and self.pcp_world_size == 1, ( "enable_return_routed_experts does not support context parallelism " "(dcp_world_size > 1 or pcp_world_size > 1)" ) - self.routed_experts_reader = RoutedExpertsReader.create() - - assert len(kv_cache_config.kv_cache_groups) > 0, ( - "enable_return_routed_experts requires at least one kv cache group" - ) - # Find the attention group for routed experts indexing. - self.routed_experts_attn_gid = 0 - for gid, group in enumerate(kv_cache_config.kv_cache_groups): - if isinstance(group.kv_cache_spec, AttentionSpec): - self.routed_experts_attn_gid = gid - break - min_block_size = min( - [ - group.kv_cache_spec.block_size - for group in kv_cache_config.kv_cache_groups - ] - ) - num_groups = len(kv_cache_config.kv_cache_groups) - self.max_num_kv_tokens = ( - kv_cache_config.num_blocks // num_groups - ) * min_block_size - dcp_size = self.vllm_config.parallel_config.decode_context_parallel_size - pcp_size = self.vllm_config.parallel_config.prefill_context_parallel_size - if pcp_size * dcp_size > 1: - self.max_num_kv_tokens *= pcp_size * dcp_size - - self.routed_experts_reader.attach_buffer( - max_num_kv_tokens=self.max_num_kv_tokens, - vllm_config=self.vllm_config, + self.routed_experts_mgr = RoutedExpertsManager( + vllm_config=vllm_config, + kv_cache_config=kv_cache_config, ) + # Block-ID snapshot taken at schedule time (before forward), + # so update_from_output can read slot data even if a later + # schedule() frees the blocks (async scheduling race). + self._re_block_ids: dict[str, list[int]] = {} self._pause_state: PauseState = PauseState.UNPAUSED @@ -990,6 +970,22 @@ class Scheduler(SchedulerInterface): request.use_structured_output and not request.is_prefill_chunk ) + # Snapshot block IDs for routed experts before forward starts. + # A concurrent schedule() may preempt requests and free blocks + # before update_from_output runs; the snapshot survives that. + # Use update() to preserve entries from the previous step that + # have not yet been consumed by update_from_output (async + # scheduling may call _update_after_schedule again before the + # prior update_from_output runs). + if self.enable_return_routed_experts: + gid = self.routed_experts_mgr.attn_gid + self._re_block_ids.update( + { + rid: self.kv_cache_manager.get_blocks(rid).get_block_ids()[gid] + for rid in num_scheduled_tokens + } + ) + # Clear the finished request IDs. # NOTE: We shouldn't do self.finished_req_ids.clear() here because # it will also affect the scheduler output. @@ -1323,6 +1319,27 @@ class Scheduler(SchedulerInterface): num_scheduled_tokens, ) + # Persist per-step routed experts into the scheduler-side slot + # buffer (CPU->CPU fancy-index assign; ~few MB per step). + # MUST precede the per-request routing reads below: stopped + # requests may terminate on tokens generated in this very step, + # whose routing was just D2H'd into model_runner_output. + routing_data = None + routing_offsets: dict[str, int] = {} + if model_runner_output.routed_experts is not None: + re = model_runner_output.routed_experts + self.routed_experts_mgr.store_batch(re.routing_data, re.slot_mapping) + routing_data = re.routing_data.astype( + self.routed_experts_mgr.routed_experts_by_slot.dtype, + copy=False, + ) + # Build offset map using model runner's request order + # (input_batch ordering), NOT scheduler dict order. + offset = 0 + for rid in model_runner_output.req_ids: + routing_offsets[rid] = offset + offset += num_scheduled_tokens[rid] + # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more, # the below loop can be a performance bottleneck. We should do our best # to avoid expensive operations inside the loop. @@ -1385,6 +1402,7 @@ class Scheduler(SchedulerInterface): pooler_output = pooler_outputs[req_index] if pooler_outputs else None kv_transfer_params = None status_before_stop = request.status + num_output_tokens_before = len(request._output_token_ids) # Check for stop and update request status. if new_token_ids: @@ -1414,10 +1432,47 @@ class Scheduler(SchedulerInterface): stopped = True routed_experts = None + if ( + self.enable_return_routed_experts + and routing_data is not None + and new_token_ids + ): + req_offset = routing_offsets[req_id] + end = req_offset + num_tokens_scheduled + block_ids = self._re_block_ids.pop(req_id, []) + if num_output_tokens_before == 0: + # Prefill completed: read full prompt routing from + # slot buffer using the block-ID snapshot taken at + # schedule time (immune to async preemption). + if ( + request.sampling_params is not None + and request.sampling_params.routed_experts_prompt_start + is not None + ): + prompt_start = ( + request.sampling_params.routed_experts_prompt_start + ) + assert prompt_start < request.num_prompt_tokens + else: + prompt_start = 0 + routed_experts = self.routed_experts_mgr.get( + block_ids, + request.num_prompt_tokens, + token_start=prompt_start, + ) + else: + if scheduled_spec_token_ids: + # Spec decode: accepted tokens at the START of + # the scheduled range, rejected at the end. + routed_experts = routing_data[ + req_offset : req_offset + len(new_token_ids) + ] + else: + # Normal decode / re-prefill: token(s) at the END. + routed_experts = routing_data[end - len(new_token_ids) : end] + finish_reason = None if stopped: - routed_experts = self._get_routed_experts(request) - # Capture finish_reason BEFORE _handle_stopped_request, which may # reset the status to WAITING for streaming requests that continue. finish_reason = request.get_finished_reason() @@ -1592,31 +1647,6 @@ class Scheduler(SchedulerInterface): self._enqueue_waiting_request(request) return False - def _get_routed_experts(self, request: Request) -> np.ndarray | None: - if not self.vllm_config.model_config.enable_return_routed_experts: - return None - - kv_blocks = self.kv_cache_manager.get_blocks(request.request_id) - block_ids = kv_blocks.get_block_ids()[self.routed_experts_attn_gid] - num_tokens = request.num_tokens - 1 - - # compute slot mapping using attention group's block_size - block_ids_array = np.array(block_ids, dtype=np.int32) - num_blocks = len(block_ids) - attn_group = self.kv_cache_config.kv_cache_groups[self.routed_experts_attn_gid] - block_size = attn_group.kv_cache_spec.block_size - - # generate block offsets - block_offsets = np.arange(0, block_size) - - # compute slot mapping: slot = block_id * block_size + offset - slot_mapping = ( - block_offsets.reshape((1, block_size)) - + block_ids_array.reshape((num_blocks, 1)) * block_size - ).flatten()[:num_tokens] - - return self.routed_experts_reader.get_routed_experts(indices=slot_mapping) - def _update_request_with_output( self, request: Request, new_token_ids: list[int] ) -> tuple[list[int], bool]: diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index 1ae89ae196..e1032cfd1f 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -175,6 +175,9 @@ class RequestState: self.stats = RequestStateStats(arrival_time=arrival_time) if log_stats else None + # Routed experts accumulation (prompt + sample chunks) + self.routed_experts_chunks: list[np.ndarray] = [] + # Stream Interval self.stream_interval = stream_interval self.sent_tokens_offset = 0 # Offset of sent tokens @@ -273,7 +276,6 @@ class RequestState: finish_reason: FinishReason | None, stop_reason: int | str | None, kv_transfer_params: dict[str, Any] | None = None, - routed_experts: np.ndarray | None = None, ) -> RequestOutput | PoolingRequestOutput | None: finished = finish_reason is not None final_only = self.output_kind == RequestOutputKind.FINAL_ONLY @@ -314,9 +316,7 @@ class RequestState: finished, ) - output = self._new_completion_output( - new_token_ids, finish_reason, stop_reason, routed_experts - ) + output = self._new_completion_output(new_token_ids, finish_reason, stop_reason) if self.parent_req is None: outputs = [output] @@ -378,7 +378,6 @@ class RequestState: token_ids: list[int], finish_reason: FinishReason | None, stop_reason: int | str | None, - routed_experts: np.ndarray | None = None, ) -> CompletionOutput: assert self.detokenizer is not None assert self.logprobs_processor is not None @@ -395,6 +394,11 @@ class RequestState: if delta and logprobs: logprobs = logprobs[-len(token_ids) :] + # Concatenate routed experts on finish + routed_experts = None + if finished and self.routed_experts_chunks: + routed_experts = np.concatenate(self.routed_experts_chunks, axis=0) + return CompletionOutput( index=self.request_index, text=text, @@ -616,7 +620,10 @@ class OutputProcessor: finish_reason = engine_core_output.finish_reason stop_reason = engine_core_output.stop_reason kv_transfer_params = engine_core_output.kv_transfer_params - routed_experts = engine_core_output.routed_experts + if engine_core_output.routed_experts is not None: + req_state.routed_experts_chunks.append( + engine_core_output.routed_experts + ) if req_state.is_prefilling: if engine_core_output.prefill_stats is not None: @@ -647,7 +654,6 @@ class OutputProcessor: finish_reason, stop_reason, kv_transfer_params, - routed_experts, ): if req_state.streaming_input: request_output.finished = False diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index 1f102ec617..9703dfa9e7 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -109,6 +109,73 @@ class LogprobsTensors(NamedTuple): ) +class RoutedExpertsTensors(NamedTuple): + """Device-side snapshot of routed experts data, pending async D2H. + + Produced by :class:`GPUModelRunner` at the end of each async-scheduled + step. The copy stream waits on the default stream, then issues + non-blocking D2H via :meth:`to_cpu_nonblocking` into a pinned CPU + buffer; :class:`AsyncGPUModelRunnerOutput.get_output` synchronizes + the copy before the scheduler reads it. + + Sliced to ``total_num_scheduled_tokens`` (step-level, across all + requests — NOT per-request). Both ``routing_data`` and + ``slot_mapping`` must be private clones when sourced from shared + capturer / prepare-input buffers, so the next forward pass / + ``_prepare_inputs`` on the default stream does not race with a + D2H still pending on the copy stream. + """ + + # (num_scheduled_tokens, num_layers, num_experts_per_tok) + routing_data: torch.Tensor + # (num_scheduled_tokens,) + slot_mapping: torch.Tensor + + def to_cpu_nonblocking(self) -> "RoutedExpertsTensors": + """Issue non-blocking D2H on the current stream. + + NOTE: ``non_blocking=True`` only delivers true overlap when the + CPU target is pinned. The current fallback here allocates a + new pageable CPU tensor per call, which silently degrades to a + synchronous copy; acceptable because the sync happens on the + dedicated copy stream, not the default stream. + """ + if self.routing_data.device.type == "cpu": + return self + return RoutedExpertsTensors( + self.routing_data.to("cpu", non_blocking=True), + self.slot_mapping.to("cpu", non_blocking=True), + ) + + def tolists(self) -> "RoutedExpertsLists": + """Convert to the numpy-backed form consumed by the scheduler. + + ``.cpu()`` is a no-op when the tensor is already on CPU, so this + is cheap for the post-D2H case; for raw device tensors it will + synchronously block, which is only reached in tests. + """ + return RoutedExpertsLists( + self.routing_data.cpu().numpy(), + self.slot_mapping.cpu().numpy(), + ) + + +class RoutedExpertsLists(NamedTuple): + """CPU-side routed experts, the form :meth:`RoutedExpertsManager.store_batch` + consumes. + + Batched per scheduler step: the leading dim is the number of tokens + scheduled across all requests in this step (``total_num_scheduled_tokens``), + not per-request tokens. ``slot_mapping[i]`` tells the scheduler which + physical KV-cache slot row ``i`` of ``routing_data`` belongs to. + """ + + # (num_scheduled_tokens, num_layers, num_experts_per_tok) + routing_data: np.ndarray + # (num_scheduled_tokens,) + slot_mapping: np.ndarray + + # [num_reqs, ] # The shape of each element depends on the pooler used PoolerOutput: TypeAlias = torch.Tensor | list[torch.Tensor] | list[torch.Tensor | None] @@ -201,6 +268,17 @@ class ModelRunnerOutput: # information related to cudagraph execution cudagraph_stats: CUDAGraphStat | None = None + # Per-step routed experts data captured by the worker. + # ``routing_data`` shape: (num_scheduled_tokens, num_layers, + # num_experts_per_tok); expert IDs as uint8/uint16. + # ``slot_mapping`` shape: (num_scheduled_tokens,); physical KV-cache + # slot for each row of routing_data. + # ``num_scheduled_tokens`` is step-level (total across all requests + # in this step), not per-request. The scheduler persists this into + # its slot buffer via ``slot_buffer[slot_mapping] = routing_data``. + # ``None`` when ``enable_return_routed_experts`` is off. + routed_experts: RoutedExpertsLists | None = None + # ModelRunnerOutput wrapper for async scheduling. class AsyncModelRunnerOutput(ABC): diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index ee1488d240..a30e9275b1 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -157,6 +157,8 @@ from vllm.v1.outputs import ( LogprobsTensors, ModelRunnerOutput, PoolerOutput, + RoutedExpertsLists, + RoutedExpertsTensors, SamplerOutput, make_empty_encoder_model_runner_output, ) @@ -236,6 +238,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): invalid_req_indices: list[int], async_output_copy_stream: torch.cuda.Stream, vocab_size: int, + routed_experts: RoutedExpertsTensors | None = None, ): self._model_runner_output = model_runner_output self._invalid_req_indices = invalid_req_indices @@ -248,6 +251,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): self._sampled_token_ids = sampled_token_ids self.vocab_size = vocab_size self._logprobs_tensors = logprobs_tensors + self._routed_experts = routed_experts # Initiate the copy on a separate stream, but do not synchronize it. default_stream = torch.cuda.current_stream() @@ -261,6 +265,11 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): if self._logprobs_tensors else None ) + self._routed_experts_cpu = ( + self._routed_experts.to_cpu_nonblocking() + if self._routed_experts is not None + else None + ) self.async_copy_ready_event.record() def get_output(self) -> ModelRunnerOutput: @@ -292,6 +301,11 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): output = self._model_runner_output output.sampled_token_ids = valid_sampled_token_ids output.logprobs = logprobs_lists + + if self._routed_experts_cpu is not None: + output.routed_experts = self._routed_experts_cpu.tolists() + del self._routed_experts + return output @@ -2192,9 +2206,18 @@ class GPUModelRunner( slot_mapping_gid_0 = slot_mappings[0] if self.routed_experts_initialized: + # Copy this step's attention slot_mapping into our private + # device buffer. The shared ``slot_mappings[attn_gid]`` is + # owned by the attention block table and will be overwritten + # 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] - self.slot_mapping = slot_mapping_attn[:num_tokens].cpu().numpy() + self.routed_experts_slot_mapping_device[:num_tokens].copy_( + slot_mapping_attn[:num_tokens] + ) + num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[ :num_reqs_padded ] @@ -3478,6 +3501,21 @@ class GPUModelRunner( invalid_req_indices = [] logprobs_lists = None if not self.use_async_scheduling: + # Sync scheduling: issue routed experts D2H into the pinned + # CPU buffer BEFORE ``_to_list`` below. ``_to_list`` does + # ``event.synchronize()`` on the async copy stream which + # waits for every D2H queued on the default stream since + # the last sync, so this enqueue is naturally covered + # without requiring its own synchronize. + if self.routed_experts_initialized: + buf = self.routed_experts_capturer.get_device_buffer() + total = scheduler_output.total_num_scheduled_tokens + self.routed_experts_cpu[:total].copy_(buf[:total], non_blocking=True) + self.routed_experts_slot_mapping_cpu[:total].copy_( + self.routed_experts_slot_mapping_device[:total], + non_blocking=True, + ) + # Get the valid generated tokens. max_gen_len = sampled_token_ids.shape[-1] if max_gen_len == 1: @@ -3878,11 +3916,7 @@ class GPUModelRunner( ) if self.routed_experts_initialized: - capturer = RoutedExpertsCapturer.get_instance() - if capturer is not None: - capturer.clear_buffer() # noqa - else: - logger.error("RoutedExpertsCapturer not 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. @@ -4415,13 +4449,6 @@ class GPUModelRunner( self.kv_connector_output = None with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"): - if self.routed_experts_initialized: - capturer = RoutedExpertsCapturer.get_instance() - if capturer is not None: - capturer.save_captured_experts(indices=self.slot_mapping) # noqa - else: - logger.error("RoutedExpertsCapturer not initialized.") - output = ModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, @@ -4434,14 +4461,47 @@ class GPUModelRunner( else None, num_nans_in_logits=num_nans_in_logits, cudagraph_stats=cudagraph_stats, + routed_experts=None, ) if not self.use_async_scheduling: + if self.routed_experts_initialized: + # Sync path: D2H was issued in ``_bookkeeping_sync`` and + # synchronized by ``_to_list``'s event.synchronize(), so + # the pinned buffers are ready to be wrapped as numpy. + total = scheduler_output.total_num_scheduled_tokens + output.routed_experts = RoutedExpertsLists( + routing_data=self.routed_experts_cpu[:total].numpy(), + slot_mapping=self.routed_experts_slot_mapping_cpu[:total].numpy(), + ) return output with record_function_or_nullcontext( "gpu_model_runner: AsyncGPUModelRunnerOutput" ): + # Async path: produce a device-side snapshot that the async + # 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. + # - ``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(), + ) + async_output = AsyncGPUModelRunnerOutput( model_runner_output=output, sampled_token_ids=sampler_output.sampled_token_ids, @@ -4449,6 +4509,7 @@ class GPUModelRunner( invalid_req_indices=invalid_req_indices, async_output_copy_stream=self._get_or_create_async_output_copy_stream(), vocab_size=self.input_batch.vocab_size, + routed_experts=routed_experts_snapshot, ) with record_function_or_nullcontext( "gpu_model_runner: set_async_sampled_token_ids" @@ -7042,9 +7103,16 @@ class GPUModelRunner( 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.""" + """Find the KV cache group index for attention layers. + + 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, AttentionSpec): + if isinstance(group.kv_cache_spec, FullAttentionSpec): return gid return 0 @@ -7053,29 +7121,41 @@ class GPUModelRunner( "Initializing routed experts capturer, enable_return_routed_experts: %s", self.model_config.enable_return_routed_experts, ) - routed_experts_capturer = RoutedExpertsCapturer.create() - self.routed_experts_attn_gid = self._get_attention_kv_cache_gid() - min_block_size = min( - [ - group.kv_cache_spec.block_size - for group in self.kv_cache_config.kv_cache_groups - ] - ) - num_groups = len(self.kv_cache_config.kv_cache_groups) - self.max_num_kv_tokens = ( - self.kv_cache_config.num_blocks // num_groups - ) * min_block_size - dcp_size = self.vllm_config.parallel_config.decode_context_parallel_size - pcp_size = self.vllm_config.parallel_config.prefill_context_parallel_size - if pcp_size * dcp_size > 1: - self.max_num_kv_tokens *= pcp_size * dcp_size - - routed_experts_capturer.init_buffer( + self.routed_experts_capturer = RoutedExpertsCapturer( max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens, - max_num_kv_tokens=self.max_num_kv_tokens, vllm_config=self.vllm_config, ) - self._bind_routed_experts_capturer(routed_experts_capturer) + self.routed_experts_attn_gid = self._get_attention_kv_cache_gid() + self._bind_routed_experts_capturer(self.routed_experts_capturer) + + # Pinned CPU buffer for non-blocking D2H of ``routing_data`` on + # the sync scheduling path. Shape / dtype mirror the device + # capturer exactly so ``copy_`` is a straight memcpy. + self.routed_experts_cpu = torch.empty( + self.routed_experts_capturer.device_buffer.shape, + dtype=self.routed_experts_capturer.device_buffer.dtype, + device="cpu", + pin_memory=self.pin_memory, + ) + # ``slot_mapping`` dtype is fixed to int64 by + # ``block_table.slot_mapping``; we mirror that here. + max_tokens = self.scheduler_config.max_num_batched_tokens + self.routed_experts_slot_mapping_cpu = torch.empty( + (max_tokens,), + dtype=torch.int64, + device="cpu", + pin_memory=self.pin_memory, + ) + # Private device buffer so the shared ``block_table.slot_mapping`` + # can be overwritten by the next ``_prepare_inputs`` while the + # D2H is still pending on the copy stream. Written in + # ``_prepare_inputs``, read in ``_bookkeeping_sync`` (sync path) + # or cloned into a snapshot (async path). + self.routed_experts_slot_mapping_device = torch.empty( + (max_tokens,), + dtype=torch.int64, + device=self.device, + ) self.routed_experts_initialized = True def _bind_routed_experts_capturer(self, capturer: RoutedExpertsCapturer) -> None: