[Bugfix][Scheduler] Refresh max tokens for streaming continuations (#57676)
Signed-off-by: 0xsensei <prblmslvr.aditya@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
0xsensei committed
Sep 29, 2026 at 02:01 UTC
d28795f1a7af4e3ce2530d4f0bdaec4ecbede693
2 files changed
+6
-11
tests/v1/streaming_input/test_scheduler_streaming.py
+5
-11
@@ -171,8 +171,7 @@ class TestStreamingScheduler(unittest.TestCase):
171
prompt_token_ids=[1, 2, 3],
172
)
173
session.num_computed_tokens = len(session.prompt_token_ids)
174
- session.max_tokens = 10 # Initial max_tokens
175
- session._output_token_ids = [1] * 10 # reach max_tokens
174
+ session.max_tokens = 2
175
176
new_request = DummyRequest(
177
request_id="session",
@@ -185,26 +184,21 @@ class TestStreamingScheduler(unittest.TestCase):
184
scheduler._update_request_as_session(session, update)
185
186
assert session.sampling_params.max_tokens == 10
188
- # _update_request_as_session clears output tokens first, so
189
- # max_tokens = num_output_tokens (0) + update.max_tokens (10) = 10
187
assert session.max_tokens == 10
188
189
session.num_computed_tokens = len(session.prompt_token_ids)
190
194
- # Simulate generating 5 more output tokens
195
- session._output_token_ids = [1] * 5
191
new_request2 = DummyRequest(
192
request_id="session",
193
prompt_token_ids=[7, 8, 9],
194
)
200
- new_request2.sampling_params = SamplingParams(max_tokens=10)
201
- new_request2.max_tokens = 10
195
+ new_request2.sampling_params = SamplingParams(max_tokens=4)
196
+ new_request2.max_tokens = 4
197
update2 = StreamingUpdate.from_request(new_request2)
198
scheduler._update_request_as_session(session, update2)
199
205
- assert session.sampling_params.max_tokens == 10
206
- # Again, output tokens are cleared first, so max_tokens = 0 + 10 = 10
207
- assert session.max_tokens == 10
200
+ assert session.sampling_params.max_tokens == 4
201
+ assert session.max_tokens == 4
202
203
def test_update_request_as_session(self):
204
scheduler = create_scheduler()
vllm/v1/core/sched/scheduler.py
+1
@@ -1648,6 +1648,7 @@ class Scheduler(SchedulerInterface):
1648
session.num_prompt_tokens = len(session.prompt_token_ids)
1649
session.arrival_time = update.arrival_time
1650
session.sampling_params = update.sampling_params
1651
+ session.max_tokens = update.max_tokens
1652
if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
1653
self.num_waiting_for_streaming_input -= 1
1654
session.status = RequestStatus.WAITING