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 11, 12:30 UTC
[Scheduer] Simplify stop checking for pooling models (#30591)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Dec 13, 2025 at 01:45 UTC
1cec5b7ea9ba72b34de9a7c7001beb8a1b8f0dc0
2 files changed
+7
-16
vllm/v1/core/sched/scheduler.py
+5
-6
@@ -1117,6 +1117,7 @@ class Scheduler(SchedulerInterface):
1117
stopped = False
1118
new_logprobs = None
1119
new_token_ids = generated_token_ids
1120
+ pooler_output = pooler_outputs[req_index] if pooler_outputs else None
1121
kv_transfer_params = None
1122
status_before_stop = request.status
1123
@@ -1125,12 +1126,10 @@ class Scheduler(SchedulerInterface):
1126
new_token_ids, stopped = self._update_request_with_output(
1127
request, new_token_ids
1128
)
1128
-
1129
- # Stop checking for pooler models.
1130
- pooler_output = None
1131
- if pooler_outputs:
1132
- pooler_output = pooler_outputs[req_index]
1133
- stopped = check_stop(request, self.max_model_len, pooler_output)
1129
+ elif request.pooling_params and pooler_output is not None:
1130
+ # Pooling stops as soon as there is output.
1131
+ request.status = RequestStatus.FINISHED_STOPPED
1132
+ stopped = True
1133
1134
if stopped:
1135
kv_transfer_params = self._free_request(request)
vllm/v1/core/sched/utils.py
+2
-10
@@ -2,8 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import contextlib
4
5
-import torch
6
-
5
from vllm.v1.request import Request, RequestStatus
6
7
@@ -39,14 +37,8 @@ def remove_all(lst: list, items_to_remove: set) -> list:
37
return [item for item in lst if item not in items_to_remove]
38
39
42
-def check_stop(
43
- request: Request, max_model_len: int, pooler_output: torch.Tensor | None = None
44
-) -> bool:
45
- if request.pooling_params:
46
- if pooler_output is not None:
47
- request.status = RequestStatus.FINISHED_STOPPED
48
- return True
49
- return False
40
+def check_stop(request: Request, max_model_len: int) -> bool:
41
+ assert not request.pooling_params
42
43
sampling_params = request.sampling_params
44
assert sampling_params is not None