Ensure token_budget respects max_num_scheduled_tokens constraint

Author: juncheoll

vllm/v1/core/sched/scheduler.py

@@ -1437,10 +1437,7 @@ def schedule(self) -> SchedulerOutput:
         token_budget = self._estimate_token_budget()
         req_index = 0
         while self.scheduled_queue:
-            if token_budget <= 0:
-                break
-
-            scheduled_request = self.scheduled_queue.popleft()
+            scheduled_request = self.scheduled_queue[0]
 
             req_id = scheduled_request.request_id
             num_new_tokens = scheduled_request.num_new_tokens
@@ -1451,6 +1448,11 @@ def schedule(self) -> SchedulerOutput:
             spec_token_ids = scheduled_request.spec_token_ids
             request_data = scheduled_request.request_data
 
+            if num_scheduled_tokens and token_budget < num_new_tokens:
+                break
+
+            self.scheduled_queue.popleft()
+
             # requests in the scheduled_queue can also be preempted or finished.
             if self.requests[req_id].status >= RequestStatus.PREEMPTED:
                 continue