@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c7560af424

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 16:30 UTC

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

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>

Hongxin Xu committed May 14, 2026 at 22:12 UTC c7560af42487b1570c4e6f4cea5df1605a4d59fc

12 files changed +610 -374

tests/model_executor/test_routed_experts_capture.py

+12 -7

@@ -25,10 +25,15 @@ def _capturer_with_buffer(

25 num_layers: int = 4,

26 num_experts_per_tok: int = 2,

27 dp_rank: int = 0,

28 + tp_size: int = 1,

29 ) -> RoutedExpertsCapturer:

29 - c = RoutedExpertsCapturer()

30 + # Bypass __init__ so the test can use a CPU buffer and skip the

31 + # VllmConfig dependency. The CUDA device-tensor allocation in the

32 + # real constructor is not what we are exercising here.

33 + c = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer)

34 c.dp_rank = dp_rank

31 - c._device_buffer = torch.full(

35 + c.tp_size = tp_size

36 + c.device_buffer = torch.full(

37 (max_tokens, num_layers, num_experts_per_tok),

38 -1,

39 dtype=torch.int32,

@@ -193,8 +198,8 @@ def test_routed_experts_capturer_single_dp_no_metadata():

198 ctx = SimpleNamespace(dp_metadata=None)

199 with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):

200 capturer.capture(layer_id=0, topk_ids=topk)

196 - assert torch.equal(capturer._device_buffer[:3, 0, :], topk)

197 - assert capturer._device_buffer[3, 0, 0].item() == -1

201 + assert torch.equal(capturer.device_buffer[:3, 0, :], topk)

202 + assert capturer.device_buffer[3, 0, 0].item() == -1

203

204

205 def test_routed_experts_capturer_dp_naive_concatenated_all_ranks():

@@ -211,7 +216,7 @@ def test_routed_experts_capturer_dp_naive_concatenated_all_ranks():

216 with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):

217 capturer.capture(layer_id=0, topk_ids=topk)

218 want = topk[2:5]

214 - assert torch.equal(capturer._device_buffer[:3, 0, :], want)

219 + assert torch.equal(capturer.device_buffer[:3, 0, :], want)

220

221

222 def test_routed_experts_capturer_dp_modular_local_tokens():

@@ -224,7 +229,7 @@ def test_routed_experts_capturer_dp_modular_local_tokens():

229 topk = torch.tensor([[10, 11], [12, 13], [14, 15]], dtype=torch.int32)

230 with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):

231 capturer.capture(layer_id=0, topk_ids=topk)

227 - assert torch.equal(capturer._device_buffer[:3, 0, :], topk)

232 + assert torch.equal(capturer.device_buffer[:3, 0, :], topk)

233

234

235 def test_routed_experts_capturer_dp_unexpected_batch_raises():

@@ -241,4 +246,4 @@ def test_routed_experts_capturer_dp_unexpected_batch_raises():

246 pytest.raises(AssertionError, match="unexpected topk_ids batch dim"),

247 ):

248 capturer.capture(layer_id=0, topk_ids=topk)

244 - assert capturer._device_buffer[0, 0, 0].item() == -1

249 + assert capturer.device_buffer[0, 0, 0].item() == -1

tests/v1/core/test_async_scheduler.py

@@ -282,6 +282,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():

282 scheduler.finished_req_ids_dict = None

283 scheduler.vllm_config = Mock()

284 scheduler.vllm_config.model_config.enable_return_routed_experts = False

285 + scheduler.enable_return_routed_experts = False

286 scheduler.recompute_kv_load_failures = False

287 scheduler.make_stats = Mock(return_value=None)

288 scheduler.max_model_len = 128

tests/v1/core/test_scheduler.py

@@ -2545,6 +2545,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():

2545 scheduler.finished_req_ids_dict = None

2546 scheduler.vllm_config = Mock()

2547 scheduler.vllm_config.model_config.enable_return_routed_experts = False

2548 + scheduler.enable_return_routed_experts = False

2549 scheduler.recompute_kv_load_failures = False

2550 scheduler.make_stats = Mock(return_value=None)

2551 scheduler.max_model_len = 128

vllm/config/vllm.py

+24

@@ -785,6 +785,30 @@ class VllmConfig:

785

786 self.parallel_config.is_moe_model = self.model_config.is_moe

787

788 + if (

789 + self.model_config is not None

790 + and self.model_config.enable_return_routed_experts

791 + ):

792 + if self.parallel_config.pipeline_parallel_size > 1:

793 + raise ValueError(

794 + "--enable-return-routed-experts is incompatible with "

795 + "pipeline parallelism (PP > 1)."

796 + )

797 +

798 + # Incompatible with any KV connector — covers both PD disaggregation

799 + # (kv_producer/kv_consumer: routing captured on P can't reach D) and

800 + # single-instance KV offload/sharing (kv_both: slot_mapping semantics

801 + # change when KV blocks live outside local GPU memory, breaking the

802 + # slot-indexed routed_experts buffer).

803 + if (

804 + self.kv_transfer_config is not None

805 + and self.kv_transfer_config.is_kv_transfer_instance

806 + ):

807 + raise ValueError(

808 + "--enable-return-routed-experts is incompatible with KV "

809 + "connectors (PD disaggregation, KV cache offload)."

810 + )

811 +

812 if self.lora_config is not None:

813 self.lora_config.verify_with_model_config(self.model_config)

814

vllm/entrypoints/serve/disagg/protocol.py

+10

@@ -160,6 +160,16 @@ class GenerateResponseChoice(BaseModel):

160 # per OpenAI spec this is the default

161 finish_reason: str | None = "stop"

162 token_ids: list[int] | None = None

163 + # Per-token expert routing decisions, base64-encoded ``.npy`` bytes

164 + # (numpy serialization). Shape after decode:

165 + # (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16

166 + # ``num_tokens - 1`` because the last sampled token has not been

167 + # forwarded yet and therefore has no routing data.

168 + # Decode:

169 + # np.load(io.BytesIO(base64.b64decode(s)))

170 + # ``None`` if (a) the request was aborted before any forward pass,

171 + # or (b) ``enable_return_routed_experts`` is off server-side.

172 + routed_experts: str | None = None

173

174

175 class GenerateResponseStreamChoice(BaseModel):

vllm/entrypoints/serve/disagg/serving.py

+16

@@ -3,10 +3,13 @@

5 import asyncio

6 +import io

7 import time

8 from collections.abc import AsyncGenerator

9 from collections.abc import Sequence as GenericSequence

11 +import numpy as np

12 +import pybase64 as base64

13 from fastapi import Request

15 from vllm.engine.protocol import EngineClient

@@ -255,11 +258,24 @@ class ServingTokens(OpenAIServing):

258 else:

259 logprobs = None

260

261 + # Encode routed_experts for transport. JSON can't carry raw

262 + # bytes, so we write the ndarray as a ``.npy`` byte stream

263 + # and base64-encode it. ``pybase64`` is ~3x faster than the

264 + # stdlib ``base64`` on large payloads thanks to SIMD.

265 + # This is the only base64 hop in the pipeline -- the

266 + # engine<->API-server link is binary msgpack + zmq.

267 + routed_experts_b64 = None

268 + if output.routed_experts is not None:

269 + buf = io.BytesIO()

270 + np.save(buf, output.routed_experts)

271 + routed_experts_b64 = base64.b64encode(buf.getvalue()).decode("ascii")

272 +

273 choice_data = GenerateResponseChoice(

274 index=output.index,

275 logprobs=logprobs,

276 finish_reason=output.finish_reason if output.finish_reason else "stop",

277 token_ids=as_list(output.token_ids),

278 + routed_experts=routed_experts_b64,

279 )

280

281 choices.append(choice_data)

vllm/model_executor/layers/fused_moe/routed_experts_capturer.py

+242 -264

@@ -5,22 +5,16 @@

6 from __future__ import annotations

8 -import fcntl

8 import logging

10 -import os

11 -import tempfile

12 -from collections.abc import Generator

13 -from contextlib import contextmanager

14 -from multiprocessing import shared_memory

15 -from unittest.mock import patch

10 import numpy as np

11 import torch

13 from vllm.config import VllmConfig

21 -from vllm.distributed import get_tensor_model_parallel_rank

14 +from vllm.distributed.parallel_state import get_tp_group

15 from vllm.forward_context import get_forward_context

16 from vllm.platforms import current_platform

17 +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig

19 logger = logging.getLogger(__name__)

@@ -42,151 +36,104 @@ def _get_num_experts_per_tok(hf_config) -> int:

36 return val

45 -# Constants

46 -_TMP_DIR = tempfile.gettempdir()

47 -_LOCK_FILE_PREFIX = os.path.join(_TMP_DIR, "vllm_routed_experts")

48 -_BUFFER_PREFIX = "vllm_routed_experts_buffer"

49 -

50 -# Global singleton instances

51 -_global_experts_capturer: RoutedExpertsCapturer | None = None

52 -_global_experts_reader: RoutedExpertsReader | None = None

53 -

54 -

55 -@contextmanager

56 -def _file_lock(lock_file: str, mode: str = "wb+") -> Generator[None, None, None]:

57 - """Context manager for file-based locking."""

58 - with open(lock_file, mode) as fp:

59 - fcntl.flock(fp, fcntl.LOCK_EX)

60 - try:

61 - yield

62 - finally:

63 - fcntl.flock(fp, fcntl.LOCK_UN)

64 -

65 -

66 -def _create_or_attach_shared_memory(

67 - name: str, size: int, lock_file: str

68 -) -> shared_memory.SharedMemory:

69 - """Create or attach to shared memory with proper locking."""

70 - # Ensure lock file exists before acquiring lock

71 - with open(lock_file, "wb"):

72 - pass

73 -

74 - with _file_lock(lock_file):

75 - try:

76 - shm = shared_memory.SharedMemory(name=name, create=True, size=size)

77 - except FileExistsError:

78 - shm = shared_memory.SharedMemory(name=name, create=False, size=size)

79 -

80 - if shm.size != size:

81 - logger.warning(

82 - "Shared memory %s size mismatch; recreating",

83 - name,

84 - )

85 - shm.close()

86 - shm.unlink()

87 - try:

88 - shm = shared_memory.SharedMemory(name=name, create=True, size=size)

89 - logger.info("Created shared memory %s", name)

90 - except FileExistsError:

91 - shm = shared_memory.SharedMemory(name=name, create=False, size=size)

92 - logger.info("Linked to existing shared memory %s", name)

93 -

94 - return shm

39 +def get_num_experts(hf_config) -> int:

40 + """Resolve ``num_experts`` across HuggingFace config naming conventions.

96 -

97 -class RoutedExpertsCapturer:

98 - """

99 - Capturer for routed experts with device and optional shared memory buffer.

100 -

101 - This class captures expert routing decisions during model forward passes

102 - and optionally stores them in shared memory for cross-process access.

42 + Different MoE model families expose this under different keys:

43 + - ``num_experts``: Mixtral, Qwen2-MoE, Qwen3-MoE

44 + - ``n_routed_experts``: DeepSeek-V2/V3

45 + - ``num_local_experts``: Mixtral (older exports)

46 """

47 + for key in ("num_experts", "n_routed_experts", "num_local_experts"):

48 + val = getattr(hf_config, key, None)

49 + if val is not None:

50 + return val

51 + raise ValueError(

52 + "Could not resolve num_experts from model config. "

53 + "Expected one of 'num_experts', 'n_routed_experts', "

54 + "or 'num_local_experts'."

55 + )

105 - _instance: RoutedExpertsCapturer | None = None

107 - def __init__(self) -> None:

108 - self._device_buffer: torch.Tensor | None = None

109 - self._shm: shared_memory.SharedMemory | None = None

110 - self._host_buffer_view: np.ndarray | None = None

111 - self._lock_file: str | None = None

112 -

113 - @classmethod

114 - def create(cls) -> RoutedExpertsCapturer:

115 - """Create a global singleton instance."""

116 - global _global_experts_capturer

117 - if _global_experts_capturer is not None:

118 - raise RuntimeError("Experts capturer already created.")

119 -

120 - _global_experts_capturer = cls()

121 - return _global_experts_capturer

122 -

123 - @staticmethod

124 - def get_instance() -> RoutedExpertsCapturer | None:

125 - """Get the global singleton instance."""

126 - return _global_experts_capturer

58 +class RoutedExpertsCapturer:

59 + """Worker-side capturer for routed experts, lives on GPU.

60 +

61 + Layer-level hooks call :meth:`capture` from inside the forward pass

62 + with the per-layer ``topk_ids`` tensor. The tensor is sliced to the

63 + tokens owned by this DP rank and written into a preallocated device

64 + buffer. At the end of the step, :class:`GPUModelRunner` reads the

65 + device buffer, issues a D2H copy into a pinned CPU buffer, and hands

66 + the result to the scheduler via :class:`RoutedExpertsLists`.

67 +

68 + The device / pinned-CPU transit buffers use ``torch.int32`` (not a

69 + narrow ``uint8``/``uint16`` sized by ``num_experts``). This keeps the

70 + SP all-gather path free of dtype casts, matches the router's native

71 + ``topk_ids`` indices dtype more closely, and costs only a few MB per

72 + worker (``max_num_batched_tokens * num_layers * top_k * 4`` bytes).

73 + The scheduler-side slot buffer

74 + (``RoutedExpertsManager.routed_experts_by_slot``) still uses the

75 + narrow dtype -- numpy fancy-index assignment in ``store_batch``

76 + narrows the data on the way in.

77 +

78 + Invariants:

79 + - One instance per worker; shape is fixed at init and covers the

80 + worst-case step (``max_num_batched_tokens`` tokens).

81 + - :meth:`clear_buffer` is called at the start of every step, so

82 + unused slots stay zero.

83 + - ``device_buffer.dtype`` is ``torch.int32``.

84 + """

128 - def init_buffer(

86 + def __init__(

87 self,

88 max_num_batched_tokens: int,

131 - max_num_kv_tokens: int,

89 vllm_config: VllmConfig,

90 ) -> None:

134 - """

135 - Initialize the device buffer and optionally shared memory buffer.

136 -

137 - Args:

138 - max_num_batched_tokens: Maximum number of tokens in a batch.

139 - max_num_kv_tokens: Maximum number of KV tokens for shared memory.

140 - vllm_config: vllm configuration containing layer and expert info.

141 - """

142 -

143 - if self._device_buffer is not None:

144 - raise RuntimeError("Device buffer has already been initialized")

145 -

91 hf_config = vllm_config.model_config.hf_text_config

147 - num_layers = hf_config.num_hidden_layers

92 num_experts_per_tok = _get_num_experts_per_tok(hf_config)

149 -

150 - # Initialize device buffer

151 - self._device_buffer = torch.zeros(

152 - (max_num_batched_tokens, num_layers, num_experts_per_tok),

93 + self.device_buffer = torch.zeros(

94 + (

95 + max_num_batched_tokens,

96 + hf_config.num_hidden_layers,

97 + num_experts_per_tok,

98 + ),

99 + # Use int32 for the device / host transit buffers: it

100 + # matches the router's native topk_ids dtype, is universally

101 + # supported by NCCL (uint8/uint16 are version-dependent),

102 + # and the extra bytes are small (few MB per worker). The

103 + # big scheduler-side slot buffer stays narrow.

104 dtype=torch.int32,

105 device=current_platform.device_type,

106 )

107 self.dp_rank = vllm_config.parallel_config.data_parallel_rank

157 -

158 - if get_tensor_model_parallel_rank() != 0:

159 - return

160 -

161 - # Initialize shared memory

162 - shape = (max_num_kv_tokens, num_layers, num_experts_per_tok)

163 - buffer_size = int(np.prod(shape)) * np.dtype(np.int32).itemsize

164 - instance_id = vllm_config.instance_id

165 - self._lock_file = f"{_LOCK_FILE_PREFIX}_{instance_id}_{self.dp_rank}.lock"

166 - shm_name = f"{_BUFFER_PREFIX}_{instance_id}_{self.dp_rank}"

167 -

168 - self._shm = _create_or_attach_shared_memory(

169 - shm_name, buffer_size, self._lock_file

170 - )

171 - self._host_buffer_view = np.ndarray(shape, dtype=np.int32, buffer=self._shm.buf)

172 - self._host_buffer_view.fill(0)

173 -

174 - logger.debug(

175 - "Created shared memory buffer '%s' with shape %s",

176 - shm_name,

177 - shape,

178 - )

108 + self.tp_size = vllm_config.parallel_config.tensor_parallel_size

109

110 def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None:

181 - """

182 - Capture expert routing decisions for a specific layer.

111 + """Capture expert routing decisions for a specific layer.

112 +

113 + Under data parallelism, ``topk_ids`` may have three different batch

114 + layouts depending on where the DP combine happens and whether

115 + Sequence Parallelism (SP) is active for the MoE layer:

116 + - ``n == total`` (naive dispatch): all DP ranks' tokens are

117 + concatenated before routing; we slice out this rank's span

118 + using the cumulative per-rank counts.

119 + - ``n == token_num_per_dp`` (modular-kernel path): DP combine

120 + happens inside ``quant_method.apply``; ``select_experts`` only

121 + ever sees this rank's tokens, so we take the whole tensor.

122 + - ``n == ceil(token_num_per_dp / tp_size)`` (SP + modular-kernel

123 + path): tokens were split along dim=0 across the TP group by

124 + ``_sequence_parallel_context``

125 + (``moe_runner_base.py:_sequence_parallel_context``), so each

126 + TP rank only sees its shard. We all-gather along dim=0 to

127 + reconstruct this DP rank's full routing tensor. SP pads with

128 + ceil-div (see ``_compute_sp_num_tokens`` in

129 + ``forward_context.py``), so the gathered tensor may contain a

130 + few trailing padding rows which are trimmed by the downstream

131 + ``[:token_num_per_dp]`` slice.

132

133 Args:

134 layer_id: The layer index.

135 topk_ids: Tensor of shape (batch_size, num_routed_experts).

136 """

188 - if self._device_buffer is None:

189 - raise RuntimeError("Buffer not initialized. Call init_buffer() first.")

137

138 ctx = get_forward_context()

139 if ctx.dp_metadata is None: # single dp

@@ -200,172 +147,203 @@ class RoutedExpertsCapturer:

147 n = topk_ids.shape[0]

148

149 if n == total:

203 - # Naive dispatch: all DP ranks' tokens concatenated before routing.

150 + # Naive dispatch: all DP ranks' tokens concatenated

151 + # before routing. This rank owns tokens

152 + # [end_loc - token_num_per_dp, end_loc).

153 cumsum = torch.cumsum(num_tokens_dp, dim=0)

154 end_loc = int(cumsum[self.dp_rank].item())

155 start_loc = end_loc - token_num_per_dp

156 elif n == token_num_per_dp:

208 - # Modular-kernel path: DP combine happens inside quant_method.apply;

209 - # select_experts only sees this rank's tokens.

157 + # Modular-kernel path: DP combine happens inside

158 + # quant_method.apply; select_experts only sees this

159 + # rank's tokens, take the whole tensor.

160 + start_loc = 0

161 + end_loc = token_num_per_dp

162 + elif (

163 + self.tp_size > 1

164 + and n != token_num_per_dp

165 + and n == (token_num_per_dp + self.tp_size - 1) // self.tp_size

166 + ):

167 + # SP + modular-kernel path. All-gather across the TP

168 + # group along dim=0 to reconstruct the full per-DP-rank

169 + # tensor; keep only the first ``token_num_per_dp`` rows

170 + # (trailing rows are SP ceil-div padding). The TP group

171 + # is always initialized on real rollout workers, and

172 + # every rank in the group reaches this branch in

173 + # lockstep (bind is per-FusedMoE layer, SP is a global

174 + # condition), so a bare all_gather here will not

175 + # deadlock -- let it raise if the precondition is

176 + # violated rather than skip silently.

177 + #

178 + # ``topk_ids`` is already whatever the router produced

179 + # (typically int32/int64, both supported by NCCL); the

180 + # downstream ``device_buffer[...] = topk_ids[...]``

181 + # setitem narrows into int32 automatically.

182 + topk_ids = get_tp_group().all_gather(topk_ids, dim=0)

183 start_loc = 0

184 end_loc = token_num_per_dp

185 else:

186 + sp_expected = (

187 + (token_num_per_dp + self.tp_size - 1) // self.tp_size

188 + if self.tp_size > 0

189 + else -1

190 + )

191 raise AssertionError(

214 - "RoutedExpertsCapturer: unexpected topk_ids batch dim "

215 - f"{n} (expected {total} or {token_num_per_dp} "

216 - f"for dp_rank={self.dp_rank})"

192 + "RoutedExpertsCapturer: unexpected topk_ids batch "

193 + f"dim {n} (expected {total}, {token_num_per_dp}, "

194 + f"or {sp_expected} for dp_rank={self.dp_rank}, "

195 + f"tp_size={self.tp_size})"

196 )

197

219 - if layer_id >= self._device_buffer.shape[1]:

198 + # Defensive: model may expose more layers than the capture buffer

199 + # was sized for (unusual, but guards against miss-config).

200 + if layer_id >= self.device_buffer.shape[1]:

201 return

202

222 - self._device_buffer[:token_num_per_dp, layer_id, :] = topk_ids[

203 + self.device_buffer[:token_num_per_dp, layer_id, :] = topk_ids[

204 start_loc:end_loc, :

205 ]

206

207 def clear_buffer(self) -> None:

227 - """Clear the device buffer."""

228 - if self._device_buffer is not None:

229 - self._device_buffer.zero_()

230 -

231 - def save_captured_experts(self, indices: np.ndarray) -> None:

208 + """Zero the device buffer. Called at the start of every step so

209 + slots belonging to finished / preempted tokens don't leak into

210 + the next step.

211 """

233 - Save captured experts from device buffer to shared memory.

212 + self.device_buffer.zero_()

213

235 - Args:

236 - indices: Array of indices indicating where to store the data.

214 + def get_device_buffer(self) -> torch.Tensor:

215 + """Return the underlying device buffer so the model runner can

216 + issue the D2H copy. The tensor is shared; callers must either

217 + clone or fully drain it before the next forward pass runs

218 + :meth:`clear_buffer`.

219 """

238 - if get_tensor_model_parallel_rank() != 0:

239 - return

240 - if self._lock_file is None:

241 - raise RuntimeError("Shared memory not initialized.")

242 - if self._host_buffer_view is None:

243 - return

244 - if self._device_buffer is None:

245 - raise RuntimeError("Device buffer not initialized.")

246 -

247 - num_tokens = len(indices)

248 - data = self._device_buffer[:num_tokens, :, :].cpu().numpy()

249 -

250 - with _file_lock(self._lock_file):

251 - self._host_buffer_view[indices, :, :] = data

252 -

253 - def cleanup(self) -> None:

254 - """Explicitly clean up shared memory resources."""

255 - if self._shm is not None:

256 - try:

257 - self._shm.close()

258 - self._shm.unlink()

259 - except Exception:

260 - logger.debug("Exception during cleanup for capturer", exc_info=True)

261 - finally:

262 - self._shm = None

263 -

264 - def __del__(self) -> None:

265 - """Clean up shared memory on destruction."""

266 - self.cleanup()

267 -

268 -

269 -class RoutedExpertsReader:

270 - """

271 - Reader for routed experts from shared memory.

272 -

273 - This class attaches to shared memory created by RoutedExpertsCapturer

274 - and reads expert routing decisions.

220 + return self.device_buffer

221 +

222 +

223 +class RoutedExpertsManager:

224 + """Scheduler-side slot-indexed buffer for routed experts.

225 +

226 + Lives on CPU in the scheduler process. Each slot corresponds to

227 + ``block_id * block_size + offset_in_block`` where ``block_id`` is

228 + drawn from the physical KV-cache block pool, so routing data is

229 + tied to physical blocks and naturally survives preemption for

230 + prefix-cached blocks (prefix hits re-expose the same slots).

231 +

232 + Data flow per step:

233 + 1. Worker D2Hs its device capture buffer into

234 + :class:`RoutedExpertsLists` and returns it via

235 + :class:`ModelRunnerOutput`.

236 + 2. Scheduler calls :meth:`store_batch` with that step's

237 + ``(routing_data, slot_mapping)`` — a single CPU->CPU

238 + fancy-index assign, ~few MB per step.

239 + 3. On request completion / abort / preemption, the scheduler

240 + calls :meth:`get` with the request's block IDs to recover

241 + the full per-token routing.

242 +

243 + Memory: ``routed_experts_by_slot`` is sized for the whole block

244 + pool (``num_blocks * block_size`` slots). For large block pools

245 + this can reach multiple GB; see the init log for the exact size.

246 """

247

277 - _instance: RoutedExpertsReader | None = None

278 -

279 - def __init__(self) -> None:

280 - self._shm: shared_memory.SharedMemory | None = None

281 - self._host_buffer_view: np.ndarray | None = None

282 - self._lock_file: str | None = None

283 -

284 - @classmethod

285 - def create(cls) -> RoutedExpertsReader:

286 - """Create a global singleton instance."""

287 - global _global_experts_reader

288 - if _global_experts_reader is not None:

289 - raise RuntimeError("Experts reader already created.")

290 -

291 - _global_experts_reader = cls()

292 - return _global_experts_reader

293 -

294 - @staticmethod

295 - def get_instance() -> RoutedExpertsReader | None:

296 - """Get the global singleton instance."""

297 - if _global_experts_reader is None:

298 - logger.info("Experts reader not initialized.")

299 - return _global_experts_reader

300 -

301 - def attach_buffer(

248 + def __init__(

249 self,

303 - max_num_kv_tokens: int,

250 vllm_config: VllmConfig,

251 + kv_cache_config: KVCacheConfig,

252 ) -> None:

306 - """

307 - Attach to an existing shared memory buffer.

308 -

309 - Args:

310 - max_num_kv_tokens: Maximum number of KV tokens.

311 - vllm_config: vllm configuration.

312 - """

313 - if self._shm is not None:

314 - logger.warning("Already attached to shared memory buffer.")

315 - return # Already attached

253 + # Pick the attention group for block/slot mapping. We require

254 + # a FullAttentionSpec group rather than any AttentionSpec to

255 + # stay consistent with the worker-side lookup in

256 + # ``GPUModelRunner._get_attention_kv_cache_gid``; hybrid models

257 + # (Mamba / linear attention) also have other AttentionSpec

258 + # groups whose slot layout differs.

259 + self.attn_gid = next(

260 + gid

261 + for gid, g in enumerate(kv_cache_config.kv_cache_groups)

262 + if isinstance(g.kv_cache_spec, FullAttentionSpec)

263 + )

264 + attn_group = kv_cache_config.kv_cache_groups[self.attn_gid]

265 + self.block_size = attn_group.kv_cache_spec.block_size

266

267 + # All kv_cache_groups share the same physical block pool, so

268 + # block IDs span [0, num_blocks) regardless of how many groups

269 + # exist. Sizing to the full pool avoids index-out-of-range

270 + # when different groups happen to land on the same block.

271 hf_config = vllm_config.model_config.hf_text_config

318 - shape = (

319 - max_num_kv_tokens,

272 + num_experts = get_num_experts(hf_config)

273 + num_experts_per_tok = _get_num_experts_per_tok(hf_config)

274 + max_num_slots = kv_cache_config.num_blocks * self.block_size

275 + # Expert IDs are 0..num_experts-1; uint8 fits 256 distinct

276 + # values so the boundary is ``<= 256`` (NOT ``< 256``). Keeping

277 + # this narrow matters because the slot buffer is sized for the

278 + # whole block pool and can reach multiple GB.

279 + expert_id_dtype = np.uint8 if num_experts <= 256 else np.uint16

280 + self.routed_experts_by_slot = np.zeros(

281 + (

282 + max_num_slots,

283 + hf_config.num_hidden_layers,

284 + num_experts_per_tok,

285 + ),

286 + dtype=expert_id_dtype,

287 + )

288 + logger.info(

289 + "RoutedExpertsManager CPU buffer: %.2f GB "

290 + "(slots=%d, layers=%d, top_k=%d, dtype=%s)",

291 + self.routed_experts_by_slot.nbytes / 1e9,

292 + max_num_slots,

293 hf_config.num_hidden_layers,

321 - _get_num_experts_per_tok(hf_config),

294 + hf_config.num_experts_per_tok,

295 + self.routed_experts_by_slot.dtype.name,

296 )

297

324 - self.dp_rank = vllm_config.parallel_config.data_parallel_rank

325 - instance_id = vllm_config.instance_id

326 - self._lock_file = f"{_LOCK_FILE_PREFIX}_{instance_id}_{self.dp_rank}.lock"

327 - shm_name = f"{_BUFFER_PREFIX}_{instance_id}_{self.dp_rank}"

328 -

329 - with _file_lock(self._lock_file, mode="rb+"):

330 - # Avoid resource_tracker registering the shared memory

331 - with patch(

332 - "multiprocessing.resource_tracker.register",

333 - lambda *args, **kwargs: None,

334 - ):

335 - self._shm = shared_memory.SharedMemory(name=shm_name)

298 + def store_batch(self, data: np.ndarray, slot_mapping: np.ndarray) -> None:

299 + """Persist one step's routed experts into the slot buffer.

300

337 - self._host_buffer_view = np.ndarray(

338 - shape, dtype=np.int32, buffer=self._shm.buf

339 - )

340 -

341 - def get_routed_experts(self, indices: np.ndarray) -> np.ndarray:

301 + Equivalent to ``slot_buffer[slot_mapping] = data``; numpy fancy

302 + indexing handles repeated / out-of-order indices. Called once

303 + per scheduler step in ``update_from_output``.

304 """

343 - Read routed expert data from shared memory.

305 + self.routed_experts_by_slot[slot_mapping] = data

306 +

307 + def get(

308 + self,

309 + block_ids: list[int],

310 + num_tokens: int,

311 + token_start: int = 0,

312 + ) -> np.ndarray:

313 + """Read routed experts data for a completed / preempted request.

314 +

315 + Reconstructs a per-token slot_mapping from the request's block

316 + IDs and returns the routing slice. Because numpy fancy indexing

317 + returns a **copy** (not a view), the returned ndarray is safe

318 + to hold across subsequent :meth:`store_batch` calls — do not

319 + replace the fancy index with a slice without re-verifying.

320

321 Args:

346 - indices: Array of indices to read.

322 + block_ids: Block IDs from the attention KV-cache group.

323 + num_tokens: Number of tokens that have gone through a forward

324 + pass and therefore have routing data written to their

325 + slots (typically ``request.num_tokens - 1``; the last

326 + sampled token has not been forwarded yet). Slots beyond

327 + ``request.num_computed_tokens`` are zero-initialized.

328 + token_start: Skip the first ``token_start`` tokens from the

329 + result. The slot_mapping is sliced before the fancy-index

330 + read, so only the requested slots are fetched — no large

331 + intermediate array is allocated. Clamped to

332 + ``[0, num_tokens]`` automatically.

333

334 Returns:

349 - Copy of the expert routing data for the given indices.

335 + Array of shape (num_tokens - token_start, num_layers,

336 + num_experts_per_tok).

337 """

351 - if self._host_buffer_view is None:

352 - raise RuntimeError("Buffer not attached. Call attach_buffer() first.")

353 - if self._lock_file is None:

354 - raise RuntimeError("Lock file not initialized.")

355 -

356 - with _file_lock(self._lock_file, mode="rb+"):

357 - return self._host_buffer_view[indices, :, :].copy()

358 -

359 - def cleanup(self) -> None:

360 - """Explicitly clean up resources (close without unlink)."""

361 - if self._shm is not None:

362 - try:

363 - self._shm.close()

364 - except Exception:

365 - logger.debug("Exception during cleanup for reader", exc_info=True)

366 - finally:

367 - self._shm = None

368 -

369 - def __del__(self) -> None:

370 - """Close shared memory on destruction (do not unlink)."""

371 - self.cleanup()

338 + bs = self.block_size

339 + block_ids_array = np.array(block_ids, dtype=np.int32)

340 + block_offsets = np.arange(bs)

341 + # slot = block_id * block_size + offset_in_block; flatten the

342 + # (num_blocks, block_size) grid and trim to num_tokens, then

343 + # skip the first token_start entries so only the requested

344 + # range is fetched in a single fancy-index read.

345 + slot_mapping = (

346 + block_ids_array.reshape(-1, 1) * bs + block_offsets.reshape(1, -1)

347 + ).flatten()[:num_tokens]

348 + slot_mapping = slot_mapping[token_start:]

349 + return self.routed_experts_by_slot[slot_mapping]

vllm/sampling_params.py

@@ -294,6 +294,13 @@ class SamplingParams(

294 """Arbitrary additional args, that can be used by custom sampling

295 implementations, plugins, etc. Not used by any in-tree sampling

296 implementations."""

297 + routed_experts_prompt_start: int = 0

298 + """When enable_return_routed_experts is active, skip the first

299 + routed_experts_prompt_start prompt tokens from the returned routing

300 + data. In multi-turn agent scenarios, set this to the length of the

301 + already-returned prefix to avoid duplicating routing for prompt tokens

302 + covered by earlier turns. Default 0 returns routing for all prompt

303 + tokens."""

304

305 # Fields used for bad words

306 bad_words: list[str] | None = None

vllm/v1/core/sched/scheduler.py

+91 -61

@@ -7,8 +7,6 @@ from collections.abc import Iterable

7 from dataclasses import replace

8 from typing import Any

10 -import numpy as np

11 -

10 from vllm import envs

11 from vllm.compilation.cuda_graph import CUDAGraphStat

12 from vllm.config import VllmConfig

@@ -28,7 +26,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadat

26 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats

27 from vllm.logger import init_logger

28 from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (

31 - RoutedExpertsReader,

29 + RoutedExpertsManager,

30 )

31 from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry

32 from vllm.multimodal.encoder_budget import MultiModalBudget

@@ -52,7 +50,7 @@ from vllm.v1.core.sched.request_queue import (

50 )

51 from vllm.v1.core.sched.utils import check_stop, remove_all

52 from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs

55 -from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig

53 +from vllm.v1.kv_cache_interface import KVCacheConfig

54 from vllm.v1.metrics.perf import ModelMetrics, PerfStats

55 from vllm.v1.metrics.stats import PrefixCacheStats, SchedulerStats

56 from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput

@@ -258,42 +256,24 @@ class Scheduler(SchedulerInterface):

256 if self.log_stats and vllm_config.observability_config.enable_mfu_metrics:

257 self.perf_metrics = ModelMetrics(vllm_config)

258

261 - if self.vllm_config.model_config.enable_return_routed_experts:

259 + self.enable_return_routed_experts = (

260 + vllm_config.model_config.enable_return_routed_experts

261 + )

262 +

263 + if self.enable_return_routed_experts:

264 assert self.dcp_world_size == 1 and self.pcp_world_size == 1, (

265 "enable_return_routed_experts does not support context parallelism "

266 "(dcp_world_size > 1 or pcp_world_size > 1)"

267 )

268

267 - self.routed_experts_reader = RoutedExpertsReader.create()

268 -

269 - assert len(kv_cache_config.kv_cache_groups) > 0, (

270 - "enable_return_routed_experts requires at least one kv cache group"

271 - )

272 - # Find the attention group for routed experts indexing.

273 - self.routed_experts_attn_gid = 0

274 - for gid, group in enumerate(kv_cache_config.kv_cache_groups):

275 - if isinstance(group.kv_cache_spec, AttentionSpec):

276 - self.routed_experts_attn_gid = gid

277 - break

278 - min_block_size = min(

279 - [

280 - group.kv_cache_spec.block_size

281 - for group in kv_cache_config.kv_cache_groups

282 - ]

283 - )

284 - num_groups = len(kv_cache_config.kv_cache_groups)

285 - self.max_num_kv_tokens = (

286 - kv_cache_config.num_blocks // num_groups

287 - ) * min_block_size

288 - dcp_size = self.vllm_config.parallel_config.decode_context_parallel_size

289 - pcp_size = self.vllm_config.parallel_config.prefill_context_parallel_size

290 - if pcp_size * dcp_size > 1:

291 - self.max_num_kv_tokens *= pcp_size * dcp_size

292 -

293 - self.routed_experts_reader.attach_buffer(

294 - max_num_kv_tokens=self.max_num_kv_tokens,

295 - vllm_config=self.vllm_config,

269 + self.routed_experts_mgr = RoutedExpertsManager(

270 + vllm_config=vllm_config,

271 + kv_cache_config=kv_cache_config,

272 )

273 + # Block-ID snapshot taken at schedule time (before forward),

274 + # so update_from_output can read slot data even if a later

275 + # schedule() frees the blocks (async scheduling race).

276 + self._re_block_ids: dict[str, list[int]] = {}

277

278 self._pause_state: PauseState = PauseState.UNPAUSED

279

@@ -990,6 +970,22 @@ class Scheduler(SchedulerInterface):

970 request.use_structured_output and not request.is_prefill_chunk

971 )

972

973 + # Snapshot block IDs for routed experts before forward starts.

974 + # A concurrent schedule() may preempt requests and free blocks

975 + # before update_from_output runs; the snapshot survives that.

976 + # Use update() to preserve entries from the previous step that

977 + # have not yet been consumed by update_from_output (async

978 + # scheduling may call _update_after_schedule again before the

979 + # prior update_from_output runs).

980 + if self.enable_return_routed_experts:

981 + gid = self.routed_experts_mgr.attn_gid

982 + self._re_block_ids.update(

983 + {

984 + rid: self.kv_cache_manager.get_blocks(rid).get_block_ids()[gid]

985 + for rid in num_scheduled_tokens

986 + }

987 + )

988 +

989 # Clear the finished request IDs.

990 # NOTE: We shouldn't do self.finished_req_ids.clear() here because

991 # it will also affect the scheduler output.

@@ -1323,6 +1319,27 @@ class Scheduler(SchedulerInterface):

1319 num_scheduled_tokens,

1320 )

1321

1322 + # Persist per-step routed experts into the scheduler-side slot

1323 + # buffer (CPU->CPU fancy-index assign; ~few MB per step).

1324 + # MUST precede the per-request routing reads below: stopped

1325 + # requests may terminate on tokens generated in this very step,

1326 + # whose routing was just D2H'd into model_runner_output.

1327 + routing_data = None

1328 + routing_offsets: dict[str, int] = {}

1329 + if model_runner_output.routed_experts is not None:

1330 + re = model_runner_output.routed_experts

1331 + self.routed_experts_mgr.store_batch(re.routing_data, re.slot_mapping)

1332 + routing_data = re.routing_data.astype(

1333 + self.routed_experts_mgr.routed_experts_by_slot.dtype,

1334 + copy=False,

1335 + )

1336 + # Build offset map using model runner's request order

1337 + # (input_batch ordering), NOT scheduler dict order.

1338 + offset = 0

1339 + for rid in model_runner_output.req_ids:

1340 + routing_offsets[rid] = offset

1341 + offset += num_scheduled_tokens[rid]

1342 +

1343 # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,

1344 # the below loop can be a performance bottleneck. We should do our best

1345 # to avoid expensive operations inside the loop.

@@ -1385,6 +1402,7 @@ class Scheduler(SchedulerInterface):

1402 pooler_output = pooler_outputs[req_index] if pooler_outputs else None

1403 kv_transfer_params = None

1404 status_before_stop = request.status

1405 + num_output_tokens_before = len(request._output_token_ids)

1406

1407 # Check for stop and update request status.

1408 if new_token_ids:

@@ -1414,10 +1432,47 @@ class Scheduler(SchedulerInterface):

1432 stopped = True

1433

1434 routed_experts = None

1435 + if (

1436 + self.enable_return_routed_experts

1437 + and routing_data is not None

1438 + and new_token_ids

1439 + ):

1440 + req_offset = routing_offsets[req_id]

1441 + end = req_offset + num_tokens_scheduled

1442 + block_ids = self._re_block_ids.pop(req_id, [])

1443 + if num_output_tokens_before == 0:

1444 + # Prefill completed: read full prompt routing from

1445 + # slot buffer using the block-ID snapshot taken at

1446 + # schedule time (immune to async preemption).

1447 + if (

1448 + request.sampling_params is not None

1449 + and request.sampling_params.routed_experts_prompt_start

1450 + is not None

1451 + ):

1452 + prompt_start = (

1453 + request.sampling_params.routed_experts_prompt_start

1454 + )

1455 + assert prompt_start < request.num_prompt_tokens

1456 + else:

1457 + prompt_start = 0

1458 + routed_experts = self.routed_experts_mgr.get(

1459 + block_ids,

1460 + request.num_prompt_tokens,

1461 + token_start=prompt_start,

1462 + )

1463 + else:

1464 + if scheduled_spec_token_ids:

1465 + # Spec decode: accepted tokens at the START of

1466 + # the scheduled range, rejected at the end.

1467 + routed_experts = routing_data[

1468 + req_offset : req_offset + len(new_token_ids)

1469 + ]

1470 + else:

1471 + # Normal decode / re-prefill: token(s) at the END.

1472 + routed_experts = routing_data[end - len(new_token_ids) : end]

1473 +

1474 finish_reason = None

1475 if stopped:

1419 - routed_experts = self._get_routed_experts(request)

1420 -

1476 # Capture finish_reason BEFORE _handle_stopped_request, which may

1477 # reset the status to WAITING for streaming requests that continue.

1478 finish_reason = request.get_finished_reason()

@@ -1592,31 +1647,6 @@ class Scheduler(SchedulerInterface):

1647 self._enqueue_waiting_request(request)

1648 return False

1649

1595 - def _get_routed_experts(self, request: Request) -> np.ndarray | None:

1596 - if not self.vllm_config.model_config.enable_return_routed_experts:

1597 - return None

1598 -

1599 - kv_blocks = self.kv_cache_manager.get_blocks(request.request_id)

1600 - block_ids = kv_blocks.get_block_ids()[self.routed_experts_attn_gid]

1601 - num_tokens = request.num_tokens - 1

1602 -

1603 - # compute slot mapping using attention group's block_size

1604 - block_ids_array = np.array(block_ids, dtype=np.int32)

1605 - num_blocks = len(block_ids)

1606 - attn_group = self.kv_cache_config.kv_cache_groups[self.routed_experts_attn_gid]

1607 - block_size = attn_group.kv_cache_spec.block_size

1608 -

1609 - # generate block offsets

1610 - block_offsets = np.arange(0, block_size)

1611 -

1612 - # compute slot mapping: slot = block_id * block_size + offset

1613 - slot_mapping = (

1614 - block_offsets.reshape((1, block_size))

1615 - + block_ids_array.reshape((num_blocks, 1)) * block_size

1616 - ).flatten()[:num_tokens]

1617 -

1618 - return self.routed_experts_reader.get_routed_experts(indices=slot_mapping)

1619 -

1650 def _update_request_with_output(

1651 self, request: Request, new_token_ids: list[int]

1652 ) -> tuple[list[int], bool]:

vllm/v1/engine/output_processor.py

+13 -7

@@ -175,6 +175,9 @@ class RequestState:

175

176 self.stats = RequestStateStats(arrival_time=arrival_time) if log_stats else None

177

178 + # Routed experts accumulation (prompt + sample chunks)

179 + self.routed_experts_chunks: list[np.ndarray] = []

180 +

181 # Stream Interval

182 self.stream_interval = stream_interval

183 self.sent_tokens_offset = 0 # Offset of sent tokens

@@ -273,7 +276,6 @@ class RequestState:

276 finish_reason: FinishReason | None,

277 stop_reason: int | str | None,

278 kv_transfer_params: dict[str, Any] | None = None,

276 - routed_experts: np.ndarray | None = None,

279 ) -> RequestOutput | PoolingRequestOutput | None:

280 finished = finish_reason is not None

281 final_only = self.output_kind == RequestOutputKind.FINAL_ONLY

@@ -314,9 +316,7 @@ class RequestState:

316 finished,

317 )

318

317 - output = self._new_completion_output(

318 - new_token_ids, finish_reason, stop_reason, routed_experts

319 - )

319 + output = self._new_completion_output(new_token_ids, finish_reason, stop_reason)

320

321 if self.parent_req is None:

322 outputs = [output]

@@ -378,7 +378,6 @@ class RequestState:

378 token_ids: list[int],

379 finish_reason: FinishReason | None,

380 stop_reason: int | str | None,

381 - routed_experts: np.ndarray | None = None,

381 ) -> CompletionOutput:

382 assert self.detokenizer is not None

383 assert self.logprobs_processor is not None

@@ -395,6 +394,11 @@ class RequestState:

394 if delta and logprobs:

395 logprobs = logprobs[-len(token_ids) :]

396

397 + # Concatenate routed experts on finish

398 + routed_experts = None

399 + if finished and self.routed_experts_chunks:

400 + routed_experts = np.concatenate(self.routed_experts_chunks, axis=0)

401 +

402 return CompletionOutput(

403 index=self.request_index,

404 text=text,

@@ -616,7 +620,10 @@ class OutputProcessor:

620 finish_reason = engine_core_output.finish_reason

621 stop_reason = engine_core_output.stop_reason

622 kv_transfer_params = engine_core_output.kv_transfer_params

619 - routed_experts = engine_core_output.routed_experts

623 + if engine_core_output.routed_experts is not None:

624 + req_state.routed_experts_chunks.append(

625 + engine_core_output.routed_experts

626 + )

627

628 if req_state.is_prefilling:

629 if engine_core_output.prefill_stats is not None:

@@ -647,7 +654,6 @@ class OutputProcessor:

654 finish_reason,

655 stop_reason,

656 kv_transfer_params,

650 - routed_experts,

657 ):

658 if req_state.streaming_input:

659 request_output.finished = False

vllm/v1/outputs.py

+78

@@ -109,6 +109,73 @@ class LogprobsTensors(NamedTuple):

109 )

110

111

112 +class RoutedExpertsTensors(NamedTuple):

113 + """Device-side snapshot of routed experts data, pending async D2H.

114 +

115 + Produced by :class:`GPUModelRunner` at the end of each async-scheduled

116 + step. The copy stream waits on the default stream, then issues

117 + non-blocking D2H via :meth:`to_cpu_nonblocking` into a pinned CPU

118 + buffer; :class:`AsyncGPUModelRunnerOutput.get_output` synchronizes

119 + the copy before the scheduler reads it.

120 +

121 + Sliced to ``total_num_scheduled_tokens`` (step-level, across all

122 + requests — NOT per-request). Both ``routing_data`` and

123 + ``slot_mapping`` must be private clones when sourced from shared

124 + capturer / prepare-input buffers, so the next forward pass /

125 + ``_prepare_inputs`` on the default stream does not race with a

126 + D2H still pending on the copy stream.

127 + """

128 +

129 + # (num_scheduled_tokens, num_layers, num_experts_per_tok)

130 + routing_data: torch.Tensor

131 + # (num_scheduled_tokens,)

132 + slot_mapping: torch.Tensor

133 +

134 + def to_cpu_nonblocking(self) -> "RoutedExpertsTensors":

135 + """Issue non-blocking D2H on the current stream.

136 +

137 + NOTE: ``non_blocking=True`` only delivers true overlap when the

138 + CPU target is pinned. The current fallback here allocates a

139 + new pageable CPU tensor per call, which silently degrades to a

140 + synchronous copy; acceptable because the sync happens on the

141 + dedicated copy stream, not the default stream.

142 + """

143 + if self.routing_data.device.type == "cpu":

144 + return self

145 + return RoutedExpertsTensors(

146 + self.routing_data.to("cpu", non_blocking=True),

147 + self.slot_mapping.to("cpu", non_blocking=True),

148 + )

149 +

150 + def tolists(self) -> "RoutedExpertsLists":

151 + """Convert to the numpy-backed form consumed by the scheduler.

152 +

153 + ``.cpu()`` is a no-op when the tensor is already on CPU, so this

154 + is cheap for the post-D2H case; for raw device tensors it will

155 + synchronously block, which is only reached in tests.

156 + """

157 + return RoutedExpertsLists(

158 + self.routing_data.cpu().numpy(),

159 + self.slot_mapping.cpu().numpy(),

160 + )

161 +

162 +

163 +class RoutedExpertsLists(NamedTuple):

164 + """CPU-side routed experts, the form :meth:`RoutedExpertsManager.store_batch`

165 + consumes.

166 +

167 + Batched per scheduler step: the leading dim is the number of tokens

168 + scheduled across all requests in this step (``total_num_scheduled_tokens``),

169 + not per-request tokens. ``slot_mapping[i]`` tells the scheduler which

170 + physical KV-cache slot row ``i`` of ``routing_data`` belongs to.

171 + """

172 +

173 + # (num_scheduled_tokens, num_layers, num_experts_per_tok)

174 + routing_data: np.ndarray

175 + # (num_scheduled_tokens,)

176 + slot_mapping: np.ndarray

177 +

178 +

179 # [num_reqs, <dynamic>]

180 # The shape of each element depends on the pooler used

181 PoolerOutput: TypeAlias = torch.Tensor | list[torch.Tensor] | list[torch.Tensor | None]

@@ -201,6 +268,17 @@ class ModelRunnerOutput:

268 # information related to cudagraph execution

269 cudagraph_stats: CUDAGraphStat | None = None

270

271 + # Per-step routed experts data captured by the worker.

272 + # ``routing_data`` shape: (num_scheduled_tokens, num_layers,

273 + # num_experts_per_tok); expert IDs as uint8/uint16.

274 + # ``slot_mapping`` shape: (num_scheduled_tokens,); physical KV-cache

275 + # slot for each row of routing_data.

276 + # ``num_scheduled_tokens`` is step-level (total across all requests

277 + # in this step), not per-request. The scheduler persists this into

278 + # its slot buffer via ``slot_buffer[slot_mapping] = routing_data``.

279 + # ``None`` when ``enable_return_routed_experts`` is off.

280 + routed_experts: RoutedExpertsLists | None = None

281 +

282

283 # ModelRunnerOutput wrapper for async scheduling.

284 class AsyncModelRunnerOutput(ABC):

vllm/v1/worker/gpu_model_runner.py

+115 -35

@@ -157,6 +157,8 @@ from vllm.v1.outputs import (

157 LogprobsTensors,

158 ModelRunnerOutput,

159 PoolerOutput,

160 + RoutedExpertsLists,

161 + RoutedExpertsTensors,

162 SamplerOutput,

163 make_empty_encoder_model_runner_output,

164 )

@@ -236,6 +238,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

238 invalid_req_indices: list[int],

239 async_output_copy_stream: torch.cuda.Stream,

240 vocab_size: int,

241 + routed_experts: RoutedExpertsTensors | None = None,

242 ):

243 self._model_runner_output = model_runner_output

244 self._invalid_req_indices = invalid_req_indices

@@ -248,6 +251,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

251 self._sampled_token_ids = sampled_token_ids

252 self.vocab_size = vocab_size

253 self._logprobs_tensors = logprobs_tensors

254 + self._routed_experts = routed_experts

255

256 # Initiate the copy on a separate stream, but do not synchronize it.

257 default_stream = torch.cuda.current_stream()

@@ -261,6 +265,11 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

265 if self._logprobs_tensors

266 else None

267 )

268 + self._routed_experts_cpu = (

269 + self._routed_experts.to_cpu_nonblocking()

270 + if self._routed_experts is not None

271 + else None

272 + )

273 self.async_copy_ready_event.record()

274

275 def get_output(self) -> ModelRunnerOutput:

@@ -292,6 +301,11 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

301 output = self._model_runner_output

302 output.sampled_token_ids = valid_sampled_token_ids

303 output.logprobs = logprobs_lists

304 +

305 + if self._routed_experts_cpu is not None:

306 + output.routed_experts = self._routed_experts_cpu.tolists()

307 + del self._routed_experts

308 +

309 return output

310

311

@@ -2192,9 +2206,18 @@ class GPUModelRunner(

2206 slot_mapping_gid_0 = slot_mappings[0]

2207

2208 if self.routed_experts_initialized:

2209 + # Copy this step's attention slot_mapping into our private

2210 + # device buffer. The shared ``slot_mappings[attn_gid]`` is

2211 + # owned by the attention block table and will be overwritten

2212 + # by the next ``_prepare_inputs``; we need a stable snapshot

2213 + # because the async D2H may still be in flight on the copy

2214 + # stream when the next step runs.

2215 attn_gid = self.routed_experts_attn_gid

2216 slot_mapping_attn = slot_mappings[attn_gid]

2197 - self.slot_mapping = slot_mapping_attn[:num_tokens].cpu().numpy()

2217 + self.routed_experts_slot_mapping_device[:num_tokens].copy_(

2218 + slot_mapping_attn[:num_tokens]

2219 + )

2220 +

2221 num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu_tensor[

2222 :num_reqs_padded

2223 ]

@@ -3478,6 +3501,21 @@ class GPUModelRunner(

3501 invalid_req_indices = []

3502 logprobs_lists = None

3503 if not self.use_async_scheduling:

3504 + # Sync scheduling: issue routed experts D2H into the pinned

3505 + # CPU buffer BEFORE ``_to_list`` below. ``_to_list`` does

3506 + # ``event.synchronize()`` on the async copy stream which

3507 + # waits for every D2H queued on the default stream since

3508 + # the last sync, so this enqueue is naturally covered

3509 + # without requiring its own synchronize.

3510 + if self.routed_experts_initialized:

3511 + buf = self.routed_experts_capturer.get_device_buffer()

3512 + total = scheduler_output.total_num_scheduled_tokens

3513 + self.routed_experts_cpu[:total].copy_(buf[:total], non_blocking=True)

3514 + self.routed_experts_slot_mapping_cpu[:total].copy_(

3515 + self.routed_experts_slot_mapping_device[:total],

3516 + non_blocking=True,

3517 + )

3518 +

3519 # Get the valid generated tokens.

3520 max_gen_len = sampled_token_ids.shape[-1]

3521 if max_gen_len == 1:

@@ -3878,11 +3916,7 @@ class GPUModelRunner(

3916 )

3917

3918 if self.routed_experts_initialized:

3881 - capturer = RoutedExpertsCapturer.get_instance()

3882 - if capturer is not None:

3883 - capturer.clear_buffer() # noqa

3884 - else:

3885 - logger.error("RoutedExpertsCapturer not initialized.")

3919 + self.routed_experts_capturer.clear_buffer()

3920

3921 # If ngram_gpu is used, we need to copy the scheduler_output to avoid

3922 # the modification has influence on the scheduler_output in engine core process.

@@ -4415,13 +4449,6 @@ class GPUModelRunner(

4449 self.kv_connector_output = None

4450

4451 with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):

4418 - if self.routed_experts_initialized:

4419 - capturer = RoutedExpertsCapturer.get_instance()

4420 - if capturer is not None:

4421 - capturer.save_captured_experts(indices=self.slot_mapping) # noqa

4422 - else:

4423 - logger.error("RoutedExpertsCapturer not initialized.")

4424 -

4452 output = ModelRunnerOutput(

4453 req_ids=req_ids_output_copy,

4454 req_id_to_index=req_id_to_index_output_copy,

@@ -4434,14 +4461,47 @@ class GPUModelRunner(

4461 else None,

4462 num_nans_in_logits=num_nans_in_logits,

4463 cudagraph_stats=cudagraph_stats,

4464 + routed_experts=None,

4465 )

4466

4467 if not self.use_async_scheduling:

4468 + if self.routed_experts_initialized:

4469 + # Sync path: D2H was issued in ``_bookkeeping_sync`` and

4470 + # synchronized by ``_to_list``'s event.synchronize(), so

4471 + # the pinned buffers are ready to be wrapped as numpy.

4472 + total = scheduler_output.total_num_scheduled_tokens

4473 + output.routed_experts = RoutedExpertsLists(

4474 + routing_data=self.routed_experts_cpu[:total].numpy(),

4475 + slot_mapping=self.routed_experts_slot_mapping_cpu[:total].numpy(),

4476 + )

4477 return output

4478

4479 with record_function_or_nullcontext(

4480 "gpu_model_runner: AsyncGPUModelRunnerOutput"

4481 ):

4482 + # Async path: produce a device-side snapshot that the async

4483 + # copy stream can D2H later. Both tensors must be private

4484 + # clones because:

4485 + # - ``routing_data`` source is the shared capturer buffer,

4486 + # which is ``clear_buffer()``-ed at the start of the

4487 + # next step on the default stream.

4488 + # - ``slot_mapping`` source is our own

4489 + # ``routed_experts_slot_mapping_device``, which the

4490 + # next ``_prepare_inputs`` overwrites on the default

4491 + # stream while the D2H is still pending on the copy

4492 + # stream.

4493 + # Without clones, the copy stream would read torn data.

4494 + routed_experts_snapshot = None

4495 + if self.routed_experts_initialized:

4496 + buf = self.routed_experts_capturer.get_device_buffer()

4497 + total = scheduler_output.total_num_scheduled_tokens

4498 + routed_experts_snapshot = RoutedExpertsTensors(

4499 + routing_data=buf[:total].clone(),

4500 + slot_mapping=self.routed_experts_slot_mapping_device[

4501 + :total

4502 + ].clone(),

4503 + )

4504 +

4505 async_output = AsyncGPUModelRunnerOutput(

4506 model_runner_output=output,

4507 sampled_token_ids=sampler_output.sampled_token_ids,

@@ -4449,6 +4509,7 @@ class GPUModelRunner(

4509 invalid_req_indices=invalid_req_indices,

4510 async_output_copy_stream=self._get_or_create_async_output_copy_stream(),

4511 vocab_size=self.input_batch.vocab_size,

4512 + routed_experts=routed_experts_snapshot,

4513 )

4514 with record_function_or_nullcontext(

4515 "gpu_model_runner: set_async_sampled_token_ids"

@@ -7042,9 +7103,16 @@ class GPUModelRunner(

7103 kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks)

7104

7105 def _get_attention_kv_cache_gid(self) -> int:

7045 - """Find the KV cache group index for attention layers."""

7106 + """Find the KV cache group index for attention layers.

7107 +

7108 + Must match :attr:`RoutedExpertsManager.attn_gid` in the scheduler:

7109 + both pick the first ``FullAttentionSpec`` group so hybrid models

7110 + (Mamba / linear-attention layers that use other AttentionSpec

7111 + subclasses) end up indexing the same slot layout on both sides.

7112 + Falls back to 0 only for legacy single-group configs.

7113 + """

7114 for gid, group in enumerate(self.kv_cache_config.kv_cache_groups):

7047 - if isinstance(group.kv_cache_spec, AttentionSpec):

7115 + if isinstance(group.kv_cache_spec, FullAttentionSpec):

7116 return gid

7117 return 0

7118

@@ -7053,29 +7121,41 @@ class GPUModelRunner(

7121 "Initializing routed experts capturer, enable_return_routed_experts: %s",

7122 self.model_config.enable_return_routed_experts,

7123 )

7056 - routed_experts_capturer = RoutedExpertsCapturer.create()

7057 - self.routed_experts_attn_gid = self._get_attention_kv_cache_gid()

7058 - min_block_size = min(

7059 - [

7060 - group.kv_cache_spec.block_size

7061 - for group in self.kv_cache_config.kv_cache_groups

7062 - ]

7063 - )

7064 - num_groups = len(self.kv_cache_config.kv_cache_groups)

7065 - self.max_num_kv_tokens = (

7066 - self.kv_cache_config.num_blocks // num_groups

7067 - ) * min_block_size

7068 - dcp_size = self.vllm_config.parallel_config.decode_context_parallel_size

7069 - pcp_size = self.vllm_config.parallel_config.prefill_context_parallel_size

7070 - if pcp_size * dcp_size > 1:

7071 - self.max_num_kv_tokens *= pcp_size * dcp_size

7072 -

7073 - routed_experts_capturer.init_buffer(

7124 + self.routed_experts_capturer = RoutedExpertsCapturer(

7125 max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens,

7075 - max_num_kv_tokens=self.max_num_kv_tokens,

7126 vllm_config=self.vllm_config,

7127 )

7078 - self._bind_routed_experts_capturer(routed_experts_capturer)

7128 + self.routed_experts_attn_gid = self._get_attention_kv_cache_gid()

7129 + self._bind_routed_experts_capturer(self.routed_experts_capturer)

7130 +

7131 + # Pinned CPU buffer for non-blocking D2H of ``routing_data`` on

7132 + # the sync scheduling path. Shape / dtype mirror the device

7133 + # capturer exactly so ``copy_`` is a straight memcpy.

7134 + self.routed_experts_cpu = torch.empty(

7135 + self.routed_experts_capturer.device_buffer.shape,

7136 + dtype=self.routed_experts_capturer.device_buffer.dtype,

7137 + device="cpu",

7138 + pin_memory=self.pin_memory,

7139 + )

7140 + # ``slot_mapping`` dtype is fixed to int64 by

7141 + # ``block_table.slot_mapping``; we mirror that here.

7142 + max_tokens = self.scheduler_config.max_num_batched_tokens

7143 + self.routed_experts_slot_mapping_cpu = torch.empty(

7144 + (max_tokens,),

7145 + dtype=torch.int64,

7146 + device="cpu",

7147 + pin_memory=self.pin_memory,

7148 + )

7149 + # Private device buffer so the shared ``block_table.slot_mapping``

7150 + # can be overwritten by the next ``_prepare_inputs`` while the

7151 + # D2H is still pending on the copy stream. Written in

7152 + # ``_prepare_inputs``, read in ``_bookkeeping_sync`` (sync path)

7153 + # or cloned into a snapshot (async path).

7154 + self.routed_experts_slot_mapping_device = torch.empty(

7155 + (max_tokens,),

7156 + dtype=torch.int64,

7157 + device=self.device,

7158 + )

7159 self.routed_experts_initialized = True

7160

7161 def _bind_routed_experts_capturer(self, capturer: RoutedExpertsCapturer) -> None: