[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>

tests/v1/streaming_input/test_scheduler_streaming.py

@@ -171,8 +171,7 @@ class TestStreamingScheduler(unittest.TestCase):
             prompt_token_ids=[1, 2, 3],
         )
         session.num_computed_tokens = len(session.prompt_token_ids)
-        session.max_tokens = 10  # Initial max_tokens
-        session._output_token_ids = [1] * 10  # reach max_tokens
+        session.max_tokens = 2
 
         new_request = DummyRequest(
             request_id="session",
@@ -185,26 +184,21 @@ class TestStreamingScheduler(unittest.TestCase):
         scheduler._update_request_as_session(session, update)
 
         assert session.sampling_params.max_tokens == 10
-        # _update_request_as_session clears output tokens first, so
-        # max_tokens = num_output_tokens (0) + update.max_tokens (10) = 10
         assert session.max_tokens == 10
 
         session.num_computed_tokens = len(session.prompt_token_ids)
 
-        # Simulate generating 5 more output tokens
-        session._output_token_ids = [1] * 5
         new_request2 = DummyRequest(
             request_id="session",
             prompt_token_ids=[7, 8, 9],
         )
-        new_request2.sampling_params = SamplingParams(max_tokens=10)
-        new_request2.max_tokens = 10
+        new_request2.sampling_params = SamplingParams(max_tokens=4)
+        new_request2.max_tokens = 4
         update2 = StreamingUpdate.from_request(new_request2)
         scheduler._update_request_as_session(session, update2)
 
-        assert session.sampling_params.max_tokens == 10
-        # Again, output tokens are cleared first, so max_tokens = 0 + 10 = 10
-        assert session.max_tokens == 10
+        assert session.sampling_params.max_tokens == 4
+        assert session.max_tokens == 4
 
     def test_update_request_as_session(self):
         scheduler = create_scheduler()

vllm/v1/core/sched/scheduler.py

@@ -1648,6 +1648,7 @@ class Scheduler(SchedulerInterface):
         session.num_prompt_tokens = len(session.prompt_token_ids)
         session.arrival_time = update.arrival_time
         session.sampling_params = update.sampling_params
+        session.max_tokens = update.max_tokens
         if session.status == RequestStatus.WAITING_FOR_STREAMING_REQ:
             self.num_waiting_for_streaming_input -= 1
         session.status = RequestStatus.WAITING