A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 03:30 UTC
[Optimization] Use Shared `CachedRequestData` Instance Across All Requests (#20232)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jun 30, 2025 at 09:07 UTC
2863befce359ee1a82afe02d1953252866aa3e96
12 files changed
+221
-232
tests/v1/core/test_scheduler.py
+67
-63
@@ -10,7 +10,7 @@ from vllm.config import (CacheConfig, KVTransferConfig, ModelConfig,
10
SchedulerConfig, SpeculativeConfig, VllmConfig)
11
from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
12
from vllm.sampling_params import SamplingParams
13
-from vllm.v1.core.sched.output import SchedulerOutput
13
+from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
14
from vllm.v1.core.sched.scheduler import Scheduler
15
from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
16
KVCacheGroupSpec)
@@ -198,7 +198,7 @@ def test_schedule(enable_prefix_caching: Optional[bool],
198
# Test initial scheduling
199
output = scheduler.schedule()
200
assert len(output.scheduled_new_reqs) == len(requests)
201
- assert len(output.scheduled_cached_reqs) == 0
201
+ assert output.scheduled_cached_reqs.num_reqs == 0
202
assert len(output.finished_req_ids) == 0
203
# Verify all requests are scheduled.
204
for req_id, num_tokens in output.num_scheduled_tokens.items():
@@ -225,7 +225,7 @@ def test_schedule_multimodal_requests():
225
226
output = scheduler.schedule()
227
assert len(output.scheduled_new_reqs) == len(requests)
228
- assert len(output.scheduled_cached_reqs) == 0
228
+ assert output.scheduled_cached_reqs.num_reqs == 0
229
assert len(output.finished_req_ids) == 0
230
for req_id, num_tokens in output.num_scheduled_tokens.items():
231
assert num_tokens == len(requests[int(req_id)].prompt_token_ids)
@@ -259,7 +259,7 @@ def test_schedule_partial_requests():
259
260
output = scheduler.schedule()
261
assert len(output.scheduled_new_reqs) == 3
262
- assert len(output.scheduled_cached_reqs) == 0
262
+ assert output.scheduled_cached_reqs.num_reqs == 0
263
assert len(output.finished_req_ids) == 0
264
265
assert scheduler.max_num_encoder_input_tokens == 1024
@@ -295,7 +295,7 @@ def test_schedule_partial_requests():
295
output = scheduler.schedule()
296
assert len(scheduler.running) == 3
297
assert len(output.scheduled_new_reqs) == 0
298
- assert len(output.scheduled_cached_reqs) == 2
298
+ assert output.scheduled_cached_reqs.num_reqs == 2
299
assert len(output.finished_req_ids) == 0
300
assert output.num_scheduled_tokens[requests[0].request_id] == 1
301
assert output.num_scheduled_tokens[requests[1].request_id] == 700
@@ -319,7 +319,7 @@ def test_no_mm_input_chunking():
319
320
output = scheduler.schedule()
321
assert len(output.scheduled_new_reqs) == 1
322
- assert len(output.scheduled_cached_reqs) == 0
322
+ assert output.scheduled_cached_reqs.num_reqs == 0
323
assert len(output.finished_req_ids) == 0
324
# We want to only see the 400 text tokens at the start scheduled
325
assert output.num_scheduled_tokens[requests[0].request_id] == 400
@@ -342,7 +342,7 @@ def test_no_mm_input_chunking():
342
output = scheduler.schedule()
343
assert len(scheduler.running) == 1
344
assert len(output.scheduled_new_reqs) == 0
345
- assert len(output.scheduled_cached_reqs) == 1
345
+ assert output.scheduled_cached_reqs.num_reqs == 1
346
assert len(output.finished_req_ids) == 0
347
assert output.num_scheduled_tokens[requests[0].request_id] == 800
348
@@ -379,7 +379,7 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):
379
380
output = scheduler.schedule()
381
assert len(output.scheduled_new_reqs) == 3
382
- assert len(output.scheduled_cached_reqs) == 0
382
+ assert output.scheduled_cached_reqs.num_reqs == 0
383
assert len(output.finished_req_ids) == 0
384
385
# The first request is scheduled partially - 400.
@@ -408,7 +408,7 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):
408
output1 = scheduler.schedule()
409
assert len(scheduler.running) == 3
410
assert len(output1.scheduled_new_reqs) == 0
411
- assert len(output1.scheduled_cached_reqs) == 3
411
+ assert output1.scheduled_cached_reqs.num_reqs == 3
412
assert len(output1.finished_req_ids) == 0
413
assert output1.num_scheduled_tokens[requests[0].request_id] == 400
414
assert output1.num_scheduled_tokens[requests[1].request_id] == 400
@@ -430,7 +430,7 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):
430
output2 = scheduler.schedule()
431
assert len(scheduler.running) == 3
432
assert len(output2.scheduled_new_reqs) == 0
433
- assert len(output2.scheduled_cached_reqs) == 3
433
+ assert output2.scheduled_cached_reqs.num_reqs == 3
434
assert len(output2.finished_req_ids) == 0
435
assert output2.num_scheduled_tokens[requests[0].request_id] == 1
436
assert output2.num_scheduled_tokens[requests[1].request_id] == 1
@@ -449,23 +449,24 @@ def test_stop_via_update_from_output():
449
scheduler.requests[req.request_id] = req
450
scheduler.running.append(req)
451
452
- scheduler_output = SchedulerOutput(scheduled_new_reqs=[],
453
- scheduled_cached_reqs=[],
454
- num_scheduled_tokens={
455
- requests[0].request_id: 1,
456
- requests[1].request_id: 2
457
- },
458
- total_num_scheduled_tokens=3,
459
- scheduled_encoder_inputs={},
460
- scheduled_spec_decode_tokens={
461
- requests[0].request_id: [],
462
- requests[1].request_id: [10]
463
- },
464
- num_common_prefix_blocks=0,
465
- finished_req_ids=set(),
466
- free_encoder_input_ids=[],
467
- structured_output_request_ids={},
468
- grammar_bitmask=None)
452
+ scheduler_output = SchedulerOutput(
453
+ scheduled_new_reqs=[],
454
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
455
+ num_scheduled_tokens={
456
+ requests[0].request_id: 1,
457
+ requests[1].request_id: 2
458
+ },
459
+ total_num_scheduled_tokens=3,
460
+ scheduled_encoder_inputs={},
461
+ scheduled_spec_decode_tokens={
462
+ requests[0].request_id: [],
463
+ requests[1].request_id: [10]
464
+ },
465
+ num_common_prefix_blocks=0,
466
+ finished_req_ids=set(),
467
+ free_encoder_input_ids=[],
468
+ structured_output_request_ids={},
469
+ grammar_bitmask=None)
470
471
model_output = ModelRunnerOutput(
472
req_ids=[req.request_id for req in requests],
@@ -501,23 +502,25 @@ def test_stop_via_update_from_output():
502
scheduler.requests[req.request_id] = req
503
scheduler.running.append(req)
504
504
- scheduler_output = SchedulerOutput(scheduled_new_reqs=[],
505
- scheduled_cached_reqs=[],
506
- num_scheduled_tokens={
507
- requests[0].request_id: 3,
508
- requests[1].request_id: 2
509
- },
510
- total_num_scheduled_tokens=5,
511
- scheduled_encoder_inputs={},
512
- scheduled_spec_decode_tokens={
513
- requests[0].request_id: [10, 42],
514
- requests[1].request_id: [13]
515
- },
516
- num_common_prefix_blocks=0,
517
- finished_req_ids=set(),
518
- free_encoder_input_ids=[],
519
- structured_output_request_ids={},
520
- grammar_bitmask=None)
505
+ scheduler_output = SchedulerOutput(
506
+ scheduled_new_reqs=[],
507
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
508
+ num_scheduled_tokens={
509
+ requests[0].request_id: 3,
510
+ requests[1].request_id: 2
511
+ },
512
+ total_num_scheduled_tokens=5,
513
+ scheduled_encoder_inputs={},
514
+ scheduled_spec_decode_tokens={
515
+ requests[0].request_id: [10, 42],
516
+ requests[1].request_id: [13]
517
+ },
518
+ num_common_prefix_blocks=0,
519
+ finished_req_ids=set(),
520
+ free_encoder_input_ids=[],
521
+ structured_output_request_ids={},
522
+ grammar_bitmask=None,
523
+ )
524
525
model_output = ModelRunnerOutput(
526
req_ids=[req.request_id for req in requests],
@@ -551,23 +554,25 @@ def test_stop_via_update_from_output():
554
scheduler.requests[req.request_id] = req
555
scheduler.running.append(req)
556
554
- scheduler_output = SchedulerOutput(scheduled_new_reqs=[],
555
- scheduled_cached_reqs=[],
556
- num_scheduled_tokens={
557
- requests[0].request_id: 3,
558
- requests[1].request_id: 1
559
- },
560
- total_num_scheduled_tokens=4,
561
- scheduled_encoder_inputs={},
562
- scheduled_spec_decode_tokens={
563
- requests[0].request_id: [10, 11],
564
- requests[1].request_id: []
565
- },
566
- num_common_prefix_blocks=0,
567
- finished_req_ids=set(),
568
- free_encoder_input_ids=[],
569
- structured_output_request_ids={},
570
- grammar_bitmask=None)
557
+ scheduler_output = SchedulerOutput(
558
+ scheduled_new_reqs=[],
559
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
560
+ num_scheduled_tokens={
561
+ requests[0].request_id: 3,
562
+ requests[1].request_id: 1
563
+ },
564
+ total_num_scheduled_tokens=4,
565
+ scheduled_encoder_inputs={},
566
+ scheduled_spec_decode_tokens={
567
+ requests[0].request_id: [10, 11],
568
+ requests[1].request_id: []
569
+ },
570
+ num_common_prefix_blocks=0,
571
+ finished_req_ids=set(),
572
+ free_encoder_input_ids=[],
573
+ structured_output_request_ids={},
574
+ grammar_bitmask=None,
575
+ )
576
577
model_output = ModelRunnerOutput(
578
req_ids=[req.request_id for req in requests],
@@ -603,7 +608,7 @@ def test_stop_via_update_from_output():
608
609
scheduler_output = SchedulerOutput(
610
scheduled_new_reqs=[],
606
- scheduled_cached_reqs=[],
611
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
612
num_scheduled_tokens={requests[0].request_id: 3},
613
total_num_scheduled_tokens=3,
614
scheduled_encoder_inputs={},
@@ -1208,7 +1213,6 @@ def assert_scheduler_empty(scheduler: Scheduler):
1213
assert len(scheduler.waiting) == 0
1214
assert len(scheduler.running) == 0
1215
assert len(scheduler.finished_req_ids) == 0
1211
- assert len(scheduler._cached_reqs_data) == 0
1216
1217
# EncoderCacheManager.
1218
assert len(scheduler.encoder_cache_manager.freed) == 0
tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py
+2
-2
@@ -66,7 +66,7 @@ def test_basic_lifecycle():
66
assert len(scheduler_output.finished_req_ids) == 1
67
assert request_id in scheduler_output.finished_req_ids
68
assert len(scheduler_output.scheduled_new_reqs) == 0
69
- assert len(scheduler_output.scheduled_cached_reqs) == 0
69
+ assert scheduler_output.scheduled_cached_reqs.num_reqs == 0
70
assert len(scheduler.finished_req_ids) == 0
71
72
# (2b): execute_model()
@@ -81,7 +81,7 @@ def test_basic_lifecycle():
81
assert len(scheduler.running) == 0
82
assert len(scheduler_output.finished_req_ids) == 0
83
assert len(scheduler_output.scheduled_new_reqs) == 0
84
- assert len(scheduler_output.scheduled_cached_reqs) == 0
84
+ assert scheduler_output.scheduled_cached_reqs.num_reqs == 0
85
assert len(scheduler.finished_req_ids) == 0
86
87
# (3b): execute_model()
tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py
+6
-6
@@ -36,7 +36,7 @@ def test_basic_lifecycle():
36
# Nothing running and empty scheduler output.
37
assert len(scheduler.running) == 0
38
assert len(scheduler_output.scheduled_new_reqs) == 0
39
- assert len(scheduler_output.scheduled_cached_reqs) == 0
39
+ assert scheduler_output.scheduled_cached_reqs.num_reqs == 0
40
assert len(scheduler_output.num_scheduled_tokens) == 0
41
assert scheduler_output.total_num_scheduled_tokens == 0
42
@@ -158,7 +158,7 @@ def test_interleaved_lifecycle():
158
assert len(scheduler.running) == 2
159
assert len(scheduler.waiting) == 1
160
assert len(scheduler_output.scheduled_new_reqs) == 1
161
- assert len(scheduler_output.scheduled_cached_reqs) == 1
161
+ assert scheduler_output.scheduled_cached_reqs.num_reqs == 1
162
163
model_runner_output = create_model_runner_output(
164
[request_local_a, request_local_b])
@@ -169,7 +169,7 @@ def test_interleaved_lifecycle():
169
assert len(scheduler.running) == 2
170
assert len(scheduler.waiting) == 1
171
assert len(scheduler_output.scheduled_new_reqs) == 0
172
- assert len(scheduler_output.scheduled_cached_reqs) == 2
172
+ assert scheduler_output.scheduled_cached_reqs.num_reqs == 2
173
174
model_runner_output = create_model_runner_output(
175
reqs=[request_local_a, request_local_b])
@@ -177,14 +177,14 @@ def test_interleaved_lifecycle():
177
assert len(scheduler.running) == 2
178
assert len(scheduler.waiting) == 1
179
assert len(scheduler_output.scheduled_new_reqs) == 0
180
- assert len(scheduler_output.scheduled_cached_reqs) == 2
180
+ assert scheduler_output.scheduled_cached_reqs.num_reqs == 2
181
182
# STEP 4: KVs arrive.
183
scheduler_output = scheduler.schedule()
184
assert len(scheduler.running) == 2
185
assert len(scheduler.waiting) == 1
186
assert len(scheduler_output.scheduled_new_reqs) == 0
187
- assert len(scheduler_output.scheduled_cached_reqs) == 2
187
+ assert scheduler_output.scheduled_cached_reqs.num_reqs == 2
188
189
model_runner_output = create_model_runner_output(
190
[request_local_a, request_local_b],
@@ -196,7 +196,7 @@ def test_interleaved_lifecycle():
196
assert len(scheduler.running) == 3
197
assert len(scheduler.waiting) == 0
198
assert len(scheduler_output.scheduled_new_reqs) == 1
199
- assert len(scheduler_output.scheduled_cached_reqs) == 2
199
+ assert scheduler_output.scheduled_cached_reqs.num_reqs == 2
200
201
model_runner_output = create_model_runner_output(
202
[request_local_a, request_local_b, request_remote])
tests/v1/kv_connector/unit/utils.py
-1
@@ -25,7 +25,6 @@ def assert_scheduler_empty(scheduler: Scheduler):
25
assert len(scheduler.running) == 0
26
assert len(scheduler.finished_req_ids) == 0
27
assert len(scheduler.finished_recving_kv_req_ids) == 0
28
- assert len(scheduler._cached_reqs_data) == 0
28
29
# EncoderCacheManager.
30
assert len(scheduler.encoder_cache_manager.freed) == 0
tests/v1/tpu/worker/test_tpu_model_runner.py
+11
-11
@@ -82,7 +82,7 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:
82
83
return SchedulerOutput(
84
scheduled_new_reqs=new_reqs,
85
- scheduled_cached_reqs=[],
85
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
86
num_scheduled_tokens=num_scheduled_tokens,
87
total_num_scheduled_tokens=total_num_scheduled_tokens,
88
scheduled_spec_decode_tokens={},
@@ -161,7 +161,7 @@ def test_update_states_request_finished(model_runner):
161
# finish req
162
scheduler_output = SchedulerOutput(
163
scheduled_new_reqs=[],
164
- scheduled_cached_reqs=[],
164
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
165
num_scheduled_tokens={},
166
total_num_scheduled_tokens=0,
167
scheduled_spec_decode_tokens={},
@@ -191,7 +191,7 @@ def test_update_states_request_resumed(model_runner):
191
# unschedule req
192
scheduler_output = SchedulerOutput(
193
scheduled_new_reqs=[],
194
- scheduled_cached_reqs=[],
194
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
195
num_scheduled_tokens={},
196
total_num_scheduled_tokens=0,
197
scheduled_spec_decode_tokens={},
@@ -209,16 +209,16 @@ def test_update_states_request_resumed(model_runner):
209
210
# resume req
211
cached_req_data = CachedRequestData(
212
- req_id=req_id,
213
- resumed_from_preemption=False,
214
- new_token_ids=[],
215
- new_block_ids=([], ),
216
- num_computed_tokens=0,
212
+ req_ids=[req_id],
213
+ resumed_from_preemption=[False],
214
+ new_token_ids=[[]],
215
+ new_block_ids=[([], )],
216
+ num_computed_tokens=[0],
217
)
218
219
scheduler_output = SchedulerOutput(
220
scheduled_new_reqs=[],
221
- scheduled_cached_reqs=[cached_req_data],
221
+ scheduled_cached_reqs=cached_req_data,
222
num_scheduled_tokens={req_id: 1},
223
total_num_scheduled_tokens=1,
224
scheduled_spec_decode_tokens={},
@@ -249,7 +249,7 @@ def test_update_states_no_changes(model_runner):
249
# schedule req
250
scheduler_output = SchedulerOutput(
251
scheduled_new_reqs=[],
252
- scheduled_cached_reqs=[],
252
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
253
num_scheduled_tokens={req_id: 1},
254
total_num_scheduled_tokens=1,
255
scheduled_spec_decode_tokens={},
@@ -284,7 +284,7 @@ def test_update_states_request_unscheduled(model_runner):
284
# unschedule req_1
285
scheduler_output = SchedulerOutput(
286
scheduled_new_reqs=[],
287
- scheduled_cached_reqs=[],
287
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
288
num_scheduled_tokens={req_ids[0]: 1},
289
total_num_scheduled_tokens=1,
290
scheduled_spec_decode_tokens={},
tests/v1/worker/test_gpu_model_runner.py
+11
-11
@@ -133,7 +133,7 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:
133
134
return SchedulerOutput(
135
scheduled_new_reqs=new_reqs,
136
- scheduled_cached_reqs=[],
136
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
137
num_scheduled_tokens=num_scheduled_tokens,
138
total_num_scheduled_tokens=total_num_scheduled_tokens,
139
scheduled_spec_decode_tokens={},
@@ -199,7 +199,7 @@ def test_update_states_request_finished(model_runner):
199
# finish req
200
scheduler_output = SchedulerOutput(
201
scheduled_new_reqs=[],
202
- scheduled_cached_reqs=[],
202
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
203
num_scheduled_tokens={},
204
total_num_scheduled_tokens=0,
205
scheduled_spec_decode_tokens={},
@@ -231,7 +231,7 @@ def test_update_states_request_resumed(model_runner):
231
# unschedule req
232
scheduler_output = SchedulerOutput(
233
scheduled_new_reqs=[],
234
- scheduled_cached_reqs=[],
234
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
235
num_scheduled_tokens={},
236
total_num_scheduled_tokens=0,
237
scheduled_spec_decode_tokens={},
@@ -249,16 +249,16 @@ def test_update_states_request_resumed(model_runner):
249
250
# resume req
251
cached_req_data = CachedRequestData(
252
- req_id=req_id,
253
- resumed_from_preemption=False,
254
- new_token_ids=[],
255
- new_block_ids=([], ),
256
- num_computed_tokens=0,
252
+ req_ids=[req_id],
253
+ resumed_from_preemption=[False],
254
+ new_token_ids=[[]],
255
+ new_block_ids=([[0]], ),
256
+ num_computed_tokens=[0],
257
)
258
259
scheduler_output = SchedulerOutput(
260
scheduled_new_reqs=[],
261
- scheduled_cached_reqs=[cached_req_data],
261
+ scheduled_cached_reqs=cached_req_data,
262
num_scheduled_tokens={req_id: 1},
263
total_num_scheduled_tokens=1,
264
scheduled_spec_decode_tokens={},
@@ -339,7 +339,7 @@ def test_update_states_no_changes(model_runner):
339
# schedule req
340
scheduler_output = SchedulerOutput(
341
scheduled_new_reqs=[],
342
- scheduled_cached_reqs=[],
342
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
343
num_scheduled_tokens={req_id: 1},
344
total_num_scheduled_tokens=1,
345
scheduled_spec_decode_tokens={},
@@ -376,7 +376,7 @@ def test_update_states_request_unscheduled(model_runner):
376
# unschedule req_1
377
scheduler_output = SchedulerOutput(
378
scheduled_new_reqs=[],
379
- scheduled_cached_reqs=[],
379
+ scheduled_cached_reqs=CachedRequestData.make_empty(),
380
num_scheduled_tokens={req_ids[0]: 1},
381
total_num_scheduled_tokens=1,
382
scheduled_spec_decode_tokens={},
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py
+23
-20
@@ -371,45 +371,48 @@ class P2pNcclConnector(KVConnectorBase_V1):
371
block_size=self._block_size)
372
self._requests_need_load.pop(new_req.req_id)
373
374
- for cached_req in scheduler_output.scheduled_cached_reqs:
374
+ cached_reqs = scheduler_output.scheduled_cached_reqs
375
+ for i, req_id in enumerate(cached_reqs.req_ids):
376
+ num_computed_tokens = cached_reqs.num_computed_tokens[i]
377
+ new_block_ids = cached_reqs.new_block_ids[i]
378
+ resumed_from_preemption = cached_reqs.resumed_from_preemption[i]
379
+
380
if self.is_producer:
381
num_scheduled_tokens = (
377
- scheduler_output.num_scheduled_tokens)[cached_req.req_id]
378
- num_tokens = (num_scheduled_tokens +
379
- cached_req.num_computed_tokens)
380
- assert cached_req.req_id in self.chunked_prefill
381
- block_ids = cached_req.new_block_ids[0]
382
- if not cached_req.resumed_from_preemption:
383
- block_ids = (self.chunked_prefill[cached_req.req_id][0] +
384
- block_ids)
385
- prompt_token_ids = self.chunked_prefill[cached_req.req_id][1]
382
+ scheduler_output.num_scheduled_tokens)[req_id]
383
+ num_tokens = (num_scheduled_tokens + num_computed_tokens)
384
+ assert req_id in self.chunked_prefill
385
+ block_ids = new_block_ids[0]
386
+ if not resumed_from_preemption:
387
+ block_ids = (self.chunked_prefill[req_id][0] + block_ids)
388
+ prompt_token_ids = self.chunked_prefill[req_id][1]
389
# the request's prompt is chunked prefill again
390
if num_tokens < len(prompt_token_ids):
388
- self.chunked_prefill[cached_req.req_id] = (
389
- block_ids, prompt_token_ids)
391
+ self.chunked_prefill[req_id] = (block_ids,
392
+ prompt_token_ids)
393
continue
394
# the request's prompt is all prefilled finally
392
- meta.add_request(request_id=cached_req.req_id,
395
+ meta.add_request(request_id=req_id,
396
token_ids=prompt_token_ids,
397
block_ids=block_ids,
398
block_size=self._block_size)
396
- self.chunked_prefill.pop(cached_req.req_id, None)
399
+ self.chunked_prefill.pop(req_id, None)
400
continue
401
402
# NOTE(rob): here we rely on the resumed requests being
403
# the first N requests in the list scheduled_cache_reqs.
401
- if not cached_req.resumed_from_preemption:
404
+ if not resumed_from_preemption:
405
break
403
- if cached_req.req_id in self._requests_need_load:
404
- request, _ = self._requests_need_load.pop(cached_req.req_id)
405
- total_tokens = cached_req.num_computed_tokens + 1
406
+ if req_id in self._requests_need_load:
407
+ request, _ = self._requests_need_load.pop(req_id)
408
+ total_tokens = num_computed_tokens + 1
409
token_ids = request.all_token_ids[:total_tokens]
410
411
# NOTE(rob): For resumed req, new_block_ids is all
412
# of the block_ids for the request.
410
- block_ids = cached_req.new_block_ids[0]
413
+ block_ids = new_block_ids[0]
414
412
- meta.add_request(request_id=cached_req.req_id,
415
+ meta.add_request(request_id=req_id,
416
token_ids=token_ids,
417
block_ids=block_ids,
418
block_size=self._block_size)
vllm/distributed/kv_transfer/kv_connector/v1/shared_storage_connector.py
+12
-7
@@ -304,23 +304,28 @@ class SharedStorageConnector(KVConnectorBase_V1):
304
block_size=self._block_size,
305
is_store=True)
306
307
- for cached_req in scheduler_output.scheduled_cached_reqs:
307
+ cached_reqs = scheduler_output.scheduled_cached_reqs
308
+ for i, req_id in enumerate(cached_reqs.req_ids):
309
+ num_computed_tokens = cached_reqs.num_computed_tokens[i]
310
+ new_token_ids = cached_reqs.new_token_ids[i]
311
+ new_block_ids = cached_reqs.new_block_ids[i]
312
+ resumed_from_preemption = cached_reqs.resumed_from_preemption[i]
313
+
314
# NOTE(rob): here we rely on the resumed requests being
315
# the first N requests in the list scheduled_cache_reqs.
310
- if not cached_req.resumed_from_preemption:
316
+ if not resumed_from_preemption:
317
break
312
- if cached_req.req_id in self._requests_need_load:
318
+ if req_id in self._requests_need_load:
319
# NOTE(rob): cached_req_data does not have the full
320
# list of token ids (only new tokens). So we look it
321
# up in the actual request object.
316
- request = self._requests_need_load[cached_req.req_id]
317
- total_tokens = (len(cached_req.new_token_ids) +
318
- cached_req.num_computed_tokens)
322
+ request = self._requests_need_load[req_id]
323
+ total_tokens = (len(new_token_ids) + num_computed_tokens)
324
token_ids = request.all_token_ids[:total_tokens]
325
326
# NOTE(rob): For resumed req, new_block_ids is all
327
# of the block_ids for the request.
323
- block_ids = cached_req.new_block_ids[0]
328
+ block_ids = new_block_ids[0]
329
330
meta.add_request(token_ids=token_ids,
331
block_ids=block_ids,
vllm/v1/core/sched/output.py
+16
-18
@@ -83,29 +83,27 @@ class NewRequestData:
83
@dataclass
84
class CachedRequestData:
85
86
- req_id: str
86
+ req_ids: list[str]
87
# If resumed_from_preemption is False, new_block_ids will be appended to
88
# the request's block IDs. If True, new_block_ids will be used as the
89
# request's block IDs instead of appending to the existing block IDs.
90
- resumed_from_preemption: bool
91
- new_token_ids: list[int]
92
- new_block_ids: tuple[list[int], ...]
93
- num_computed_tokens: int
90
+ resumed_from_preemption: list[bool]
91
+ new_token_ids: list[list[int]]
92
+ new_block_ids: list[tuple[list[int], ...]]
93
+ num_computed_tokens: list[int]
94
+
95
+ @property
96
+ def num_reqs(self) -> int:
97
+ return len(self.req_ids)
98
99
@classmethod
96
- def from_request(
97
- cls,
98
- request: Request,
99
- resumed_from_preemption: bool,
100
- new_token_ids: list[int],
101
- new_block_ids: tuple[list[int], ...],
102
- ) -> CachedRequestData:
100
+ def make_empty(cls) -> CachedRequestData:
101
return cls(
104
- req_id=request.request_id,
105
- resumed_from_preemption=resumed_from_preemption,
106
- new_token_ids=new_token_ids,
107
- new_block_ids=new_block_ids,
108
- num_computed_tokens=request.num_computed_tokens,
102
+ req_ids=[],
103
+ resumed_from_preemption=[],
104
+ new_token_ids=[],
105
+ new_block_ids=[],
106
+ num_computed_tokens=[],
107
)
108
109
@@ -119,7 +117,7 @@ class SchedulerOutput:
117
# list of the requests that have been scheduled before.
118
# Since the request's data is already cached in the worker processes,
119
# we only send the diff to minimize the communication cost.
122
- scheduled_cached_reqs: list[CachedRequestData]
120
+ scheduled_cached_reqs: CachedRequestData
121
122
# req_id -> num_scheduled_tokens
123
# Number of tokens scheduled for each request.
vllm/v1/core/sched/scheduler.py
+43
-65
@@ -3,8 +3,9 @@
3
4
from __future__ import annotations
5
6
+import itertools
7
import time
7
-from collections import defaultdict, deque
8
+from collections import defaultdict
9
from collections.abc import Iterable
10
from typing import Any, Optional, Union
11
@@ -117,12 +118,6 @@ class Scheduler(SchedulerInterface):
118
# KV Connector: requests in process of async KV loading or recving
119
self.finished_recving_kv_req_ids: set[str] = set()
120
120
- # OPTIMIZATION: Cache the CachedRequestData objects to avoid creating
121
- # them at each scheduling step.
122
- # Request id -> deque of CachedRequestData
123
- self._cached_reqs_data: dict[
124
- str, deque[CachedRequestData]] = defaultdict(deque)
125
-
121
# Encoder-related.
122
# Calculate encoder cache size if applicable
123
# NOTE: For now we use the same budget for both compute and space.
@@ -547,27 +542,16 @@ class Scheduler(SchedulerInterface):
542
req_to_new_block_ids[req.request_id])
543
for req in scheduled_new_reqs
544
]
550
- resumed_reqs_data = [
551
- self._make_cached_request_data(
552
- req,
553
- num_scheduled_tokens[req.request_id],
554
- len(scheduled_spec_decode_tokens.get(req.request_id, ())),
555
- req_to_new_block_ids[req.request_id],
556
- resumed_from_preemption=True,
557
- ) for req in scheduled_resumed_reqs
558
- ]
559
- running_reqs_data = [
560
- self._make_cached_request_data(
561
- req,
562
- num_scheduled_tokens[req.request_id],
563
- len(scheduled_spec_decode_tokens.get(req.request_id, ())),
564
- req_to_new_block_ids[req.request_id],
565
- resumed_from_preemption=False,
566
- ) for req in scheduled_running_reqs
567
- ]
545
+ cached_reqs_data = self._make_cached_request_data(
546
+ scheduled_running_reqs,
547
+ scheduled_resumed_reqs,
548
+ num_scheduled_tokens,
549
+ scheduled_spec_decode_tokens,
550
+ req_to_new_block_ids,
551
+ )
552
scheduler_output = SchedulerOutput(
553
scheduled_new_reqs=new_reqs_data,
570
- scheduled_cached_reqs=resumed_reqs_data + running_reqs_data,
554
+ scheduled_cached_reqs=cached_reqs_data,
555
num_scheduled_tokens=num_scheduled_tokens,
556
total_num_scheduled_tokens=total_num_scheduled_tokens,
557
scheduled_spec_decode_tokens=scheduled_spec_decode_tokens,
@@ -613,34 +597,39 @@ class Scheduler(SchedulerInterface):
597
598
def _make_cached_request_data(
599
self,
616
- request: Request,
617
- num_scheduled_tokens: int,
618
- num_scheduled_spec_tokens: int,
619
- new_block_ids: tuple[list[int], ...],
620
- resumed_from_preemption: bool,
600
+ running_reqs: list[Request],
601
+ resumed_reqs: list[Request],
602
+ num_scheduled_tokens: dict[str, int],
603
+ spec_decode_tokens: dict[str, list[int]],
604
+ req_to_new_block_ids: dict[str, tuple[list[int], ...]],
605
) -> CachedRequestData:
622
- # OPTIMIZATION: Cache the CachedRequestData objects to avoid creating
623
- # them at each scheduling step.
624
- num_computed_tokens = request.num_computed_tokens
625
- num_regular_tokens = num_scheduled_tokens - num_scheduled_spec_tokens
626
- new_token_ids = request.all_token_ids[
627
- num_computed_tokens:num_computed_tokens + num_regular_tokens]
628
-
629
- req_data_queue = self._cached_reqs_data.get(request.request_id)
630
- if req_data_queue:
631
- req_data = req_data_queue.popleft()
632
- req_data.resumed_from_preemption = resumed_from_preemption
633
- req_data.new_token_ids = new_token_ids
634
- req_data.new_block_ids = new_block_ids
635
- req_data.num_computed_tokens = num_computed_tokens
636
- else:
637
- # No cached request data, or all cached request data has been
638
- # used by the scheduled requests.
639
- req_data = CachedRequestData.from_request(request,
640
- resumed_from_preemption,
641
- new_token_ids,
642
- new_block_ids)
643
- return req_data
606
+ req_ids: list[str] = []
607
+ new_token_ids: list[list[int]] = []
608
+ new_block_ids: list[tuple[list[int], ...]] = []
609
+ num_computed_tokens: list[int] = []
610
+
611
+ for req in itertools.chain(running_reqs, resumed_reqs):
612
+ req_id = req.request_id
613
+ req_ids.append(req_id)
614
+ num_tokens = (num_scheduled_tokens[req_id] -
615
+ len(spec_decode_tokens.get(req_id, ())))
616
+ token_ids = req.all_token_ids[req.num_computed_tokens:req.
617
+ num_computed_tokens + num_tokens]
618
+ new_token_ids.append(token_ids)
619
+ new_block_ids.append(req_to_new_block_ids[req_id])
620
+ num_computed_tokens.append(req.num_computed_tokens)
621
+ # Because resumed_reqs is usually empty, it is more efficient to do
622
+ # in-place appending so that we don't need to allocate a new list.
623
+ resumed_from_preemption = [False] * len(running_reqs)
624
+ resumed_from_preemption += [True] * len(resumed_reqs)
625
+
626
+ return CachedRequestData(
627
+ req_ids=req_ids,
628
+ resumed_from_preemption=resumed_from_preemption,
629
+ new_token_ids=new_token_ids,
630
+ new_block_ids=new_block_ids,
631
+ num_computed_tokens=num_computed_tokens,
632
+ )
633
634
def _try_schedule_encoder_inputs(
635
self,
@@ -870,19 +859,11 @@ class Scheduler(SchedulerInterface):
859
860
if not stopped:
861
new_running.append(request)
862
+ self.running = new_running
863
864
# KV Connector: update state for finished KV Transfers.
865
self._update_from_kv_xfer_finished(model_runner_output)
866
877
- # Return the cached request data to the queue so they can be reused.
878
- for req_data in scheduler_output.scheduled_cached_reqs:
879
- # NOTE(rob): since we free stopped reqs above, adding stopped reqs
880
- # to _cached_reqs_data will cause a memory leak.
881
- if req_data.req_id not in self.finished_req_ids:
882
- self._cached_reqs_data[req_data.req_id].append(req_data)
883
-
884
- self.running = new_running
885
-
867
# Create EngineCoreOutputs for all clients that have requests with
868
# outputs in this step.
869
engine_core_outputs = {
@@ -965,13 +946,11 @@ class Scheduler(SchedulerInterface):
946
self._free_request(request)
947
948
def _free_request(self, request: Request) -> Optional[dict[str, Any]]:
968
-
949
assert request.is_finished()
950
951
delay_free_blocks, kv_xfer_params = self._connector_finished(request)
952
self.encoder_cache_manager.free(request)
953
request_id = request.request_id
974
- self._cached_reqs_data.pop(request_id, None)
954
self.finished_req_ids.add(request_id)
955
if self.finished_req_ids_dict is not None:
956
self.finished_req_ids_dict[request.client_index].add(request_id)
@@ -983,7 +962,6 @@ class Scheduler(SchedulerInterface):
962
963
def _free_blocks(self, request: Request):
964
assert request.is_finished()
986
- assert request.request_id not in self._cached_reqs_data
965
self.kv_cache_manager.free(request)
966
self.kv_cache_manager.free_block_hashes(request)
967
del self.requests[request.request_id]
vllm/v1/worker/gpu_model_runner.py
+17
-17
@@ -470,34 +470,36 @@ class GPUModelRunner(LoRAModelRunnerMixin):
470
req_ids_to_add.append(req_id)
471
472
# Update the states of the running/resumed requests.
473
- for req_data in scheduler_output.scheduled_cached_reqs:
474
- req_id = req_data.req_id
473
+ req_data = scheduler_output.scheduled_cached_reqs
474
+ for i, req_id in enumerate(req_data.req_ids):
475
req_state = self.requests[req_id]
476
+ num_computed_tokens = req_data.num_computed_tokens[i]
477
+ new_token_ids = req_data.new_token_ids[i]
478
+ new_block_ids = req_data.new_block_ids[i]
479
+ resumed_from_preemption = req_data.resumed_from_preemption[i]
480
481
# Update the cached states.
478
- num_computed_tokens = req_data.num_computed_tokens
482
req_state.num_computed_tokens = num_computed_tokens
483
# Add the sampled token(s) from the previous step (if any).
484
# This doesn't include "unverified" tokens like spec decode tokens.
482
- num_new_tokens = (num_computed_tokens +
483
- len(req_data.new_token_ids) -
485
+ num_new_tokens = (num_computed_tokens + len(new_token_ids) -
486
req_state.num_tokens)
487
if num_new_tokens == 1:
488
# Avoid slicing list in most common case.
487
- req_state.output_token_ids.append(req_data.new_token_ids[-1])
489
+ req_state.output_token_ids.append(new_token_ids[-1])
490
elif num_new_tokens > 0:
491
req_state.output_token_ids.extend(
490
- req_data.new_token_ids[-num_new_tokens:])
492
+ new_token_ids[-num_new_tokens:])
493
# Update the block IDs.
492
- if not req_data.resumed_from_preemption:
494
+ if not resumed_from_preemption:
495
# Append the new blocks to the existing block IDs.
494
- for block_ids, new_block_ids in zip(req_state.block_ids,
495
- req_data.new_block_ids):
496
- block_ids.extend(new_block_ids)
496
+ for block_ids, new_ids in zip(req_state.block_ids,
497
+ new_block_ids):
498
+ block_ids.extend(new_ids)
499
else:
500
# The request is resumed from preemption.
501
# Replace the existing block IDs with the new ones.
500
- req_state.block_ids = req_data.new_block_ids
502
+ req_state.block_ids = new_block_ids
503
504
req_index = self.input_batch.req_id_to_index.get(req_id)
505
if req_index is None:
@@ -510,14 +512,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
512
# Update the persistent batch.
513
self.input_batch.num_computed_tokens_cpu[req_index] = (
514
num_computed_tokens)
513
- self.input_batch.block_table.append_row(req_data.new_block_ids,
514
- req_index)
515
+ self.input_batch.block_table.append_row(new_block_ids, req_index)
516
# Add new_token_ids to token_ids_cpu.
517
start_token_index = num_computed_tokens
517
- end_token_index = num_computed_tokens + len(req_data.new_token_ids)
518
+ end_token_index = num_computed_tokens + len(new_token_ids)
519
self.input_batch.token_ids_cpu[
519
- req_index,
520
- start_token_index:end_token_index] = req_data.new_token_ids
520
+ req_index, start_token_index:end_token_index] = new_token_ids
521
self.input_batch.num_tokens_no_spec[req_index] = end_token_index
522
# Add spec_token_ids to token_ids_cpu.
523
spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(
vllm/v1/worker/tpu_model_runner.py
+13
-11
@@ -418,21 +418,24 @@ class TPUModelRunner(LoRAModelRunnerMixin):
418
req_ids_to_add.append(req_id)
419
420
# Update the states of the running/resumed requests.
421
- for req_data in scheduler_output.scheduled_cached_reqs:
422
- req_id = req_data.req_id
421
+ req_data = scheduler_output.scheduled_cached_reqs
422
+ for i, req_id in enumerate(req_data.req_ids):
423
req_state = self.requests[req_id]
424
+ num_computed_tokens = req_data.num_computed_tokens[i]
425
+ new_block_ids = req_data.new_block_ids[i]
426
+ resumed_from_preemption = req_data.resumed_from_preemption[i]
427
428
# Update the cached states.
426
- req_state.num_computed_tokens = req_data.num_computed_tokens
427
- if not req_data.resumed_from_preemption:
429
+ req_state.num_computed_tokens = num_computed_tokens
430
+ if not resumed_from_preemption:
431
# Append the new blocks to the existing block IDs.
429
- for block_ids, new_block_ids in zip(req_state.block_ids,
430
- req_data.new_block_ids):
431
- block_ids.extend(new_block_ids)
432
+ for block_ids, new_ids in zip(req_state.block_ids,
433
+ new_block_ids):
434
+ block_ids.extend(new_ids)
435
else:
436
# The request is resumed from preemption.
437
# Replace the existing block IDs with the new ones.
435
- req_state.block_ids = req_data.new_block_ids
438
+ req_state.block_ids = new_block_ids
439
440
req_index = self.input_batch.req_id_to_index.get(req_id)
441
if req_index is None:
@@ -444,9 +447,8 @@ class TPUModelRunner(LoRAModelRunnerMixin):
447
448
# Update the persistent batch.
449
self.input_batch.num_computed_tokens_cpu[req_index] = (
447
- req_data.num_computed_tokens)
448
- self.input_batch.block_table.append_row(req_data.new_block_ids,
449
- req_index)
450
+ num_computed_tokens)
451
+ self.input_batch.block_table.append_row(new_block_ids, req_index)
452
453
# Add the new or resumed requests to the persistent batch.
454
# The smaller empty indices are filled first.