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, 11:00 UTC
[v1] Introduce KVCacheBlocks as interface between Scheduler and KVCacheManager (#17479)
Signed-off-by: Chen Zhang <zhangch99@outlook.com>
Chen Zhang committed
May 6, 2025 at 23:50 UTC
aabcd2cae305c08c396889be9c6930840b1adead
4 files changed
+121
-99
tests/v1/core/test_kv_cache_utils.py
+3
-3
@@ -542,7 +542,7 @@ def test_allocate_with_lookahead():
542
num_tokens=3,
543
num_lookahead_tokens=2, # Total required: 3+2=5 tokens
544
)
545
- assert len(blocks) == 2 # ceil(5/4)=2 blocks
545
+ assert len(blocks.blocks) == 2 # ceil(5/4)=2 blocks
546
547
# Test case 2: With precomputed blocks
548
kv_cache_manager = KVCacheManager(kv_cache_config=config,
@@ -553,7 +553,7 @@ def test_allocate_with_lookahead():
553
num_tokens=3,
554
num_lookahead_tokens=2,
555
)
556
- assert len(blocks) == 2
556
+ assert len(blocks.blocks) == 2
557
558
# Test case 3: With precomputed blocks
559
# required_blocks = ceil((3 + 4) / 4) = 2
@@ -564,4 +564,4 @@ def test_allocate_with_lookahead():
564
num_tokens=3,
565
num_lookahead_tokens=4,
566
)
567
- assert len(blocks) == 2
567
+ assert len(blocks.blocks) == 2
tests/v1/core/test_prefix_caching.py
+70
-69
@@ -79,10 +79,10 @@ def test_prefill(hash_algo):
79
req0 = make_request("0", all_token_ids)
80
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
81
assert len(manager.req_to_block_hashes[req0.request_id]) == 3
82
- assert not computed_blocks
82
+ assert not computed_blocks.blocks
83
assert num_computed_tokens == 0
84
blocks = manager.allocate_slots(req0, 55, computed_blocks)
85
- assert [b.block_id for b in blocks] == [1, 2, 3, 4]
85
+ assert blocks.get_block_ids() == [1, 2, 3, 4]
86
87
# Check full block metadata
88
parent_block_hash = None
@@ -105,12 +105,12 @@ def test_prefill(hash_algo):
105
req1 = make_request("1", common_token_ids + unique_token_ids)
106
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
107
assert len(manager.req_to_block_hashes[req1.request_id]) == 3
108
- assert [b.block_id for b in computed_blocks] == [1, 2, 3]
108
+ assert computed_blocks.get_block_ids() == [1, 2, 3]
109
assert num_computed_tokens == 3 * 16
110
num_new_tokens = 53 - 3 * 16
111
blocks = manager.allocate_slots(req1, num_new_tokens, computed_blocks)
112
- assert [b.block_id for b in blocks] == [5]
113
- for block in computed_blocks:
112
+ assert blocks.get_block_ids() == [5]
113
+ for block in computed_blocks.blocks:
114
assert block.ref_cnt == 2
115
116
# At this point, we should have 5 free blocks left.
@@ -137,11 +137,11 @@ def test_prefill(hash_algo):
137
req2 = make_request("2", common_token_ids + unique_token_ids)
138
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
139
assert len(manager.req_to_block_hashes[req2.request_id]) == 3
140
- assert [b.block_id for b in computed_blocks] == [1, 2, 3]
140
+ assert computed_blocks.get_block_ids() == [1, 2, 3]
141
assert num_computed_tokens == 3 * 16
142
num_new_tokens = 53 - 3 * 16
143
blocks = manager.allocate_slots(req2, num_new_tokens, computed_blocks)
144
- assert [b.block_id for b in blocks] == [6]
144
+ assert blocks.get_block_ids() == [6]
145
146
# Although we only have 6 free blocks, we have 8 blocks in
147
# the free block queue due to lazy removal.
@@ -159,11 +159,11 @@ def test_prefill(hash_algo):
159
# Cache miss and eviction.
160
req3 = make_request("3", [99] * (16 * 10))
161
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req3)
162
- assert not computed_blocks
162
+ assert not computed_blocks.blocks
163
assert num_computed_tokens == 0
164
blocks = manager.allocate_slots(req3, 16 * 10, computed_blocks)
165
# This block ID order also checks the eviction order.
166
- assert [b.block_id for b in blocks] == [7, 8, 9, 10, 4, 5, 6, 3, 2, 1]
166
+ assert blocks.get_block_ids() == [7, 8, 9, 10, 4, 5, 6, 3, 2, 1]
167
assert manager.block_pool.free_block_queue.num_free_blocks == 0
168
assert manager.block_pool.free_block_queue.free_list_head is None
169
assert manager.block_pool.free_block_queue.free_list_tail is None
@@ -195,11 +195,11 @@ def test_prefill_plp():
195
req0 = make_request("0", all_token_ids, prompt_logprobs=5)
196
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
197
assert len(manager.req_to_block_hashes[req0.request_id]) == 3
198
- assert not computed_blocks
198
+ assert not computed_blocks.blocks
199
assert num_computed_tokens == 0
200
blocks = manager.allocate_slots(req0, 55, computed_blocks)
201
- assert [b.block_id for b in blocks] == [1, 2, 3, 4]
202
- req0_block_hashes = [b.block_hash for b in blocks]
201
+ assert blocks.get_block_ids() == [1, 2, 3, 4]
202
+ req0_block_hashes = [b.block_hash for b in blocks.blocks]
203
204
# Check full block metadata
205
parent_block_hash = None
@@ -223,12 +223,12 @@ def test_prefill_plp():
223
req1 = make_request("1", common_token_ids + unique_token_ids)
224
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
225
assert len(manager.req_to_block_hashes[req1.request_id]) == 3
226
- assert [b.block_id for b in computed_blocks] == [1, 2, 3]
226
+ assert computed_blocks.get_block_ids() == [1, 2, 3]
227
assert num_computed_tokens == 3 * 16
228
num_new_tokens = 53 - 3 * 16
229
blocks = manager.allocate_slots(req1, num_new_tokens, computed_blocks)
230
- assert [b.block_id for b in blocks] == [5]
231
- for block in computed_blocks:
230
+ assert blocks.get_block_ids() == [5]
231
+ for block in computed_blocks.blocks:
232
assert block.ref_cnt == 2
233
234
# At this point, we should have 5 free blocks left.
@@ -257,12 +257,12 @@ def test_prefill_plp():
257
prompt_logprobs=5)
258
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
259
assert len(manager.req_to_block_hashes[req2.request_id]) == 3
260
- assert not computed_blocks
260
+ assert not computed_blocks.blocks
261
assert num_computed_tokens == 0
262
blocks = manager.allocate_slots(req2, 55, computed_blocks)
263
- block_ids = [b.block_id for b in blocks]
263
+ block_ids = blocks.get_block_ids()
264
# Duplicate cached blocks have different ids but same hashes vs request #0
265
- assert [b.block_hash for b in blocks] == req0_block_hashes
265
+ assert [b.block_hash for b in blocks.blocks] == req0_block_hashes
266
assert block_ids != [1, 2, 3, 4]
267
268
# Request #2 block hashes are valid since request #0 hashes are.
@@ -288,17 +288,17 @@ def test_decode():
288
unique_token_ids = [3] * 7
289
req0 = make_request("0", common_token_ids + unique_token_ids)
290
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
291
- assert not computed_blocks
291
+ assert not computed_blocks.blocks
292
assert num_computed_tokens == 0
293
blocks = manager.allocate_slots(req0, 55, computed_blocks)
294
- assert [b.block_id for b in blocks] == [1, 2, 3, 4]
294
+ assert blocks.get_block_ids() == [1, 2, 3, 4]
295
296
# Append slots without allocating a new block.
297
req0.num_computed_tokens = 55
298
for _ in range(4):
299
req0.append_output_token_ids(8)
300
new_blocks = manager.allocate_slots(req0, 4)
301
- assert new_blocks is not None and len(new_blocks) == 0
301
+ assert new_blocks is not None and len(new_blocks.blocks) == 0
302
assert manager.req_to_blocks[req0.request_id][-1].block_hash is None
303
304
# Append slots with allocating a new block.
@@ -308,7 +308,7 @@ def test_decode():
308
for _ in range(9 + 10):
309
req0.append_output_token_ids(7)
310
new_blocks = manager.allocate_slots(req0, 19)
311
- assert new_blocks is not None and len(new_blocks) == 1
311
+ assert new_blocks is not None and len(new_blocks.blocks) == 1
312
assert manager.req_to_blocks[req0.request_id][-2].block_hash is not None
313
assert manager.req_to_blocks[req0.request_id][-1].block_hash is None
314
@@ -323,19 +323,19 @@ def test_evict():
323
last_token_id = 5 * 16 + 7
324
req0 = make_request("0", list(range(last_token_id)))
325
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
326
- assert not computed_blocks
326
+ assert not computed_blocks.blocks
327
assert num_computed_tokens == 0
328
blocks = manager.allocate_slots(req0, 5 * 16 + 7, computed_blocks)
329
- assert len(blocks) == 6 # 5 full + 1 partial
329
+ assert len(blocks.blocks) == 6 # 5 full + 1 partial
330
331
# 3 blocks.
332
req1 = make_request("1", list(range(last_token_id,
333
last_token_id + 3 * 16)))
334
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
335
- assert not computed_blocks
335
+ assert not computed_blocks.blocks
336
assert num_computed_tokens == 0
337
blocks = manager.allocate_slots(req1, 3 * 16, computed_blocks)
338
- assert len(blocks) == 3 # 3 full blocks
338
+ assert len(blocks.blocks) == 3 # 3 full blocks
339
last_token_id += 3 * 16
340
341
# 10 - (6 + 3) == 1
@@ -352,10 +352,10 @@ def test_evict():
352
# Touch the first 2 blocks.
353
req2 = make_request("2", list(range(2 * 16 + 3)))
354
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
355
- assert [b.block_id for b in computed_blocks] == [1, 2]
355
+ assert computed_blocks.get_block_ids() == [1, 2]
356
assert num_computed_tokens == 2 * 16
357
blocks = manager.allocate_slots(req2, 3, computed_blocks)
358
- assert [b.block_id for b in blocks] == [10]
358
+ assert blocks.get_block_ids() == [10]
359
assert manager.block_pool.free_block_queue.num_free_blocks == 7
360
361
@@ -375,10 +375,10 @@ def test_hash_block_correct_reuse():
375
num_tokens = block_size * 1
376
req = make_request("0", list(range(num_tokens)))
377
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)
378
- assert not computed_blocks
378
+ assert not computed_blocks.blocks
379
assert num_computed_tokens == 0
380
blocks = manager.allocate_slots(req, num_tokens, computed_blocks)
381
- assert len(blocks) == 1
381
+ assert len(blocks.blocks) == 1
382
383
# Deallocate the block.
384
manager.free(req)
@@ -387,12 +387,13 @@ def test_hash_block_correct_reuse():
387
# block is cleared.
388
req = make_request("1", list(range(num_tokens - 1)))
389
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)
390
- assert not computed_blocks
390
+ assert not computed_blocks.blocks
391
assert num_computed_tokens == 0
392
blocks = manager.allocate_slots(req, num_tokens - 1, computed_blocks)
393
- assert len(blocks) == 1
393
+ assert len(blocks.blocks) == 1
394
395
- assert manager.block_pool.blocks[blocks[0].block_id].block_hash is None
395
+ assert manager.block_pool.blocks[
396
+ blocks.blocks[0].block_id].block_hash is None
397
398
399
def test_computed_blocks_not_evicted():
@@ -411,20 +412,20 @@ def test_computed_blocks_not_evicted():
412
num_tokens = block_size * 1
413
req0 = make_request("0", list(range(num_tokens)))
414
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
414
- assert not computed_blocks
415
+ assert not computed_blocks.blocks
416
assert num_computed_tokens == 0
417
blocks = manager.allocate_slots(req0, num_tokens, computed_blocks)
417
- assert len(blocks) == 1
418
- assert blocks[0].block_id == 1
418
+ assert len(blocks.blocks) == 1
419
+ assert blocks.blocks[0].block_id == 1
420
421
# Allocate another block.
422
req1 = make_request("1", list(range(num_tokens, num_tokens * 2)))
423
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
423
- assert not computed_blocks
424
+ assert not computed_blocks.blocks
425
assert num_computed_tokens == 0
426
blocks = manager.allocate_slots(req1, num_tokens, computed_blocks)
426
- assert len(blocks) == 1
427
- assert blocks[0].block_id == 2
427
+ assert len(blocks.blocks) == 1
428
+ assert blocks.blocks[0].block_id == 2
429
430
# Free the blocks.
431
manager.free(req0)
@@ -434,14 +435,14 @@ def test_computed_blocks_not_evicted():
435
# cached block rather than the first one.
436
req2 = make_request("2", list(range(num_tokens * 2)))
437
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
437
- assert len(computed_blocks) == 1
438
- assert computed_blocks[0].block_id == 1
438
+ assert len(computed_blocks.blocks) == 1
439
+ assert computed_blocks.blocks[0].block_id == 1
440
assert num_computed_tokens == block_size
441
442
blocks = manager.allocate_slots(req2, num_tokens * 2 - num_tokens,
443
computed_blocks)
443
- assert len(blocks) == 1
444
- assert blocks[0].block_id == 2
444
+ assert len(blocks.blocks) == 1
445
+ assert blocks.blocks[0].block_id == 2
446
447
448
def test_basic_prefix_caching_disabled():
@@ -458,10 +459,10 @@ def test_basic_prefix_caching_disabled():
459
req1 = make_request("1", list(range(10))) # 2 blocks and some more
460
461
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
461
- assert not computed_blocks
462
+ assert not computed_blocks.blocks
463
assert num_computed_tokens == 0
464
blocks = manager.allocate_slots(req1, 10, computed_blocks)
464
- assert len(blocks) == 3
465
+ assert len(blocks.blocks) == 3
466
467
# Free the blocks.
468
manager.free(req1)
@@ -469,15 +470,15 @@ def test_basic_prefix_caching_disabled():
470
# No caching.
471
req2 = make_request("2", list(range(16))) # shared prefix
472
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
472
- assert not computed_blocks
473
+ assert not computed_blocks.blocks
474
assert num_computed_tokens == 0
475
blocks = manager.allocate_slots(req2, 16, computed_blocks)
475
- assert len(blocks) == 4
476
+ assert len(blocks.blocks) == 4
477
478
# New requests should not have any blocks.
479
req3 = make_request("3", list(range(4)))
480
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req3)
480
- assert not computed_blocks
481
+ assert not computed_blocks.blocks
482
assert num_computed_tokens == 0
483
blocks = manager.allocate_slots(req3, 4, computed_blocks)
484
assert not blocks
@@ -569,7 +570,7 @@ def test_mm_prefix_caching():
570
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
571
572
# Completed block should have hashes with extra keys.
572
- assert not computed_blocks
573
+ assert not computed_blocks.blocks
574
assert num_computed_tokens == 0
575
block_hashes = manager.req_to_block_hashes[req0.request_id]
576
assert len(block_hashes) == 3
@@ -578,14 +579,14 @@ def test_mm_prefix_caching():
579
assert block_hashes[2].extra_keys == ("bbb", )
580
581
blocks = manager.allocate_slots(req0, 59, computed_blocks)
581
- assert [b.block_id for b in blocks] == [1, 2, 3, 4]
582
+ assert blocks.get_block_ids() == [1, 2, 3, 4]
583
req0.num_computed_tokens = 59
584
585
# Append slots without allocating a new block.
586
for _ in range(5):
587
req0.append_output_token_ids(8)
588
new_blocks = manager.allocate_slots(req0, 5)
588
- assert new_blocks is not None and len(new_blocks) == 0
589
+ assert new_blocks is not None and len(new_blocks.blocks) == 0
590
591
# The just completed block should have hashes with extra keys.
592
assert len(block_hashes) == 4
@@ -603,7 +604,7 @@ def test_mm_prefix_caching():
604
mm_positions=mm_positions,
605
mm_hashes=mm_hashes)
606
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
606
- assert len(computed_blocks) == 3
607
+ assert len(computed_blocks.blocks) == 3
608
assert num_computed_tokens == 3 * 16
609
610
@@ -626,7 +627,7 @@ def test_cache_key_salting():
627
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
628
629
# Completed block should have hashes with extra keys.
629
- assert not computed_blocks
630
+ assert not computed_blocks.blocks
631
assert num_computed_tokens == 0
632
block_hashes = manager.req_to_block_hashes[req0.request_id]
633
assert len(block_hashes) == 3
@@ -635,14 +636,14 @@ def test_cache_key_salting():
636
assert block_hashes[2].extra_keys is None
637
638
blocks = manager.allocate_slots(req0, 59, computed_blocks)
638
- assert [b.block_id for b in blocks] == [1, 2, 3, 4]
639
+ assert blocks.get_block_ids() == [1, 2, 3, 4]
640
req0.num_computed_tokens = 59
641
642
# Append slots without allocating a new block.
643
for _ in range(5):
644
req0.append_output_token_ids(8)
645
new_blocks = manager.allocate_slots(req0, 5)
645
- assert new_blocks is not None and len(new_blocks) == 0
646
+ assert new_blocks is not None and len(new_blocks.blocks) == 0
647
648
# Now one more block that should not have extra keys.
649
assert len(block_hashes) == 4
@@ -653,14 +654,14 @@ def test_cache_key_salting():
654
req1 = make_request("1", token_ids, cache_salt="salt1")
655
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
656
# Should match only a prefix of 3 blocks.
656
- assert len(computed_blocks) == 3
657
+ assert len(computed_blocks.blocks) == 3
658
assert num_computed_tokens == 3 * block_size
659
660
# Test cache miss with same content but different salt.
661
token_ids = common_token_ids + [4] * 11
662
req2 = make_request("2", token_ids, cache_salt="salt2")
663
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
663
- assert len(computed_blocks) == 0
664
+ assert len(computed_blocks.blocks) == 0
665
assert num_computed_tokens == 0
666
block_hashes = manager.req_to_block_hashes[req2.request_id]
667
assert len(block_hashes) == 3
@@ -685,7 +686,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
686
common_token_ids = [i for i in range(3) for _ in range(16)]
687
req0 = make_request("0", common_token_ids)
688
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
688
- assert not computed_blocks
689
+ assert not computed_blocks.blocks
690
assert num_computed_tokens == 0
691
manager.allocate_slots(req0, 48, computed_blocks)
692
block_part0 = manager.req_to_blocks[req0.request_id]
@@ -693,7 +694,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
694
# | Common-0 | Common-1 | Common-2 | Req1-3 | Req1-4 | Req1-5 | ... |
695
req1 = make_request("1", common_token_ids * 2)
696
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
696
- assert computed_blocks == block_part0
697
+ assert computed_blocks.blocks == block_part0
698
assert num_computed_tokens == 3 * 16
699
manager.allocate_slots(req1, 48, computed_blocks)
700
block_part1 = manager.req_to_blocks[req1.request_id]
@@ -707,7 +708,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
708
# | Req1-5(F)| Req2-0 | Req2-1 | ... |
709
req2 = make_request("2", [7] * block_size * 2)
710
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
710
- assert not computed_blocks
711
+ assert not computed_blocks.blocks
712
assert num_computed_tokens == 0
713
manager.allocate_slots(req2, block_size * 2, computed_blocks)
714
@@ -717,7 +718,7 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
718
assert manager.block_pool.free_block_queue.num_free_blocks == 5
719
req3 = make_request("3", common_token_ids * 3)
720
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req3)
720
- assert computed_blocks == block_part1
721
+ assert computed_blocks.blocks == block_part1
722
assert num_computed_tokens == 6 * 16
723
# Req3 cannot be allocated.
724
assert manager.allocate_slots(req3, 48, computed_blocks) is None
@@ -739,16 +740,16 @@ def test_reset_prefix_cache():
740
all_token_ids = full_block_token_ids + unique_token_ids
741
req0 = make_request("0", all_token_ids)
742
blocks = manager.allocate_slots(req0, 55)
742
- assert [b.block_id for b in blocks] == [1, 2, 3, 4]
743
+ assert blocks.get_block_ids() == [1, 2, 3, 4]
744
745
unique_token_ids = [4] * 7
746
all_token_ids = full_block_token_ids + unique_token_ids
747
req1 = make_request("1", all_token_ids)
748
computed_blocks, _ = manager.get_computed_blocks(req1)
749
assert len(manager.req_to_block_hashes[req1.request_id]) == 3
749
- assert len(computed_blocks) == 3
750
+ assert len(computed_blocks.blocks) == 3
751
blocks = manager.allocate_slots(req1, 7, computed_blocks)
751
- assert [b.block_id for b in blocks] == [5]
752
+ assert blocks.get_block_ids() == [5]
753
754
# Failed to reset prefix cache because some blocks are not freed yet.
755
assert not manager.reset_prefix_cache()
@@ -776,7 +777,7 @@ def test_prefix_cache_stats_disabled():
777
# Call all functions that check whether log_stats is disabled.
778
req = make_request("0", list(range(16)))
779
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)
779
- assert not computed_blocks
780
+ assert not computed_blocks.blocks
781
assert num_computed_tokens == 0
782
manager.allocate_slots(req, 16, computed_blocks)
783
manager.reset_prefix_cache()
@@ -866,7 +867,7 @@ def test_eagle_enabled_removes_last_block():
867
# Should retain 1 block:
868
# 1. Original 3 blocks → pop last hash → 2 matched blocks
869
# 2. drop last matched block → 1 remaining block
869
- assert len(computed_blocks) == 1
870
+ assert len(computed_blocks.blocks) == 1
871
assert num_tokens == 1 * block_size # 16 tokens
872
873
@@ -892,7 +893,7 @@ def test_eagle_with_partial_blocks():
893
req_eagle = make_request("partial_eagle", token_ids)
894
computed_blocks, num_tokens = manager.get_computed_blocks(req_eagle)
895
# Original match: 2 full blocks → Eagle removes 1 → 1 remaining
895
- assert len(computed_blocks) == 1
896
+ assert len(computed_blocks.blocks) == 1
897
assert num_tokens == 1 * block_size
898
899
@@ -934,7 +935,7 @@ def test_eagle_with_sliding_window():
935
req_eagle = make_request("partial_eagle", token_ids)
936
computed_blocks, num_tokens = manager.get_computed_blocks(req_eagle)
937
# Original match: 2 full blocks → Eagle removes 1 → 1 remaining
937
- assert len(computed_blocks) == 1
938
+ assert len(computed_blocks.blocks) == 1
939
assert num_tokens == 1 * block_size
940
941
# Evict the first block in the request
@@ -948,5 +949,5 @@ def test_eagle_with_sliding_window():
949
# Cache miss. The only hit prefix is [NULL_BLOCK, BLOCK_2] if eagle is
950
# not considered. But after dropping the last matched block due to eagle,
951
# there will be no matched prefix.
951
- assert len(computed_blocks) == 0
952
+ assert len(computed_blocks.blocks) == 0
953
assert num_tokens == 0
vllm/v1/core/kv_cache_manager.py
+44
-21
@@ -2,6 +2,7 @@
2
3
from collections import defaultdict
4
from collections.abc import Iterable
5
+from dataclasses import dataclass
6
from typing import Optional
7
8
from vllm.distributed.kv_events import KVCacheEvent
@@ -18,6 +19,24 @@ from vllm.v1.request import Request, RequestStatus
19
logger = init_logger(__name__)
20
21
22
+@dataclass
23
+class KVCacheBlocks:
24
+ blocks: list[KVCacheBlock]
25
+
26
+ def __add__(self, other: "KVCacheBlocks") -> "KVCacheBlocks":
27
+ """Adds two KVCacheBlocks instances."""
28
+ return KVCacheBlocks(self.blocks + other.blocks)
29
+
30
+ @classmethod
31
+ def create_empty(cls) -> "KVCacheBlocks":
32
+ """Creates a new KVCacheBlocks instance with no blocks."""
33
+ return cls([])
34
+
35
+ def get_block_ids(self) -> list[int]:
36
+ """Converts the KVCacheBlocks instance to a list of block IDs."""
37
+ return [block.block_id for block in self.blocks]
38
+
39
+
40
class KVCacheManager:
41
42
def __init__(
@@ -94,8 +113,8 @@ class KVCacheManager:
113
self.prefix_cache_stats = PrefixCacheStats()
114
return stats
115
97
- def get_computed_blocks(
98
- self, request: Request) -> tuple[list[KVCacheBlock], int]:
116
+ def get_computed_blocks(self,
117
+ request: Request) -> tuple[KVCacheBlocks, int]:
118
"""Get the computed (cached) blocks for the request.
119
Note that the computed blocks must be full.
120
@@ -109,7 +128,7 @@ class KVCacheManager:
128
"""
129
if not self.enable_caching:
130
# Prefix caching is disabled.
112
- return [], 0
131
+ return KVCacheBlocks.create_empty(), 0
132
133
# The block hashes for the request may already be computed
134
# if the scheduler has tried to schedule the request before.
@@ -124,7 +143,7 @@ class KVCacheManager:
143
self.prefix_cache_stats.requests += 1
144
# When the request requires prompt logprobs, we skip prefix caching.
145
if request.sampling_params.prompt_logprobs is not None:
127
- return [], 0
146
+ return KVCacheBlocks.create_empty(), 0
147
148
if len(block_hashes) * self.block_size == request.num_tokens:
149
# When prompt length is divisible by the block size and all
@@ -157,15 +176,15 @@ class KVCacheManager:
176
# sharing, `num_computed_tokens` is always a multiple of
177
# `block_size`.
178
num_computed_tokens = len(computed_blocks) * self.block_size
160
- return computed_blocks, num_computed_tokens
179
+ return KVCacheBlocks(computed_blocks), num_computed_tokens
180
181
def allocate_slots(
182
self,
183
request: Request,
184
num_tokens: int,
166
- new_computed_blocks: Optional[list[KVCacheBlock]] = None,
185
+ new_computed_blocks: Optional[KVCacheBlocks] = None,
186
num_lookahead_tokens: int = 0,
168
- ) -> Optional[list[KVCacheBlock]]:
187
+ ) -> Optional[KVCacheBlocks]:
188
"""Add slots for a request with new tokens to append.
189
190
Args:
@@ -173,7 +192,7 @@ class KVCacheManager:
192
num_tokens: The number of tokens to allocate, including external
193
tokens. Note that this does not include tokens that have
194
already been computed locally (i.e. new_computed_blocks).
176
- new_computed_blocks: A list of new computed blocks just hitting the
195
+ new_computed_blocks: The new computed blocks just hitting the
196
prefix caching.
197
num_lookahead_tokens: The number of speculative tokens to allocate.
198
This is used by spec decode proposers with kv-cache such
@@ -199,7 +218,10 @@ class KVCacheManager:
218
if num_tokens == 0:
219
raise ValueError("num_tokens must be greater than 0")
220
202
- new_computed_blocks = new_computed_blocks or []
221
+ if new_computed_blocks is not None:
222
+ new_computed_block_list = new_computed_blocks.blocks
223
+ else:
224
+ new_computed_block_list = []
225
226
req_blocks = self.req_to_blocks[request.request_id]
227
@@ -216,17 +238,18 @@ class KVCacheManager:
238
# The number of computed tokens is the number of computed tokens plus
239
# the new prefix caching hits
240
num_computed_tokens = (request.num_computed_tokens +
219
- len(new_computed_blocks) * self.block_size)
241
+ len(new_computed_block_list) * self.block_size)
242
num_required_blocks = cdiv(
243
num_computed_tokens + num_tokens + num_lookahead_tokens,
244
self.block_size)
245
num_new_blocks = (num_required_blocks - len(req_blocks) -
224
- len(new_computed_blocks))
246
+ len(new_computed_block_list))
247
248
# If a computed block of a request is an eviction candidate (in the
249
# free queue and ref_cnt == 0), it cannot be counted as a free block
250
# when allocating this request.
229
- num_evictable_computed_blocks = sum(1 for blk in new_computed_blocks
251
+ num_evictable_computed_blocks = sum(1
252
+ for blk in new_computed_block_list
253
if blk.ref_cnt == 0)
254
if (num_new_blocks > self.block_pool.get_num_free_blocks() -
255
num_evictable_computed_blocks):
@@ -235,15 +258,15 @@ class KVCacheManager:
258
259
# Touch the computed blocks to make sure they won't be evicted.
260
if self.enable_caching:
238
- self.block_pool.touch(new_computed_blocks)
261
+ self.block_pool.touch(new_computed_block_list)
262
else:
240
- assert not new_computed_blocks, (
263
+ assert not new_computed_block_list, (
264
"Computed blocks should be empty when "
265
"prefix caching is disabled")
266
267
# Append the new computed blocks to the request blocks until now to
268
# avoid the case where the new blocks cannot be allocated.
246
- req_blocks.extend(new_computed_blocks)
269
+ req_blocks.extend(new_computed_block_list)
270
271
# Start to handle new blocks
272
@@ -267,12 +290,12 @@ class KVCacheManager:
290
req_blocks.extend(new_blocks)
291
292
if not self.enable_caching:
270
- return new_blocks
293
+ return KVCacheBlocks(new_blocks)
294
272
- # Use `new_computed_blocks` for a new request, and `num_cached_block`
273
- # for a running request.
274
- num_cached_blocks = self.num_cached_block.get(request.request_id,
275
- len(new_computed_blocks))
295
+ # Use `new_computed_block_list` for a new request, and
296
+ # `num_cached_block` for a running request.
297
+ num_cached_blocks = self.num_cached_block.get(
298
+ request.request_id, len(new_computed_block_list))
299
# Speculated tokens might be rejected in the future, so we does
300
# not cache any speculated tokens. We only cache blocks with
301
# generated (accepted) tokens.
@@ -291,7 +314,7 @@ class KVCacheManager:
314
315
self.num_cached_block[
316
request.request_id] = num_full_blocks_after_append
294
- return new_blocks
317
+ return KVCacheBlocks(new_blocks)
318
319
def free(self, request: Request) -> None:
320
"""Free the blocks allocated for the request.
vllm/v1/core/sched/scheduler.py
+4
-6
@@ -261,9 +261,8 @@ class Scheduler(SchedulerInterface):
261
# Therefore, we might introduce some additional
262
# cycle to fill in the bitmask, which could be a big no-op.
263
structured_output_request_ids[request.request_id] = req_index
264
- req_to_new_block_ids[request.request_id] = [
265
- b.block_id for b in new_blocks
266
- ]
264
+ req_to_new_block_ids[request.request_id] = (
265
+ new_blocks.get_block_ids())
266
num_scheduled_tokens[request.request_id] = num_new_tokens
267
token_budget -= num_new_tokens
268
req_index += 1
@@ -407,9 +406,8 @@ class Scheduler(SchedulerInterface):
406
407
if self.lora_config and request.lora_request:
408
scheduled_loras.add(request.lora_request.lora_int_id)
410
- req_to_new_block_ids[request.request_id] = [
411
- b.block_id for b in computed_blocks + new_blocks
412
- ]
409
+ req_to_new_block_ids[request.request_id] = (
410
+ computed_blocks + new_blocks).get_block_ids()
411
num_scheduled_tokens[request.request_id] = num_new_tokens
412
token_budget -= num_new_tokens
413
request.status = RequestStatus.RUNNING