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 13, 20:45 UTC
[AsyncScheduling] Don't schedule past request max_tokens (#27922)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Nov 4, 2025 at 09:06 UTC
938a81692ea318e59ead4750e7e7425bfd6a4896
3 files changed
+14
-4
tests/v1/core/test_async_scheduler.py
+7
@@ -34,15 +34,20 @@ def test_stop_by_max_tokens(max_tokens: int):
34
requests = create_requests(num_requests=2, max_tokens=max_tokens)
35
req0, req1 = requests
36
37
+ expected_total_num_scheduled_tokens = 0
38
sched_outputs: deque[SchedulerOutput] = deque()
39
scheduler.add_request(req0)
40
sched_outputs.append(scheduler.schedule())
41
+ expected_total_num_scheduled_tokens += req0.num_prompt_tokens + max_tokens - 1
42
43
scheduler.add_request(req1)
44
sched_outputs.append(scheduler.schedule())
45
+ expected_total_num_scheduled_tokens += req1.num_prompt_tokens + max_tokens - 1
46
47
+ total_num_scheduled_tokens = 0
48
while sched_outputs:
49
sched_output = sched_outputs.popleft()
50
+ total_num_scheduled_tokens += sched_output.total_num_scheduled_tokens
51
model_runner_output = _make_model_runner_output(sched_output)
52
scheduler.update_from_output(sched_output, model_runner_output)
53
@@ -53,6 +58,8 @@ def test_stop_by_max_tokens(max_tokens: int):
58
assert scheduler.get_num_unfinished_requests() == 0
59
assert req0.num_output_tokens == max_tokens
60
assert req1.num_output_tokens == max_tokens
61
+ # Ensure we aren't scheduling more tokens than necessary.
62
+ assert total_num_scheduled_tokens == expected_total_num_scheduled_tokens
63
64
65
def test_abort():
tests/v1/e2e/test_spec_decode.py
-1
@@ -155,7 +155,6 @@ def test_suffix_decoding_acceptance(
155
)
156
157
# Run several times and check that the accepted tokens increase.
158
- spec_llm.chat(test_prompts, sampling_config)
158
num_draft = []
159
num_accept = []
160
for i in range(10): # Run multiple times to warm up the cache.
vllm/v1/core/sched/scheduler.py
+7
-3
@@ -217,10 +217,14 @@ class Scheduler(SchedulerInterface):
217
num_new_tokens = self.scheduler_config.long_prefill_token_threshold
218
num_new_tokens = min(num_new_tokens, token_budget)
219
220
- # Make sure the input position does not exceed the max model len.
221
- # This is necessary when using spec decoding.
220
+ # Make sure the input position does not exceed the max model len or
221
+ # request's max_tokens.
222
+ # This is necessary when using spec decoding and/or async scheduling.
223
+ max_total_tokens = min(
224
+ request.num_prompt_tokens + request.max_tokens, self.max_model_len
225
+ )
226
num_new_tokens = min(
223
- num_new_tokens, self.max_model_len - 1 - request.num_computed_tokens
227
+ num_new_tokens, max_total_tokens - 1 - request.num_computed_tokens
228
)
229
230
# Schedule encoder inputs.