[Perf] Skip/shrink all_token_ids copy in scheduler for non-async and V2 runner (#45840)
Signed-off-by: amanchugh89 <amanchugh.89@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
aman committed
Jun 20, 2026 at 23:36 UTC
6e919960af42f79d6811d84b2d4316212fcf59cb
3 files changed
+49
-11
tests/v1/core/test_scheduler.py
+37
@@ -144,6 +144,43 @@ def test_async_scheduling_pp_allows_rescheduling_with_output_placeholders():
144
assert req.request_id in output.num_scheduled_tokens
145
146
147
+def test_cached_request_data_resumed_all_token_ids_mrv1_only():
148
+ """all_token_ids carries a resumed request's token ids to the connector
149
+ for the V1 model runner, but is skipped entirely for the V2 model runner.
150
+ """
151
+ from vllm.v1.core.kv_cache_manager import KVCacheBlocks
152
+
153
+ scheduler = create_scheduler()
154
+ (req,) = create_requests(num_requests=1, num_tokens=8)
155
+ req.append_output_token_ids([101, 102, 103])
156
+
157
+ # A resumed request was not scheduled in the previous step.
158
+ assert req.request_id not in scheduler.prev_step_scheduled_req_ids
159
+
160
+ empty_blocks = KVCacheBlocks(blocks=((),))
161
+
162
+ def make_cached():
163
+ return scheduler._make_cached_request_data(
164
+ running_reqs=[],
165
+ resumed_reqs=[req],
166
+ num_scheduled_tokens={req.request_id: 1},
167
+ spec_decode_tokens={},
168
+ req_to_new_blocks={req.request_id: empty_blocks},
169
+ )
170
+
171
+ # V1 model runner: the full token id list is propagated.
172
+ assert not scheduler.use_v2_model_runner
173
+ cached = make_cached()
174
+ assert req.request_id in cached.resumed_req_ids
175
+ assert cached.all_token_ids[req.request_id] == list(req.all_token_ids)
176
+
177
+ # V2 model runner: all_token_ids is skipped entirely.
178
+ scheduler.use_v2_model_runner = True
179
+ cached = make_cached()
180
+ assert req.request_id in cached.resumed_req_ids
181
+ assert cached.all_token_ids == {}
182
+
183
+
184
def test_schedule_partial_requests():
185
"""Test scheduling behavior with partial requests.
186
vllm/v1/core/sched/output.py
+2
-2
@@ -118,8 +118,8 @@ class CachedRequestData:
118
# NOTE(woosuk): new_token_ids is only used for pipeline parallelism.
119
# When PP is not used, new_token_ids will be empty.
120
new_token_ids: list[list[int]]
121
- # For requests not scheduled in the last step, propagate the token ids to the
122
- # connector. Won't contain requests that were scheduled in the prior step.
121
+ # MRV1-only: For requests not scheduled in the last step, propagate the token ids
122
+ # to the connector. Won't contain requests scheduled in the prior step.
123
all_token_ids: dict[str, list[int]]
124
new_block_ids: list[tuple[list[int], ...] | None]
125
num_computed_tokens: list[int]
vllm/v1/core/sched/scheduler.py
+10
-9
@@ -101,6 +101,7 @@ class Scheduler(SchedulerInterface):
101
self.finished_req_ids_dict: dict[int, set[str]] | None = (
102
defaultdict(set) if include_finished_set else None
103
)
104
+ # Track requests scheduled in prior step (MRV1-only).
105
self.prev_step_scheduled_req_ids: set[str] = set()
106
107
# Scheduling constraints.
@@ -1010,8 +1011,8 @@ class Scheduler(SchedulerInterface):
1011
1012
# Construct the scheduler output.
1013
if self.use_v2_model_runner:
1013
- scheduled_new_reqs = scheduled_new_reqs + scheduled_resumed_reqs
1014
- scheduled_resumed_reqs = []
1014
+ scheduled_new_reqs.extend(scheduled_resumed_reqs)
1015
+ scheduled_resumed_reqs.clear()
1016
new_reqs_data = [
1017
NewRequestData.from_request(
1018
req,
@@ -1037,9 +1038,10 @@ class Scheduler(SchedulerInterface):
1038
req_to_new_blocks,
1039
)
1040
1040
- # Record the request ids that were scheduled in this step.
1041
- self.prev_step_scheduled_req_ids.clear()
1042
- self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())
1041
+ # Record the request ids that were scheduled in this step (MRV1-only).
1042
+ if not self.use_v2_model_runner:
1043
+ self.prev_step_scheduled_req_ids.clear()
1044
+ self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())
1045
1046
new_block_ids_to_zero = (
1047
(self.kv_cache_manager.take_new_block_ids() or None)
@@ -1252,12 +1254,11 @@ class Scheduler(SchedulerInterface):
1254
req.num_computed_tokens : req.num_computed_tokens + num_tokens
1255
]
1256
new_token_ids.append(token_ids)
1255
- scheduled_in_prev_step = req_id in self.prev_step_scheduled_req_ids
1257
if idx >= num_running_reqs:
1257
- assert not scheduled_in_prev_step
1258
resumed_req_ids.add(req_id)
1259
- if not scheduled_in_prev_step:
1260
- all_token_ids[req_id] = req.all_token_ids.copy()
1259
+ if not self.use_v2_model_runner: # noqa: SIM102
1260
+ if req_id not in self.prev_step_scheduled_req_ids:
1261
+ all_token_ids[req_id] = req.all_token_ids.copy()
1262
new_block_ids.append(
1263
req_to_new_blocks[req_id].get_block_ids(allow_none=True)
1264
)