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 9, 19:30 UTC
[Cleanup] Move scheduler `get_routed_experts` logic to separate method (#32706)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jan 22, 2026 at 07:46 UTC
bc14663e6a058722542789a044f70863edcb5a6a
1 file changed
+31
-23
vllm/v1/core/sched/scheduler.py
+31
-23
@@ -1197,28 +1197,7 @@ class Scheduler(SchedulerInterface):
1197
1198
routed_experts = None
1199
if stopped:
1200
- if self.vllm_config.model_config.enable_return_routed_experts:
1201
- kv_blocks = self.kv_cache_manager.get_blocks(request.request_id)
1202
- block_ids = kv_blocks.get_block_ids()[0]
1203
- num_tokens = request.num_tokens - 1
1204
-
1205
- # compute slot mapping
1206
- block_ids_array = np.array(block_ids, dtype=np.int32)
1207
- num_blocks = len(block_ids)
1208
- block_size = self.block_size
1209
-
1210
- # generate block offsets
1211
- block_offsets = np.arange(0, block_size)
1212
-
1213
- # compute slot mapping: slot = block_id * block_size + offset
1214
- slot_mapping = (
1215
- block_offsets.reshape((1, block_size))
1216
- + block_ids_array.reshape((num_blocks, 1)) * block_size
1217
- ).flatten()[:num_tokens]
1218
-
1219
- routed_experts = self.routed_experts_reader.get_routed_experts(
1220
- indices=slot_mapping
1221
- )
1200
+ routed_experts = self._get_routed_experts(request)
1201
kv_transfer_params = self._free_request(request)
1202
if status_before_stop == RequestStatus.RUNNING:
1203
stopped_running_reqs.add(request)
@@ -1250,7 +1229,12 @@ class Scheduler(SchedulerInterface):
1229
1230
# Get prompt logprobs for this request.
1231
prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
1253
- if new_token_ids or pooler_output is not None or kv_transfer_params:
1232
+ if (
1233
+ new_token_ids
1234
+ or pooler_output is not None
1235
+ or kv_transfer_params
1236
+ or stopped
1237
+ ):
1238
# Add EngineCoreOutput for this Request.
1239
outputs[request.client_index].append(
1240
EngineCoreOutput(
@@ -1351,6 +1335,30 @@ class Scheduler(SchedulerInterface):
1335
1336
return engine_core_outputs
1337
1338
+ def _get_routed_experts(self, request: Request) -> np.ndarray | None:
1339
+ if not self.vllm_config.model_config.enable_return_routed_experts:
1340
+ return None
1341
+
1342
+ kv_blocks = self.kv_cache_manager.get_blocks(request.request_id)
1343
+ block_ids = kv_blocks.get_block_ids()[0]
1344
+ num_tokens = request.num_tokens - 1
1345
+
1346
+ # compute slot mapping
1347
+ block_ids_array = np.array(block_ids, dtype=np.int32)
1348
+ num_blocks = len(block_ids)
1349
+ block_size = self.block_size
1350
+
1351
+ # generate block offsets
1352
+ block_offsets = np.arange(0, block_size)
1353
+
1354
+ # compute slot mapping: slot = block_id * block_size + offset
1355
+ slot_mapping = (
1356
+ block_offsets.reshape((1, block_size))
1357
+ + block_ids_array.reshape((num_blocks, 1)) * block_size
1358
+ ).flatten()[:num_tokens]
1359
+
1360
+ return self.routed_experts_reader.get_routed_experts(indices=slot_mapping)
1361
+
1362
def _update_request_with_output(
1363
self, request: Request, new_token_ids: list[int]
1364
) -> tuple[list[int], bool]: