Merge branch 'main' into wentao-optimize-async-scheduling-copy

Signed-off-by: yewentao256 <zhyanwentao@126.com>

tests/v1/core/test_scheduler.py

@@ -2568,13 +2568,19 @@ def test_priority_scheduling_preemption_and_resumption_when_out_of_kv(
 
     # Preempted request resumed in scheduled_cached_reqs
     assert len(scheduled_cached_reqs.resumed_req_ids) == 1
-    assert len(scheduled_cached_reqs.all_token_ids) == 1
     assert scheduled_cached_reqs.req_ids[0] == request_low.request_id
     assert request_low.request_id in scheduled_cached_reqs.resumed_req_ids
-    assert request_low.request_id in scheduled_cached_reqs.all_token_ids
-    # Resumed tokens include 30 prompt tokens and 2 decoded tokens
-    assert len(scheduled_cached_reqs.all_token_ids[request_low.request_id]) == 32
-    assert scheduled_cached_reqs.all_token_ids[request_low.request_id][31] == 100
+    # `output_token_ids` is only propagated for the v1 model runner.
+    if not scheduler.use_v2_model_runner:
+        assert len(scheduled_cached_reqs.output_token_ids) == 1
+        assert request_low.request_id in scheduled_cached_reqs.output_token_ids
+        # request_low produced two decode tokens before being preempted.
+        assert scheduled_cached_reqs.output_token_ids[request_low.request_id] == [
+            100,
+            100,
+        ]
+    else:
+        assert len(scheduled_cached_reqs.output_token_ids) == 0
 
 
 @pytest.mark.parametrize(
@@ -3422,13 +3428,19 @@ def test_priority_scheduling_ec_connector_preemption_and_resumption(
 
     # Preempted request resumed in scheduled_cached_reqs
     assert len(scheduled_cached_reqs.resumed_req_ids) == 1
-    assert len(scheduled_cached_reqs.all_token_ids) == 1
     assert scheduled_cached_reqs.req_ids[0] == request_low.request_id
     assert request_low.request_id in scheduled_cached_reqs.resumed_req_ids
-    assert request_low.request_id in scheduled_cached_reqs.all_token_ids
-    ## Resumed tokens include 94 prompt tokens and 2 decoded tokens
-    assert len(scheduled_cached_reqs.all_token_ids[request_low.request_id]) == 96
-    assert scheduled_cached_reqs.all_token_ids[request_low.request_id][95] == 100
+    # `output_token_ids` is only propagated for the v1 model runner.
+    if not scheduler.use_v2_model_runner:
+        assert len(scheduled_cached_reqs.output_token_ids) == 1
+        assert request_low.request_id in scheduled_cached_reqs.output_token_ids
+        # request_low produced two decode tokens before being preempted.
+        assert scheduled_cached_reqs.output_token_ids[request_low.request_id] == [
+            100,
+            100,
+        ]
+    else:
+        assert len(scheduled_cached_reqs.output_token_ids) == 0
     assert scheduler.running[0].request_id == request_low.request_id
     assert request_high.request_id in output.finished_req_ids
 

tests/v1/worker/test_gpu_model_runner.py

@@ -300,7 +300,7 @@ def test_update_states_request_resumed(model_runner, dist_init):
         req_ids=[req_id],
         resumed_req_ids=set(),
         new_token_ids=[[]],
-        all_token_ids={},
+        output_token_ids={},
         new_block_ids=[([0],)],
         num_computed_tokens=[0],
         num_output_tokens=[0],

vllm/v1/core/sched/output.py

@@ -118,7 +118,7 @@ class CachedRequestData:
     new_token_ids: list[list[int]]
     # For requests not scheduled in the last step, propagate the token ids to the
     # connector. Won't contain requests that were scheduled in the prior step.
-    all_token_ids: dict[str, list[int]]
+    output_token_ids: dict[str, list[int]]
     new_block_ids: list[tuple[list[int], ...] | None]
     num_computed_tokens: list[int]
     num_output_tokens: list[int]
@@ -126,15 +126,15 @@ class CachedRequestData:
     # Version of dataclass repr with token IDs obfuscated.
     def anon_repr(self) -> str:
         new_token_ids_lens = [len(toks) for toks in self.new_token_ids]
-        all_token_ids_lens = {
-            req_id: len(toks) for req_id, toks in self.all_token_ids.items()
+        output_token_ids_lens = {
+            req_id: len(toks) for req_id, toks in self.output_token_ids.items()
         }
         return (
             f"CachedRequestData("
             f"req_ids={self.req_ids},"
             f"resumed_req_ids={self.resumed_req_ids},"
             f"new_token_ids_lens={new_token_ids_lens},"
-            f"all_token_ids_lens={all_token_ids_lens},"
+            f"output_token_ids_lens={output_token_ids_lens},"
             f"new_block_ids={self.new_block_ids},"
             f"num_computed_tokens={self.num_computed_tokens},"
             f"num_output_tokens={self.num_output_tokens}"
@@ -168,7 +168,7 @@ class CachedRequestData:
             req_ids=[],
             resumed_req_ids=set(),
             new_token_ids=[],
-            all_token_ids={},
+            output_token_ids={},
             new_block_ids=[],
             num_computed_tokens=[],
             num_output_tokens=[],

vllm/v1/core/sched/scheduler.py

@@ -1056,7 +1056,7 @@ class Scheduler(SchedulerInterface):
         req_ids: list[str] = []
         new_token_ids: list[list[int]] = []
         new_block_ids: list[tuple[list[int], ...] | None] = []
-        all_token_ids: dict[str, list[int]] = {}
+        output_token_ids: dict[str, list[int]] = {}
         num_computed_tokens: list[int] = []
         num_output_tokens: list[int] = []
         resumed_req_ids = set()
@@ -1085,8 +1085,15 @@ class Scheduler(SchedulerInterface):
             if idx >= num_running_reqs:
                 assert not scheduled_in_prev_step
                 resumed_req_ids.add(req_id)
-            if not scheduled_in_prev_step:
-                all_token_ids[req_id] = req.all_token_ids.copy()
+            if not scheduled_in_prev_step and self.scheduler_config.async_scheduling:
+                assert req.num_output_placeholders == 0, (
+                    "Unexpected output placeholders for request not scheduled "
+                    "in the previous step."
+                )
+            if not self.use_v2_model_runner and not scheduled_in_prev_step:
+                num_out = req.num_output_tokens + req.num_output_placeholders
+                if num_out > 0:
+                    output_token_ids[req_id] = req.output_token_ids.copy()
             new_block_ids.append(
                 req_to_new_blocks[req_id].get_block_ids(allow_none=True)
             )
@@ -1099,7 +1106,7 @@ class Scheduler(SchedulerInterface):
             req_ids=req_ids,
             resumed_req_ids=resumed_req_ids,
             new_token_ids=new_token_ids,
-            all_token_ids=all_token_ids,
+            output_token_ids=output_token_ids,
             new_block_ids=new_block_ids,
             num_computed_tokens=num_computed_tokens,
             num_output_tokens=num_output_tokens,

vllm/v1/worker/gpu_model_runner.py

@@ -1314,8 +1314,10 @@ class GPUModelRunner(
                 if self.use_async_scheduling and num_output_tokens > 0:
                     # We must recover the output token ids for resumed requests in the
                     # async scheduling case, so that correct input_ids are obtained.
-                    resumed_token_ids = req_data.all_token_ids[req_id]
-                    req_state.output_token_ids = resumed_token_ids[-num_output_tokens:]
+                    resumed_output_token_ids = req_data.output_token_ids[req_id]
+                    req_state.output_token_ids = resumed_output_token_ids[
+                        -num_output_tokens:
+                    ]
 
                 reqs_to_add.append(req_state)
                 # Track resumed requests for ngram_gpu full tensor copy