iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[RFC] Replace shared-memory routed experts with ModelRunnerOutput transfer and HTTP support (#39568)

Browse Source 
Signed-off-by: xhx1022 <1737006628@qq.com>
Signed-off-by: arlenxu <arlenxu@tencent.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: arlenxu <arlenxu@tencent.com>
Co-authored-by: Junjie Zhang <junj.jay.zhang@gmail.com>
This commit is contained in:
Hongxin Xu
2026-05-14 14:12:30 +00:00
committed by GitHub
co-authored by Claude Opus 4.6 arlenxu Junjie Zhang
parent 2317682f95
commit c7560af424
12 changed files with 605 additions and 369 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/model_executor/test_routed_experts_capture.py
+12 -7
View File
@@ -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
tests/v1/core/test_async_scheduler.py
+1
View File
@@ -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
tests/v1/core/test_scheduler.py
+1
View File
@@ -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
vllm/config/vllm.py
+24
View File
@@ -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)
vllm/entrypoints/serve/disagg/protocol.py
+10
View File
@@ -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):
vllm/entrypoints/serve/disagg/serving.py
+16
View File
@@ -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)
vllm/model_executor/layers/fused_moe/routed_experts_capturer.py
+237 -259
View File
@@ -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]
vllm/sampling_params.py
+7
View File
@@ -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
vllm/v1/core/sched/scheduler.py
+91 -61
View File
@@ -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]:
vllm/v1/engine/output_processor.py
+13 -7
View File
@@ -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
vllm/v1/outputs.py
+78
View File
@@ -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, <dynamic>]
# 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):
vllm/v1/worker/gpu_model_runner.py
+115 -35
View File
@@ -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: