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 10, 07:45 UTC
[Feature] Support recording expert indices for rollout router replay (#28284)
Signed-off-by: xhx1022 <1737006628@qq.com> Signed-off-by: Hongxin Xu <70438206+xhx1022@users.noreply.github.com> Signed-off-by: arlenxu <arlenxu@tencent.com> Co-authored-by: 22quinn <33176974+22quinn@users.noreply.github.com> Co-authored-by: arlenxu <arlenxu@tencent.com>
Hongxin Xu committed
Jan 12, 2026 at 22:23 UTC
49e6b86c91cddc4a4b2152c6d8834780bbd50deb
11 files changed
+463
-3
vllm/config/model.py
+2
@@ -198,6 +198,8 @@ class ModelConfig:
198
graph and always execute the model in eager mode. If False, we will use
199
CUDA graph and eager execution in hybrid for maximal performance and
200
flexibility."""
201
+ enable_return_routed_experts: bool = False
202
+ """Whether to return routed experts."""
203
max_logprobs: int = 20
204
"""Maximum number of log probabilities to return when `logprobs` is
205
specified in `SamplingParams`. The default value comes the default for the
vllm/config/vllm.py
+1
@@ -1352,6 +1352,7 @@ class VllmConfig:
1352
f"disable_custom_all_reduce={self.parallel_config.disable_custom_all_reduce}, " # noqa
1353
f"quantization={self.model_config.quantization}, "
1354
f"enforce_eager={self.model_config.enforce_eager}, "
1355
+ f"enable_return_routed_experts={self.model_config.enable_return_routed_experts}, " # noqa
1356
f"kv_cache_dtype={self.cache_config.cache_dtype}, "
1357
f"device_config={self.device_config.device}, "
1358
f"structured_outputs_config={self.structured_outputs_config!r}, "
vllm/engine/arg_utils.py
+6
@@ -354,6 +354,7 @@ class EngineArgs:
354
"""Arguments for vLLM engine."""
355
356
model: str = ModelConfig.model
357
+ enable_return_routed_experts: bool = ModelConfig.enable_return_routed_experts
358
model_weights: str = ModelConfig.model_weights
359
served_model_name: str | list[str] | None = ModelConfig.served_model_name
360
tokenizer: str | None = ModelConfig.tokenizer
@@ -657,6 +658,10 @@ class EngineArgs:
658
**model_kwargs["allow_deprecated_quantization"],
659
)
660
model_group.add_argument("--enforce-eager", **model_kwargs["enforce_eager"])
661
+ model_group.add_argument(
662
+ "--enable-return-routed-experts",
663
+ **model_kwargs["enable_return_routed_experts"],
664
+ )
665
model_group.add_argument("--max-logprobs", **model_kwargs["max_logprobs"])
666
model_group.add_argument("--logprobs-mode", **model_kwargs["logprobs_mode"])
667
model_group.add_argument(
@@ -1239,6 +1244,7 @@ class EngineArgs:
1244
quantization=self.quantization,
1245
allow_deprecated_quantization=self.allow_deprecated_quantization,
1246
enforce_eager=self.enforce_eager,
1247
+ enable_return_routed_experts=self.enable_return_routed_experts,
1248
max_logprobs=self.max_logprobs,
1249
logprobs_mode=self.logprobs_mode,
1250
disable_sliding_window=self.disable_sliding_window,
vllm/entrypoints/llm.py
+3
@@ -158,6 +158,7 @@ class LLM:
158
enforce_eager: Whether to enforce eager execution. If True, we will
159
disable CUDA graph and always execute the model in eager mode.
160
If False, we will use CUDA graph and eager execution in hybrid.
161
+ enable_return_routed_experts: Whether to return routed experts.
162
disable_custom_all_reduce: See
163
[ParallelConfig][vllm.config.ParallelConfig].
164
hf_token: The token to use as HTTP bearer authorization for remote files
@@ -209,6 +210,7 @@ class LLM:
210
swap_space: float = 4,
211
cpu_offload_gb: float = 0,
212
enforce_eager: bool = False,
213
+ enable_return_routed_experts: bool = False,
214
disable_custom_all_reduce: bool = False,
215
hf_token: bool | str | None = None,
216
hf_overrides: HfOverrides | None = None,
@@ -317,6 +319,7 @@ class LLM:
319
swap_space=swap_space,
320
cpu_offload_gb=cpu_offload_gb,
321
enforce_eager=enforce_eager,
322
+ enable_return_routed_experts=enable_return_routed_experts,
323
disable_custom_all_reduce=disable_custom_all_reduce,
324
hf_token=hf_token,
325
hf_overrides=hf_overrides,
vllm/model_executor/layers/fused_moe/layer.py
+22
@@ -35,6 +35,9 @@ from vllm.model_executor.layers.fused_moe.fused_moe_router import FusedMoERouter
35
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
36
init_aiter_topK_meta_data,
37
)
38
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
39
+ RoutedExpertsCapturer,
40
+)
41
from vllm.model_executor.layers.fused_moe.routing_simulator import RoutingSimulator
42
from vllm.model_executor.layers.quantization.base_config import (
43
QuantizationConfig,
@@ -701,6 +704,13 @@ class FusedMoE(CustomOp):
704
def shared_experts(self) -> torch.nn.Module | None:
705
return None
706
707
+ @property
708
+ def layer_id(self):
709
+ # Delayed import to avoid circular dependency
710
+ from vllm.model_executor.models.utils import extract_layer_index
711
+
712
+ return extract_layer_index(self.layer_name)
713
+
714
@property
715
def gate(self) -> torch.nn.Module | None:
716
return None
@@ -1650,6 +1660,18 @@ class FusedMoE(CustomOp):
1660
1661
assert topk_ids.dtype == indices_type or indices_type is None
1662
1663
+ if (
1664
+ self.vllm_config.model_config is not None
1665
+ and self.vllm_config.model_config.enable_return_routed_experts
1666
+ ):
1667
+ # In dummy runs, the capturer is not initialized.
1668
+ capturer = RoutedExpertsCapturer.get_instance()
1669
+ if capturer is not None: # in dummmy_run may be None
1670
+ capturer.capture( # noqa
1671
+ layer_id=self.layer_id,
1672
+ topk_ids=topk_ids,
1673
+ )
1674
+
1675
return topk_weights, topk_ids
1676
1677
def must_reduce_shared_expert_outputs(self) -> bool:
vllm/model_executor/layers/fused_moe/routed_experts_capturer.py
new
+324
@@ -0,0 +1,324 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# Adapted from
4
+# https://github.com/sgl-project/sglang/blob/bed301a5acaa9577c9aa706468bdf242f6a43051/python/sglang/srt/layers/moe/routed_experts_capturer.py
5
+
6
+from __future__ import annotations
7
+
8
+import fcntl
9
+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
16
+
17
+import numpy as np
18
+import torch
19
+
20
+from vllm.config import ModelConfig
21
+from vllm.distributed import get_tensor_model_parallel_rank
22
+
23
+logger = logging.getLogger(__name__)
24
+
25
+# Constants
26
+_TMP_DIR = tempfile.gettempdir()
27
+_LOCK_FILE_PREFIX = os.path.join(_TMP_DIR, "vllm_routed_experts")
28
+_BUFFER_PREFIX = "vllm_routed_experts_buffer"
29
+
30
+# Global singleton instances
31
+_global_experts_capturer: RoutedExpertsCapturer | None = None
32
+_global_experts_reader: RoutedExpertsReader | None = None
33
+
34
+
35
+@contextmanager
36
+def _file_lock(lock_file: str, mode: str = "wb+") -> Generator[None, None, None]:
37
+ """Context manager for file-based locking."""
38
+ with open(lock_file, mode) as fp:
39
+ fcntl.flock(fp, fcntl.LOCK_EX)
40
+ try:
41
+ yield
42
+ finally:
43
+ fcntl.flock(fp, fcntl.LOCK_UN)
44
+
45
+
46
+def _create_or_attach_shared_memory(
47
+ name: str, size: int, lock_file: str
48
+) -> shared_memory.SharedMemory:
49
+ """Create or attach to shared memory with proper locking."""
50
+ # Ensure lock file exists before acquiring lock
51
+ with open(lock_file, "wb"):
52
+ pass
53
+
54
+ with _file_lock(lock_file):
55
+ try:
56
+ shm = shared_memory.SharedMemory(name=name, create=True, size=size)
57
+ except FileExistsError:
58
+ shm = shared_memory.SharedMemory(name=name, create=False, size=size)
59
+
60
+ if shm.size != size:
61
+ logger.warning(
62
+ "Shared memory %s size mismatch; recreating",
63
+ name,
64
+ )
65
+ shm.close()
66
+ shm.unlink()
67
+ try:
68
+ shm = shared_memory.SharedMemory(name=name, create=True, size=size)
69
+ logger.info("Created shared memory %s", name)
70
+ except FileExistsError:
71
+ shm = shared_memory.SharedMemory(name=name, create=False, size=size)
72
+ logger.info("Linked to existing shared memory %s", name)
73
+
74
+ return shm
75
+
76
+
77
+class RoutedExpertsCapturer:
78
+ """
79
+ Capturer for routed experts with device and optional shared memory buffer.
80
+
81
+ This class captures expert routing decisions during model forward passes
82
+ and optionally stores them in shared memory for cross-process access.
83
+ """
84
+
85
+ _instance: RoutedExpertsCapturer | None = None
86
+
87
+ def __init__(self) -> None:
88
+ self._device_buffer: torch.Tensor | None = None
89
+ self._shm: shared_memory.SharedMemory | None = None
90
+ self._host_buffer_view: np.ndarray | None = None
91
+ self._lock_file: str | None = None
92
+ self._shm_name: str | None = None
93
+
94
+ @classmethod
95
+ def create(cls) -> RoutedExpertsCapturer:
96
+ """Create a global singleton instance."""
97
+ global _global_experts_capturer
98
+ if _global_experts_capturer is not None:
99
+ raise RuntimeError("Experts capturer already created.")
100
+
101
+ _global_experts_capturer = cls()
102
+ return _global_experts_capturer
103
+
104
+ @staticmethod
105
+ def get_instance() -> RoutedExpertsCapturer | None:
106
+ """Get the global singleton instance."""
107
+ return _global_experts_capturer
108
+
109
+ def init_buffer(
110
+ self,
111
+ max_num_batched_tokens: int,
112
+ max_num_kv_tokens: int,
113
+ model_config: ModelConfig,
114
+ instance_id: str,
115
+ ) -> None:
116
+ """
117
+ Initialize the device buffer and optionally shared memory buffer.
118
+
119
+ Args:
120
+ max_num_batched_tokens: Maximum number of tokens in a batch.
121
+ max_num_kv_tokens: Maximum number of KV tokens for shared memory.
122
+ model_config: Model configuration containing layer and expert info.
123
+ instance_id: Unique identifier for the shared memory buffer.
124
+ """
125
+
126
+ if self._device_buffer is not None:
127
+ raise RuntimeError("Device buffer has already been initialized")
128
+
129
+ hf_config = model_config.hf_text_config
130
+ num_layers = hf_config.num_hidden_layers
131
+ num_experts_per_tok = hf_config.num_experts_per_tok
132
+
133
+ # Initialize device buffer
134
+ self._device_buffer = torch.zeros(
135
+ (max_num_batched_tokens, num_layers, num_experts_per_tok),
136
+ dtype=torch.int32,
137
+ device="cuda",
138
+ )
139
+
140
+ if get_tensor_model_parallel_rank() != 0:
141
+ return
142
+
143
+ # Initialize shared memory
144
+ shape = (max_num_kv_tokens, num_layers, num_experts_per_tok)
145
+ buffer_size = int(np.prod(shape)) * np.dtype(np.int32).itemsize
146
+
147
+ self._lock_file = f"{_LOCK_FILE_PREFIX}_{instance_id}.lock"
148
+ self._shm_name = f"{_BUFFER_PREFIX}_{instance_id}"
149
+
150
+ self._shm = _create_or_attach_shared_memory(
151
+ self._shm_name, buffer_size, self._lock_file
152
+ )
153
+ self._host_buffer_view = np.ndarray(shape, dtype=np.int32, buffer=self._shm.buf)
154
+ self._host_buffer_view.fill(0)
155
+
156
+ logger.debug(
157
+ "Created shared memory buffer '%s' with shape %s",
158
+ self._shm.name,
159
+ shape,
160
+ )
161
+
162
+ def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None:
163
+ """
164
+ Capture expert routing decisions for a specific layer.
165
+
166
+ Args:
167
+ layer_id: The layer index.
168
+ topk_ids: Tensor of shape (batch_size, num_routed_experts).
169
+ """
170
+ if self._device_buffer is None:
171
+ raise RuntimeError("Buffer not initialized. Call init_buffer() first.")
172
+
173
+ if layer_id >= self._device_buffer.shape[1]:
174
+ return
175
+
176
+ batch_size = topk_ids.shape[0]
177
+ self._device_buffer[:batch_size, layer_id, :] = topk_ids
178
+
179
+ def clear_buffer(self) -> None:
180
+ """Clear the device buffer."""
181
+ if self._device_buffer is not None:
182
+ self._device_buffer.zero_()
183
+
184
+ def save_captured_experts(self, indices: np.ndarray) -> None:
185
+ """
186
+ Save captured experts from device buffer to shared memory.
187
+
188
+ Args:
189
+ indices: Array of indices indicating where to store the data.
190
+ """
191
+ if get_tensor_model_parallel_rank() != 0:
192
+ return
193
+ if self._lock_file is None:
194
+ raise RuntimeError("Shared memory not initialized.")
195
+ if self._host_buffer_view is None:
196
+ return
197
+ if self._device_buffer is None:
198
+ raise RuntimeError("Device buffer not initialized.")
199
+
200
+ num_tokens = len(indices)
201
+ data = self._device_buffer[:num_tokens, :, :].cpu().numpy()
202
+
203
+ with _file_lock(self._lock_file):
204
+ self._host_buffer_view[indices, :, :] = data
205
+
206
+ def cleanup(self) -> None:
207
+ """Explicitly clean up shared memory resources."""
208
+ if self._shm is not None:
209
+ try:
210
+ self._shm.close()
211
+ self._shm.unlink()
212
+ except Exception:
213
+ logger.debug("Exception during cleanup for capturer", exc_info=True)
214
+ finally:
215
+ self._shm = None
216
+
217
+ def __del__(self) -> None:
218
+ """Clean up shared memory on destruction."""
219
+ self.cleanup()
220
+
221
+
222
+class RoutedExpertsReader:
223
+ """
224
+ Reader for routed experts from shared memory.
225
+
226
+ This class attaches to shared memory created by RoutedExpertsCapturer
227
+ and reads expert routing decisions.
228
+ """
229
+
230
+ _instance: RoutedExpertsReader | None = None
231
+
232
+ def __init__(self) -> None:
233
+ self._shm: shared_memory.SharedMemory | None = None
234
+ self._host_buffer_view: np.ndarray | None = None
235
+ self._lock_file: str | None = None
236
+
237
+ @classmethod
238
+ def create(cls) -> RoutedExpertsReader:
239
+ """Create a global singleton instance."""
240
+ global _global_experts_reader
241
+ if _global_experts_reader is not None:
242
+ raise RuntimeError("Experts reader already created.")
243
+
244
+ _global_experts_reader = cls()
245
+ return _global_experts_reader
246
+
247
+ @staticmethod
248
+ def get_instance() -> RoutedExpertsReader | None:
249
+ """Get the global singleton instance."""
250
+ if _global_experts_reader is None:
251
+ logger.info("Experts reader not initialized.")
252
+ return _global_experts_reader
253
+
254
+ def attach_buffer(
255
+ self,
256
+ max_num_kv_tokens: int,
257
+ model_config: ModelConfig,
258
+ instance_id: str,
259
+ ) -> None:
260
+ """
261
+ Attach to an existing shared memory buffer.
262
+
263
+ Args:
264
+ max_num_kv_tokens: Maximum number of KV tokens.
265
+ model_config: Model configuration.
266
+ instance_id: Unique identifier for the shared memory buffer.
267
+ """
268
+ if self._shm is not None:
269
+ logger.warning("Already attached to shared memory buffer.")
270
+ return # Already attached
271
+
272
+ hf_config = model_config.hf_text_config
273
+ shape = (
274
+ max_num_kv_tokens,
275
+ hf_config.num_hidden_layers,
276
+ hf_config.num_experts_per_tok,
277
+ )
278
+
279
+ self._lock_file = f"{_LOCK_FILE_PREFIX}_{instance_id}.lock"
280
+ shm_name = f"{_BUFFER_PREFIX}_{instance_id}"
281
+
282
+ with _file_lock(self._lock_file, mode="rb+"):
283
+ # Avoid resource_tracker registering the shared memory
284
+ with patch(
285
+ "multiprocessing.resource_tracker.register",
286
+ lambda *args, **kwargs: None,
287
+ ):
288
+ self._shm = shared_memory.SharedMemory(name=shm_name)
289
+
290
+ self._host_buffer_view = np.ndarray(
291
+ shape, dtype=np.int32, buffer=self._shm.buf
292
+ )
293
+
294
+ def get_routed_experts(self, indices: np.ndarray) -> np.ndarray:
295
+ """
296
+ Read routed expert data from shared memory.
297
+
298
+ Args:
299
+ indices: Array of indices to read.
300
+
301
+ Returns:
302
+ Copy of the expert routing data for the given indices.
303
+ """
304
+ if self._host_buffer_view is None:
305
+ raise RuntimeError("Buffer not attached. Call attach_buffer() first.")
306
+ if self._lock_file is None:
307
+ raise RuntimeError("Lock file not initialized.")
308
+
309
+ with _file_lock(self._lock_file, mode="rb+"):
310
+ return self._host_buffer_view[indices, :, :].copy()
311
+
312
+ def cleanup(self) -> None:
313
+ """Explicitly clean up resources (close without unlink)."""
314
+ if self._shm is not None:
315
+ try:
316
+ self._shm.close()
317
+ except Exception:
318
+ logger.debug("Exception during cleanup for reader", exc_info=True)
319
+ finally:
320
+ self._shm = None
321
+
322
+ def __del__(self) -> None:
323
+ """Close shared memory on destruction (do not unlink)."""
324
+ self.cleanup()
vllm/outputs.py
+3
@@ -6,6 +6,7 @@ from collections.abc import Sequence as GenericSequence
6
from dataclasses import dataclass
7
from typing import Any, Generic
8
9
+import numpy as np
10
import torch
11
from typing_extensions import TypeVar
12
@@ -42,6 +43,7 @@ class CompletionOutput:
43
token_ids: GenericSequence[int]
44
cumulative_logprob: float | None
45
logprobs: SampleLogprobs | None
46
+ routed_experts: np.ndarray | None = None # [seq_len,layer_num,topk]
47
finish_reason: str | None = None
48
stop_reason: int | str | None = None
49
lora_request: LoRARequest | None = None
@@ -54,6 +56,7 @@ class CompletionOutput:
56
f"CompletionOutput(index={self.index}, "
57
f"text={self.text!r}, "
58
f"token_ids={self.token_ids}, "
59
+ f"routed_experts={self.routed_experts}, "
60
f"cumulative_logprob={self.cumulative_logprob}, "
61
f"logprobs={self.logprobs}, "
62
f"finish_reason={self.finish_reason}, "
vllm/v1/core/sched/scheduler.py
+50
-1
@@ -6,6 +6,8 @@ from collections import defaultdict
6
from collections.abc import Iterable
7
from typing import Any
8
9
+import numpy as np
10
+
11
from vllm import envs
12
from vllm.compilation.cuda_graph import CUDAGraphStat
13
from vllm.config import VllmConfig
@@ -24,6 +26,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import (
26
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
27
from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
28
from vllm.logger import init_logger
29
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
30
+ RoutedExpertsReader,
31
+)
32
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
33
from vllm.v1.core.encoder_cache_manager import (
34
EncoderCacheManager,
@@ -219,11 +224,31 @@ class Scheduler(SchedulerInterface):
224
)
225
self.use_pp = self.parallel_config.pipeline_parallel_size > 1
226
self.use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER
222
-
227
self.perf_metrics: ModelMetrics | None = None
228
if self.log_stats and vllm_config.observability_config.enable_mfu_metrics:
229
self.perf_metrics = ModelMetrics(vllm_config)
230
231
+ if self.vllm_config.model_config.enable_return_routed_experts:
232
+ assert self.dcp_world_size == 1 and self.pcp_world_size == 1, (
233
+ "enable_return_routed_experts does not support context parallelism "
234
+ "(dcp_world_size > 1 or pcp_world_size > 1)"
235
+ )
236
+
237
+ self.routed_experts_reader = RoutedExpertsReader.create()
238
+
239
+ assert len(kv_cache_config.kv_cache_groups) > 0, (
240
+ "enable_return_routed_experts requires at least one kv cache group"
241
+ )
242
+ self.max_num_kv_tokens = (
243
+ kv_cache_config.num_blocks // len(kv_cache_config.kv_cache_groups) + 1
244
+ ) * self.block_size
245
+
246
+ self.routed_experts_reader.attach_buffer(
247
+ max_num_kv_tokens=self.max_num_kv_tokens,
248
+ model_config=self.vllm_config.model_config,
249
+ instance_id=self.vllm_config.instance_id,
250
+ )
251
+
252
def schedule(self) -> SchedulerOutput:
253
# NOTE(woosuk) on the scheduling algorithm:
254
# There's no "decoding phase" nor "prefill phase" in the scheduler.
@@ -1162,7 +1187,30 @@ class Scheduler(SchedulerInterface):
1187
request.status = RequestStatus.FINISHED_STOPPED
1188
stopped = True
1189
1190
+ routed_experts = None
1191
if stopped:
1192
+ if self.vllm_config.model_config.enable_return_routed_experts:
1193
+ kv_blocks = self.kv_cache_manager.get_blocks(request.request_id)
1194
+ block_ids = kv_blocks.get_block_ids()[0]
1195
+ num_tokens = request.num_tokens - 1
1196
+
1197
+ # compute slot mapping
1198
+ block_ids_array = np.array(block_ids, dtype=np.int32)
1199
+ num_blocks = len(block_ids)
1200
+ block_size = self.block_size
1201
+
1202
+ # generate block offsets
1203
+ block_offsets = np.arange(0, block_size)
1204
+
1205
+ # compute slot mapping: slot = block_id * block_size + offset
1206
+ slot_mapping = (
1207
+ block_offsets.reshape((1, block_size))
1208
+ + block_ids_array.reshape((num_blocks, 1)) * block_size
1209
+ ).flatten()[:num_tokens]
1210
+
1211
+ routed_experts = self.routed_experts_reader.get_routed_experts(
1212
+ indices=slot_mapping
1213
+ )
1214
kv_transfer_params = self._free_request(request)
1215
if status_before_stop == RequestStatus.RUNNING:
1216
stopped_running_reqs.add(request)
@@ -1209,6 +1257,7 @@ class Scheduler(SchedulerInterface):
1257
kv_transfer_params=kv_transfer_params,
1258
trace_headers=request.trace_headers,
1259
num_cached_tokens=request.num_cached_tokens,
1260
+ routed_experts=routed_experts,
1261
num_nans_in_logits=request.num_nans_in_logits,
1262
)
1263
)
vllm/v1/engine/__init__.py
+2
-1
@@ -7,6 +7,7 @@ from collections.abc import Mapping
7
from typing import Any
8
9
import msgspec
10
+import numpy as np
11
import torch
12
13
from vllm.lora.request import LoRARequest
@@ -139,7 +140,7 @@ class EngineCoreOutput(
140
trace_headers: Mapping[str, str] | None = None
141
# The number of tokens with prefix cache hits.
142
num_cached_tokens: int = 0
142
-
143
+ routed_experts: np.ndarray | None = None
144
# The number of NaNs in logits.
145
# A value greater than 0 indicates that the output is corrupted.
146
num_nans_in_logits: int = 0
vllm/v1/engine/output_processor.py
+9
-1
@@ -7,6 +7,7 @@ from collections.abc import Iterable
7
from dataclasses import dataclass
8
from typing import Any, cast
9
10
+import numpy as np
11
import torch
12
13
from vllm.lora.request import LoRARequest
@@ -213,6 +214,7 @@ class RequestState:
214
finish_reason: FinishReason | None,
215
stop_reason: int | str | None,
216
kv_transfer_params: dict[str, Any] | None = None,
217
+ routed_experts: np.ndarray | None = None,
218
) -> RequestOutput | PoolingRequestOutput | None:
219
finished = finish_reason is not None
220
final_only = self.output_kind == RequestOutputKind.FINAL_ONLY
@@ -253,7 +255,9 @@ class RequestState:
255
finished,
256
)
257
256
- output = self._new_completion_output(new_token_ids, finish_reason, stop_reason)
258
+ output = self._new_completion_output(
259
+ new_token_ids, finish_reason, stop_reason, routed_experts
260
+ )
261
262
if self.parent_req is None:
263
outputs = [output]
@@ -316,6 +320,7 @@ class RequestState:
320
token_ids: list[int],
321
finish_reason: FinishReason | None,
322
stop_reason: int | str | None,
323
+ routed_experts: np.ndarray | None = None,
324
) -> CompletionOutput:
325
assert self.detokenizer is not None
326
assert self.logprobs_processor is not None
@@ -336,6 +341,7 @@ class RequestState:
341
index=self.request_index,
342
text=text,
343
token_ids=token_ids,
344
+ routed_experts=routed_experts,
345
logprobs=logprobs,
346
cumulative_logprob=self.logprobs_processor.cumulative_logprob,
347
finish_reason=str(finish_reason) if finished else None,
@@ -527,6 +533,7 @@ class OutputProcessor:
533
finish_reason = engine_core_output.finish_reason
534
stop_reason = engine_core_output.stop_reason
535
kv_transfer_params = engine_core_output.kv_transfer_params
536
+ routed_experts = engine_core_output.routed_experts
537
req_state.num_cached_tokens = engine_core_output.num_cached_tokens
538
req_state.is_prefilling = False
539
@@ -552,6 +559,7 @@ class OutputProcessor:
559
finish_reason,
560
stop_reason,
561
kv_transfer_params,
562
+ routed_experts,
563
):
564
if req_state.queue is not None:
565
# AsyncLLM: put into queue for handling by generate().
vllm/v1/worker/gpu_model_runner.py
+41
@@ -50,6 +50,9 @@ from vllm.forward_context import (
50
from vllm.logger import init_logger
51
from vllm.lora.layers import LoRAMapping, LoRAMappingType
52
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
53
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
54
+ RoutedExpertsCapturer,
55
+)
56
from vllm.model_executor.layers.rotary_embedding import (
57
MRotaryEmbedding,
58
XDRotaryEmbedding,
@@ -1633,6 +1636,8 @@ class GPUModelRunner(
1636
return blk_table_tensor, slot_mapping
1637
1638
block_table_gid_0, slot_mapping_gid_0 = _get_block_table_and_slot_mapping(0)
1639
+ if self.model_config.enable_return_routed_experts:
1640
+ self.slot_mapping = slot_mapping_gid_0[:num_tokens].cpu().numpy()
1641
cm_base = CommonAttentionMetadata(
1642
query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1],
1643
query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1],
@@ -3112,6 +3117,13 @@ class GPUModelRunner(
3117
"after execute_model() returns None."
3118
)
3119
3120
+ if self.vllm_config.model_config.enable_return_routed_experts:
3121
+ capturer = RoutedExpertsCapturer.get_instance()
3122
+ if capturer is not None:
3123
+ capturer.clear_buffer() # noqa
3124
+ else:
3125
+ logger.error("RoutedExpertsCapturer not initialized.")
3126
+
3127
if scheduler_output.preempted_req_ids and has_kv_transfer_group():
3128
get_kv_transfer_group().handle_preemptions(
3129
scheduler_output.preempted_req_ids
@@ -3485,6 +3497,13 @@ class GPUModelRunner(
3497
self.eplb_step()
3498
3499
with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
3500
+ if self.model_config.enable_return_routed_experts:
3501
+ capturer = RoutedExpertsCapturer.get_instance()
3502
+ if capturer is not None:
3503
+ capturer.save_captured_experts(indices=self.slot_mapping) # noqa
3504
+ else:
3505
+ logger.error("RoutedExpertsCapturer not initialized.")
3506
+
3507
output = ModelRunnerOutput(
3508
req_ids=req_ids_output_copy,
3509
req_id_to_index=req_id_to_index_output_copy,
@@ -5646,6 +5665,28 @@ class GPUModelRunner(
5665
kv_transfer_group.register_kv_caches(kv_caches)
5666
kv_transfer_group.set_host_xfer_buffer_ops(copy_kv_blocks)
5667
5668
+ if self.model_config.enable_return_routed_experts:
5669
+ self.init_routed_experts_capturer()
5670
+
5671
+ def init_routed_experts_capturer(self):
5672
+ logger.info(
5673
+ "Initializing routed experts capturer, enable_return_routed_experts: %s",
5674
+ self.model_config.enable_return_routed_experts,
5675
+ )
5676
+ routed_experts_capturer = RoutedExpertsCapturer.create()
5677
+ block_size = self.cache_config.block_size
5678
+ self.max_num_kv_tokens = (
5679
+ self.kv_cache_config.num_blocks // len(self.kv_cache_config.kv_cache_groups)
5680
+ + 1
5681
+ ) * block_size
5682
+
5683
+ routed_experts_capturer.init_buffer(
5684
+ max_num_batched_tokens=self.scheduler_config.max_num_batched_tokens,
5685
+ max_num_kv_tokens=self.max_num_kv_tokens,
5686
+ model_config=self.model_config,
5687
+ instance_id=self.vllm_config.instance_id,
5688
+ )
5689
+
5690
def may_add_encoder_only_layers_to_kv_cache_config(self) -> None:
5691
"""
5692
Add encoder-only layers to the KV cache config.