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, 07:15 UTC
[v1] Hybrid Memory Allocator (#17996)
Signed-off-by: Chen Zhang <zhangch99@outlook.com>
Chen Zhang committed
Jun 6, 2025 at 11:47 UTC
f8a1a2d108d290791ae1245b2ee309f38fdd7619
21 files changed
+1607
-442
tests/v1/core/test_kv_cache_utils.py
+210
-35
@@ -15,8 +15,8 @@ from vllm.v1.core.kv_cache_manager import KVCacheManager
15
from vllm.v1.core.kv_cache_utils import (
16
FreeKVCacheBlockQueue, KVCacheBlock, PrefixCachingMetrics,
17
estimate_max_model_len, generate_block_hash_extra_keys,
18
- get_max_concurrency_for_kv_cache_config, hash_block_tokens,
19
- hash_request_tokens, unify_kv_cache_configs)
18
+ get_kv_cache_config, get_max_concurrency_for_kv_cache_config,
19
+ hash_block_tokens, hash_request_tokens, unify_kv_cache_configs)
20
from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
21
KVCacheGroupSpec, KVCacheTensor,
22
SlidingWindowSpec)
@@ -63,6 +63,20 @@ def new_kv_cache_spec(block_size=16,
63
sliding_window=sliding_window)
64
65
66
+def new_sliding_window_spec(block_size=16,
67
+ num_kv_heads=2,
68
+ head_size=64,
69
+ dtype=torch.float32,
70
+ use_mla=False,
71
+ sliding_window=1):
72
+ return SlidingWindowSpec(block_size=block_size,
73
+ num_kv_heads=num_kv_heads,
74
+ head_size=head_size,
75
+ dtype=dtype,
76
+ use_mla=use_mla,
77
+ sliding_window=sliding_window)
78
+
79
+
80
def test_none_hash(monkeypatch):
81
import vllm.v1.core.kv_cache_utils
82
@@ -403,10 +417,10 @@ def test_unify_kv_cache_configs():
417
same_kv_cache_config = [
418
KVCacheConfig(
419
num_blocks=10,
406
- tensors={
407
- "layer1": KVCacheTensor(100),
408
- "layer2": KVCacheTensor(100),
409
- },
420
+ kv_cache_tensors=[
421
+ KVCacheTensor(size=100, shared_by=["layer1"]),
422
+ KVCacheTensor(size=100, shared_by=["layer2"]),
423
+ ],
424
kv_cache_groups=[
425
KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
426
KVCacheGroupSpec(["layer2"],
@@ -415,10 +429,10 @@ def test_unify_kv_cache_configs():
429
),
430
KVCacheConfig(
431
num_blocks=20,
418
- tensors={
419
- "layer1": KVCacheTensor(100),
420
- "layer2": KVCacheTensor(100),
421
- },
432
+ kv_cache_tensors=[
433
+ KVCacheTensor(size=100, shared_by=["layer1"]),
434
+ KVCacheTensor(size=100, shared_by=["layer2"]),
435
+ ],
436
kv_cache_groups=[
437
KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
438
KVCacheGroupSpec(["layer2"],
@@ -433,10 +447,10 @@ def test_unify_kv_cache_configs():
447
need_sort_kv_cache_config = [
448
KVCacheConfig(
449
num_blocks=10,
436
- tensors={
437
- "layer1": KVCacheTensor(100),
438
- "layer2": KVCacheTensor(100),
439
- },
450
+ kv_cache_tensors=[
451
+ KVCacheTensor(size=100, shared_by=["layer1"]),
452
+ KVCacheTensor(size=100, shared_by=["layer2"]),
453
+ ],
454
kv_cache_groups=[
455
KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
456
KVCacheGroupSpec(["layer2"],
@@ -445,10 +459,10 @@ def test_unify_kv_cache_configs():
459
),
460
KVCacheConfig(
461
num_blocks=20,
448
- tensors={
449
- "layer1": KVCacheTensor(100),
450
- "layer2": KVCacheTensor(100),
451
- },
462
+ kv_cache_tensors=[
463
+ KVCacheTensor(size=100, shared_by=["layer1"]),
464
+ KVCacheTensor(size=100, shared_by=["layer2"]),
465
+ ],
466
kv_cache_groups=[
467
KVCacheGroupSpec(["layer2"],
468
new_kv_cache_spec(num_kv_heads=4)),
@@ -464,10 +478,10 @@ def test_unify_kv_cache_configs():
478
diff_kv_cache_config = [
479
KVCacheConfig(
480
num_blocks=10,
467
- tensors={
468
- "layer1": KVCacheTensor(100),
469
- "layer2": KVCacheTensor(100),
470
- },
481
+ kv_cache_tensors=[
482
+ KVCacheTensor(size=100, shared_by=["layer1"]),
483
+ KVCacheTensor(size=100, shared_by=["layer2"]),
484
+ ],
485
kv_cache_groups=[
486
KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
487
KVCacheGroupSpec(["layer2"],
@@ -476,10 +490,10 @@ def test_unify_kv_cache_configs():
490
),
491
KVCacheConfig(
492
num_blocks=20,
479
- tensors={
480
- "layer1": KVCacheTensor(100),
481
- "layer2": KVCacheTensor(100),
482
- },
493
+ kv_cache_tensors=[
494
+ KVCacheTensor(size=100, shared_by=["layer1"]),
495
+ KVCacheTensor(size=100, shared_by=["layer2"]),
496
+ ],
497
kv_cache_groups=[
498
KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
499
KVCacheGroupSpec(["layer2"],
@@ -636,7 +650,7 @@ def test_get_max_concurrency_for_kv_cache_config():
650
651
kv_cache_config_full_attention = KVCacheConfig(
652
num_blocks=int(1024 * 1.5),
639
- tensors={},
653
+ kv_cache_tensors=[],
654
kv_cache_groups=[
655
KVCacheGroupSpec([f"layer_{i}" for i in range(32)],
656
full_attention_spec),
@@ -648,7 +662,7 @@ def test_get_max_concurrency_for_kv_cache_config():
662
663
kv_cache_config_sliding_window = KVCacheConfig(
664
num_blocks=129 * 3,
651
- tensors={},
665
+ kv_cache_tensors=[],
666
kv_cache_groups=[
667
KVCacheGroupSpec([f"layer_{i}" for i in range(32)],
668
sliding_window_spec),
@@ -660,7 +674,7 @@ def test_get_max_concurrency_for_kv_cache_config():
674
675
kv_cache_config_hybrid_model = KVCacheConfig(
676
num_blocks=(1024 + 129) * 3,
663
- tensors={},
677
+ kv_cache_tensors=[],
678
kv_cache_groups=[
679
KVCacheGroupSpec([f"layer_{i}" for i in range(32)],
680
full_attention_spec),
@@ -678,9 +692,9 @@ def test_allocate_with_lookahead():
692
block_size = 4
693
config = KVCacheConfig(
694
num_blocks=10,
681
- tensors={
682
- "layer1": KVCacheTensor(100),
683
- },
695
+ kv_cache_tensors=[
696
+ KVCacheTensor(size=100, shared_by=["layer1"]),
697
+ ],
698
kv_cache_groups=[
699
KVCacheGroupSpec(["layer1"],
700
new_kv_cache_spec(block_size=block_size)),
@@ -702,7 +716,7 @@ def test_allocate_with_lookahead():
716
num_new_tokens=3,
717
num_lookahead_tokens=2, # Total required: 3+2=5 tokens
718
)
705
- assert len(blocks.blocks) == 2 # ceil(5/4)=2 blocks
719
+ assert len(blocks.get_block_ids()[0]) == 2 # ceil(5/4)=2 blocks
720
721
# Test case 2: With precomputed blocks
722
kv_cache_manager = KVCacheManager(kv_cache_config=config,
@@ -713,7 +727,7 @@ def test_allocate_with_lookahead():
727
num_new_tokens=3,
728
num_lookahead_tokens=2,
729
)
716
- assert len(blocks.blocks) == 2
730
+ assert len(blocks.get_block_ids()[0]) == 2
731
732
# Test case 3: With precomputed blocks
733
# required_blocks = ceil((3 + 4) / 4) = 2
@@ -724,4 +738,165 @@ def test_allocate_with_lookahead():
738
num_new_tokens=3,
739
num_lookahead_tokens=4,
740
)
727
- assert len(blocks.blocks) == 2
741
+ assert len(blocks.get_block_ids()[0]) == 2
742
+
743
+
744
+def test_get_kv_cache_config():
745
+ # pass max_model_len to pass check_enough_kv_cache_memory
746
+ model_config = ModelConfig(max_model_len=16)
747
+ vllm_config = VllmConfig(model_config=model_config)
748
+
749
+ mem_per_block_per_layer = 16 * 2 * 64 * 4 * 2
750
+ # all layers are full attention -> single group
751
+ kv_cache_specs_full = {
752
+ 'layer_1': new_kv_cache_spec(),
753
+ 'layer_2': new_kv_cache_spec(),
754
+ }
755
+ kv_cache_config_full = get_kv_cache_config(
756
+ vllm_config, kv_cache_specs_full, mem_per_block_per_layer * 2 * 32)
757
+ assert kv_cache_config_full == KVCacheConfig(
758
+ num_blocks=32,
759
+ kv_cache_tensors=[
760
+ KVCacheTensor(size=mem_per_block_per_layer * 32,
761
+ shared_by=["layer_1"]),
762
+ KVCacheTensor(size=mem_per_block_per_layer * 32,
763
+ shared_by=["layer_2"]),
764
+ ],
765
+ kv_cache_groups=[
766
+ KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())
767
+ ])
768
+
769
+ # all layers are sliding window -> single group
770
+ kv_cache_specs_sliding = {
771
+ 'layer_1': new_sliding_window_spec(),
772
+ 'layer_2': new_sliding_window_spec(),
773
+ }
774
+ kv_cache_config_sliding = get_kv_cache_config(
775
+ vllm_config, kv_cache_specs_sliding, mem_per_block_per_layer * 2 * 32)
776
+ assert kv_cache_config_sliding == KVCacheConfig(
777
+ num_blocks=32,
778
+ kv_cache_tensors=[
779
+ KVCacheTensor(size=mem_per_block_per_layer * 32,
780
+ shared_by=["layer_1"]),
781
+ KVCacheTensor(size=mem_per_block_per_layer * 32,
782
+ shared_by=["layer_2"]),
783
+ ],
784
+ kv_cache_groups=[
785
+ KVCacheGroupSpec(["layer_1", "layer_2"], new_sliding_window_spec())
786
+ ])
787
+
788
+ # full + sliding, but disable_hybrid_kv_cache_manager
789
+ vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = True
790
+ kv_cache_specs_hybrid = {
791
+ 'layer_1': new_kv_cache_spec(),
792
+ 'layer_2': new_sliding_window_spec(),
793
+ }
794
+ kv_cache_config_hybrid = get_kv_cache_config(
795
+ vllm_config, kv_cache_specs_hybrid, mem_per_block_per_layer * 2 * 32)
796
+ assert kv_cache_config_hybrid == KVCacheConfig(
797
+ num_blocks=32,
798
+ kv_cache_tensors=[
799
+ KVCacheTensor(size=mem_per_block_per_layer * 32,
800
+ shared_by=["layer_1"]),
801
+ KVCacheTensor(size=mem_per_block_per_layer * 32,
802
+ shared_by=["layer_2"]),
803
+ ],
804
+ kv_cache_groups=[
805
+ KVCacheGroupSpec(["layer_1", "layer_2"],
806
+ new_kv_cache_spec(sliding_window=1)),
807
+ ],
808
+ )
809
+ vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False
810
+
811
+ # full + sliding, with hybrid_kv_cache_manager
812
+ kv_cache_specs_hybrid = {
813
+ 'layer_1': new_kv_cache_spec(),
814
+ 'layer_2': new_sliding_window_spec(),
815
+ }
816
+ kv_cache_config_hybrid = get_kv_cache_config(
817
+ vllm_config, kv_cache_specs_hybrid, mem_per_block_per_layer * 2 * 32)
818
+ assert kv_cache_config_hybrid == KVCacheConfig(
819
+ num_blocks=64,
820
+ kv_cache_tensors=[
821
+ KVCacheTensor(size=mem_per_block_per_layer * 64,
822
+ shared_by=["layer_1", "layer_2"]),
823
+ ],
824
+ kv_cache_groups=[
825
+ KVCacheGroupSpec(["layer_1"], new_kv_cache_spec()),
826
+ KVCacheGroupSpec(["layer_2"], new_sliding_window_spec()),
827
+ ],
828
+ )
829
+
830
+ # 2 full + 4 sliding, 2 layers per group
831
+ kv_cache_specs_hybrid = {
832
+ 'layer_1': new_kv_cache_spec(),
833
+ 'layer_2': new_kv_cache_spec(),
834
+ 'layer_3': new_sliding_window_spec(),
835
+ 'layer_4': new_sliding_window_spec(),
836
+ 'layer_5': new_sliding_window_spec(),
837
+ 'layer_6': new_sliding_window_spec(),
838
+ }
839
+ kv_cache_config_hybrid = get_kv_cache_config(
840
+ vllm_config, kv_cache_specs_hybrid, mem_per_block_per_layer * 2 * 32)
841
+ assert kv_cache_config_hybrid == KVCacheConfig(
842
+ num_blocks=32,
843
+ kv_cache_tensors=[
844
+ KVCacheTensor(size=mem_per_block_per_layer * 32,
845
+ shared_by=["layer_1", "layer_3", "layer_5"]),
846
+ KVCacheTensor(size=mem_per_block_per_layer * 32,
847
+ shared_by=["layer_2", "layer_4", "layer_6"]),
848
+ ],
849
+ kv_cache_groups=[
850
+ KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec()),
851
+ KVCacheGroupSpec(["layer_3", "layer_4"],
852
+ new_sliding_window_spec()),
853
+ KVCacheGroupSpec(["layer_5", "layer_6"],
854
+ new_sliding_window_spec()),
855
+ ],
856
+ )
857
+
858
+ # 3 full + 7 sliding, pad to 3 full + 9 sliding
859
+ kv_cache_specs_hybrid = {
860
+ 'layer_1': new_kv_cache_spec(),
861
+ 'layer_2': new_kv_cache_spec(),
862
+ 'layer_3': new_kv_cache_spec(),
863
+ 'layer_4': new_sliding_window_spec(),
864
+ 'layer_5': new_sliding_window_spec(),
865
+ 'layer_6': new_sliding_window_spec(),
866
+ 'layer_7': new_sliding_window_spec(),
867
+ 'layer_8': new_sliding_window_spec(),
868
+ 'layer_9': new_sliding_window_spec(),
869
+ 'layer_10': new_sliding_window_spec(),
870
+ }
871
+ kv_cache_config_hybrid = get_kv_cache_config(
872
+ vllm_config, kv_cache_specs_hybrid, mem_per_block_per_layer * 3 * 32)
873
+ assert kv_cache_config_hybrid == KVCacheConfig(
874
+ num_blocks=32,
875
+ kv_cache_tensors=[
876
+ KVCacheTensor(
877
+ size=mem_per_block_per_layer * 32,
878
+ shared_by=["layer_1", "layer_4", "layer_7", "layer_10"]),
879
+ KVCacheTensor(size=mem_per_block_per_layer * 32,
880
+ shared_by=["layer_2", "layer_5", "layer_8"]),
881
+ KVCacheTensor(size=mem_per_block_per_layer * 32,
882
+ shared_by=["layer_3", "layer_6", "layer_9"]),
883
+ ],
884
+ kv_cache_groups=[
885
+ KVCacheGroupSpec(["layer_1", "layer_2", "layer_3"],
886
+ new_kv_cache_spec()),
887
+ KVCacheGroupSpec(["layer_4", "layer_5", "layer_6"],
888
+ new_sliding_window_spec()),
889
+ KVCacheGroupSpec(["layer_7", "layer_8", "layer_9"],
890
+ new_sliding_window_spec()),
891
+ KVCacheGroupSpec(["layer_10"], new_sliding_window_spec()),
892
+ ],
893
+ )
894
+
895
+ # different hidden size, unimplemented
896
+ kv_cache_specs_hybrid = {
897
+ 'layer_1': new_kv_cache_spec(head_size=128),
898
+ 'layer_2': new_kv_cache_spec(),
899
+ }
900
+ with pytest.raises(NotImplementedError):
901
+ get_kv_cache_config(vllm_config, kv_cache_specs_hybrid,
902
+ mem_per_block_per_layer * 2 * 32)
tests/v1/core/test_prefix_caching.py
+350
-100
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Compare the with and without prefix caching."""
4
5
+import copy
6
from typing import Optional
7
8
import pytest
@@ -13,8 +14,8 @@ from vllm.sampling_params import SamplingParams
14
from vllm.utils import sha256
15
from vllm.v1.core.block_pool import BlockPool
16
from vllm.v1.core.kv_cache_manager import KVCacheManager, Request
16
-from vllm.v1.core.kv_cache_utils import (BlockHash, KVCacheBlock,
17
- hash_block_tokens)
17
+from vllm.v1.core.kv_cache_utils import (BlockHash, BlockHashWithGroupId,
18
+ KVCacheBlock, hash_block_tokens)
19
from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
20
KVCacheGroupSpec, SlidingWindowSpec)
21
@@ -47,7 +48,7 @@ def make_request(request_id,
48
def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig:
49
return KVCacheConfig(
50
num_blocks=num_blocks,
50
- tensors={},
51
+ kv_cache_tensors=[],
52
kv_cache_groups=[
53
KVCacheGroupSpec(
54
["layer"],
@@ -57,6 +58,38 @@ def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig:
58
)
59
60
61
+def make_kv_cache_config_hybrid_model(block_size: int,
62
+ num_blocks: int) -> KVCacheConfig:
63
+ return KVCacheConfig(
64
+ num_blocks=num_blocks,
65
+ kv_cache_tensors=[],
66
+ kv_cache_groups=[
67
+ KVCacheGroupSpec(
68
+ ["layer1"],
69
+ FullAttentionSpec(block_size, 1, 1, torch.float32, False),
70
+ ),
71
+ KVCacheGroupSpec(
72
+ ["layer2"],
73
+ SlidingWindowSpec(block_size,
74
+ 1,
75
+ 1,
76
+ torch.float32,
77
+ False,
78
+ sliding_window=2 * block_size),
79
+ ),
80
+ KVCacheGroupSpec(
81
+ ["layer3"],
82
+ SlidingWindowSpec(block_size,
83
+ 1,
84
+ 1,
85
+ torch.float32,
86
+ False,
87
+ sliding_window=2 * block_size),
88
+ ),
89
+ ],
90
+ )
91
+
92
+
93
@pytest.mark.parametrize("hash_algo", ["sha256", "hash"])
94
def test_prefill(hash_algo):
95
manager = KVCacheManager(
@@ -79,10 +112,10 @@ def test_prefill(hash_algo):
112
req0 = make_request("0", all_token_ids)
113
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
114
assert len(manager.req_to_block_hashes[req0.request_id]) == 3
82
- assert not computed_blocks.blocks
115
+ assert not computed_blocks.blocks[0]
116
assert num_computed_tokens == 0
117
blocks = manager.allocate_slots(req0, 55,
85
- len(computed_blocks.blocks) * 16,
118
+ len(computed_blocks.blocks[0]) * 16,
119
computed_blocks)
120
assert blocks.get_block_ids() == [[1, 2, 3, 4]]
121
@@ -92,7 +125,8 @@ def test_prefill(hash_algo):
125
block_tokens = tuple(all_token_ids[(block_id - 1) * 16:block_id * 16])
126
block_hash = hash_block_tokens(hash_fn, parent_block_hash,
127
block_tokens)
95
- assert manager.block_pool.blocks[block_id].block_hash == block_hash
128
+ assert manager.block_pool.blocks[
129
+ block_id].block_hash.block_hash == block_hash
130
assert manager.block_pool.blocks[block_id].ref_cnt == 1
131
parent_block_hash = block_hash.hash_value
132
@@ -111,10 +145,10 @@ def test_prefill(hash_algo):
145
assert num_computed_tokens == 3 * 16
146
num_new_tokens = 53 - 3 * 16
147
blocks = manager.allocate_slots(req1, num_new_tokens,
114
- len(computed_blocks.blocks) * 16,
148
+ len(computed_blocks.blocks[0]) * 16,
149
computed_blocks)
150
assert blocks.get_block_ids() == [[5]]
117
- for block in computed_blocks.blocks:
151
+ for block in computed_blocks.blocks[0]:
152
assert block.ref_cnt == 2
153
154
# At this point, we should have 5 free blocks left.
@@ -145,7 +179,7 @@ def test_prefill(hash_algo):
179
assert num_computed_tokens == 3 * 16
180
num_new_tokens = 53 - 3 * 16
181
blocks = manager.allocate_slots(req2, num_new_tokens,
148
- len(computed_blocks.blocks) * 16,
182
+ len(computed_blocks.blocks[0]) * 16,
183
computed_blocks)
184
assert blocks.get_block_ids() == [[6]]
185
@@ -165,10 +199,10 @@ def test_prefill(hash_algo):
199
# Cache miss and eviction.
200
req3 = make_request("3", [99] * (16 * 10))
201
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req3)
168
- assert not computed_blocks.blocks
202
+ assert not computed_blocks.blocks[0]
203
assert num_computed_tokens == 0
204
blocks = manager.allocate_slots(req3, 16 * 10,
171
- len(computed_blocks.blocks) * 16,
205
+ len(computed_blocks.blocks[0]) * 16,
206
computed_blocks)
207
# This block ID order also checks the eviction order.
208
assert blocks.get_block_ids() == [[7, 8, 9, 10, 4, 5, 6, 3, 2, 1]]
@@ -177,6 +211,138 @@ def test_prefill(hash_algo):
211
assert manager.block_pool.free_block_queue.free_list_tail is None
212
213
214
+def test_prefill_hybrid_model():
215
+ block_size = 16
216
+ manager = KVCacheManager(
217
+ make_kv_cache_config_hybrid_model(block_size, 21),
218
+ max_model_len=8192,
219
+ enable_caching=True,
220
+ )
221
+
222
+ hash_fn = hash
223
+
224
+ # Complete 3 blocks (48 tokens)
225
+ common_token_ids = [i for i in range(3) for _ in range(block_size)]
226
+
227
+ # Fully cache miss
228
+ # Incomplete 1 block (7 tokens)
229
+ unique_token_ids = [3] * 7
230
+ all_token_ids = common_token_ids + unique_token_ids
231
+ req0 = make_request("0", all_token_ids)
232
+ computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
233
+ assert len(manager.req_to_block_hashes[req0.request_id]) == 3
234
+ assert not computed_blocks.blocks[0]
235
+ assert num_computed_tokens == 0
236
+ blocks = manager.allocate_slots(req0, 55,
237
+ len(computed_blocks.blocks[0]) * 16,
238
+ computed_blocks)
239
+ assert blocks.get_block_ids() == [[1, 2, 3, 4], [5, 6, 7, 8],
240
+ [9, 10, 11, 12]]
241
+
242
+ # Check full block metadata
243
+ parent_block_hash = None
244
+ for length, block_ids in zip((1, 2, 3),
245
+ ((1, 5, 9), (2, 6, 10), (3, 7, 11))):
246
+ block_tokens = tuple(all_token_ids[(length - 1) * 16:length * 16])
247
+ block_hash = hash_block_tokens(hash_fn, parent_block_hash,
248
+ block_tokens)
249
+ for block_id in block_ids:
250
+ assert manager.block_pool.blocks[
251
+ block_id].block_hash.block_hash == block_hash
252
+ assert manager.block_pool.blocks[block_id].ref_cnt == 1
253
+ parent_block_hash = block_hash.hash_value
254
+
255
+ # Check partial block metadata
256
+ for block_id in (4, 8, 12):
257
+ assert manager.block_pool.blocks[block_id].block_hash is None
258
+ assert manager.block_pool.blocks[block_id].ref_cnt == 1
259
+
260
+ # Cache hit in the common prefix
261
+ # Incomplete 1 block (5 tokens)
262
+ unique_token_ids = [3] * 5
263
+ req1 = make_request("1", common_token_ids + unique_token_ids)
264
+ computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
265
+ assert len(manager.req_to_block_hashes[req1.request_id]) == 3
266
+ assert computed_blocks.get_block_ids() == [[1, 2, 3], [0, 6, 7],
267
+ [0, 10, 11]]
268
+ assert num_computed_tokens == 3 * 16
269
+ num_new_tokens = 53 - 3 * 16
270
+ blocks = manager.allocate_slots(req1, num_new_tokens,
271
+ len(computed_blocks.blocks[0]) * 16,
272
+ computed_blocks)
273
+ assert blocks.get_block_ids() == [[13], [14], [15]]
274
+ for block_per_group in computed_blocks.blocks:
275
+ for block in block_per_group:
276
+ if block != manager.block_pool.null_block:
277
+ assert block.ref_cnt == 2
278
+
279
+ block_hashes = manager.req_to_block_hashes[req1.request_id]
280
+ manager.free(req0)
281
+ manager.free(req1)
282
+
283
+ cached_block_hash_to_block_bak = copy.copy(
284
+ manager.block_pool.cached_block_hash_to_block)
285
+
286
+ def test_partial_request_hit(request_id: str,
287
+ hash_to_evict: list[BlockHashWithGroupId],
288
+ expect_hit_length: int):
289
+ req = make_request(request_id, common_token_ids + unique_token_ids)
290
+ for hash_with_group_id in hash_to_evict:
291
+ manager.block_pool.cached_block_hash_to_block.pop(
292
+ hash_with_group_id)
293
+ computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)
294
+ assert len(manager.req_to_block_hashes[req.request_id]) == 3
295
+ assert num_computed_tokens == expect_hit_length * block_size
296
+ for block_per_group in computed_blocks.blocks:
297
+ assert len(block_per_group) == num_computed_tokens // block_size
298
+ for hash_with_group_id in hash_to_evict:
299
+ manager.block_pool.cached_block_hash_to_block[
300
+ hash_with_group_id] = cached_block_hash_to_block_bak[
301
+ hash_with_group_id]
302
+ manager.free(req)
303
+
304
+ # Evict the blocks outside sliding window, does not affect the hit length.
305
+ test_partial_request_hit("2", [
306
+ BlockHashWithGroupId(block_hashes[0], 1),
307
+ BlockHashWithGroupId(block_hashes[0], 2)
308
+ ], 3)
309
+
310
+ # Evict the first block of full attention, makes total cache miss.
311
+ test_partial_request_hit("3", [
312
+ BlockHashWithGroupId(block_hashes[0], 0),
313
+ ], 0)
314
+
315
+ # Evict the last block of all layers, reduces the hit length to 2.
316
+ test_partial_request_hit("4", [
317
+ BlockHashWithGroupId(block_hashes[2], 0),
318
+ BlockHashWithGroupId(block_hashes[2], 1),
319
+ BlockHashWithGroupId(block_hashes[2], 2),
320
+ ], 2)
321
+
322
+ # Evict the last block of full attention, reduces the hit length to 2.
323
+ test_partial_request_hit("5", [BlockHashWithGroupId(block_hashes[2], 0)],
324
+ 2)
325
+
326
+ # Evict the last block of sliding window, reduces the hit length to 2.
327
+ test_partial_request_hit("6", [BlockHashWithGroupId(block_hashes[2], 1)],
328
+ 2)
329
+
330
+ # Evict the last block of sliding window, reduces the hit length to 2.
331
+ test_partial_request_hit("7", [BlockHashWithGroupId(block_hashes[2], 2)],
332
+ 2)
333
+
334
+ # Evict different set of blocks for full attention and sliding window makes
335
+ # total cache miss.
336
+ # The cache hit length of full attention is 1 * block_size.
337
+ # The cache hit length of sliding window is 2 * block_size.
338
+ # Then it is cache miss as the two type of layers have different hit length.
339
+ test_partial_request_hit("8", [
340
+ BlockHashWithGroupId(block_hashes[2], 0),
341
+ BlockHashWithGroupId(block_hashes[0], 1),
342
+ BlockHashWithGroupId(block_hashes[0], 2),
343
+ ], 0)
344
+
345
+
346
def test_prefill_plp():
347
'''Test prefill with APC and some prompt logprobs (plp) requests.
348
@@ -203,13 +369,13 @@ def test_prefill_plp():
369
req0 = make_request("0", all_token_ids, prompt_logprobs=5)
370
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
371
assert len(manager.req_to_block_hashes[req0.request_id]) == 0
206
- assert not computed_blocks.blocks
372
+ assert not computed_blocks.blocks[0]
373
assert num_computed_tokens == 0
374
blocks = manager.allocate_slots(req0, 55,
209
- len(computed_blocks.blocks) * 16,
375
+ len(computed_blocks.blocks[0]) * 16,
376
computed_blocks)
377
assert blocks.get_block_ids() == [[1, 2, 3, 4]]
212
- req0_block_hashes = [b.block_hash for b in blocks.blocks]
378
+ req0_block_hashes = [b.block_hash for b in blocks.blocks[0]]
379
380
# Check full block metadata
381
parent_block_hash = None
@@ -217,7 +383,8 @@ def test_prefill_plp():
383
block_tokens = tuple(all_token_ids[(block_id - 1) * 16:block_id * 16])
384
block_hash = hash_block_tokens(hash_fn, parent_block_hash,
385
block_tokens)
220
- assert manager.block_pool.blocks[block_id].block_hash == block_hash
386
+ assert manager.block_pool.blocks[
387
+ block_id].block_hash.block_hash == block_hash
388
assert manager.block_pool.blocks[block_id].ref_cnt == 1
389
parent_block_hash = block_hash.hash_value
390
@@ -237,10 +404,10 @@ def test_prefill_plp():
404
assert num_computed_tokens == 3 * 16
405
num_new_tokens = 53 - 3 * 16
406
blocks = manager.allocate_slots(req1, num_new_tokens,
240
- len(computed_blocks.blocks) * 16,
407
+ len(computed_blocks.blocks[0]) * 16,
408
computed_blocks)
409
assert blocks.get_block_ids() == [[5]]
243
- for block in computed_blocks.blocks:
410
+ for block in computed_blocks.blocks[0]:
411
assert block.ref_cnt == 2
412
413
# At this point, we should have 5 free blocks left.
@@ -269,14 +436,14 @@ def test_prefill_plp():
436
prompt_logprobs=5)
437
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
438
assert len(manager.req_to_block_hashes[req2.request_id]) == 0
272
- assert not computed_blocks.blocks
439
+ assert not computed_blocks.blocks[0]
440
assert num_computed_tokens == 0
441
blocks = manager.allocate_slots(req2, 55,
275
- len(computed_blocks.blocks) * 16,
442
+ len(computed_blocks.blocks[0]) * 16,
443
computed_blocks)
444
block_ids = blocks.get_block_ids()
445
# Duplicate cached blocks have different ids but same hashes vs request #0
279
- assert [b.block_hash for b in blocks.blocks] == req0_block_hashes
446
+ assert [b.block_hash for b in blocks.blocks[0]] == req0_block_hashes
447
assert block_ids != [[1, 2, 3, 4]]
448
449
# Request #2 block hashes are valid since request #0 hashes are.
@@ -302,10 +469,10 @@ def test_decode():
469
unique_token_ids = [3] * 7
470
req0 = make_request("0", common_token_ids + unique_token_ids)
471
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
305
- assert not computed_blocks.blocks
472
+ assert not computed_blocks.blocks[0]
473
assert num_computed_tokens == 0
474
blocks = manager.allocate_slots(req0, 55,
308
- len(computed_blocks.blocks) * 16,
475
+ len(computed_blocks.blocks[0]) * 16,
476
computed_blocks)
477
assert blocks.get_block_ids() == [[1, 2, 3, 4]]
478
@@ -314,10 +481,10 @@ def test_decode():
481
for _ in range(4):
482
req0.append_output_token_ids(8)
483
new_blocks = manager.allocate_slots(req0, 4,
317
- len(computed_blocks.blocks) * 16,
484
+ len(computed_blocks.blocks[0]) * 16,
485
computed_blocks)
319
- assert new_blocks is not None and len(new_blocks.blocks) == 0
320
- assert manager.single_type_manager.req_to_blocks[
486
+ assert new_blocks is not None and len(new_blocks.blocks[0]) == 0
487
+ assert manager.coordinator.single_type_managers[0].req_to_blocks[
488
req0.request_id][-1].block_hash is None
489
490
# Append slots with allocating a new block.
@@ -327,12 +494,12 @@ def test_decode():
494
for _ in range(9 + 10):
495
req0.append_output_token_ids(7)
496
new_blocks = manager.allocate_slots(req0, 19,
330
- len(computed_blocks.blocks) * 16,
497
+ len(computed_blocks.blocks[0]) * 16,
498
computed_blocks)
332
- assert new_blocks is not None and len(new_blocks.blocks) == 1
333
- assert manager.single_type_manager.req_to_blocks[
499
+ assert new_blocks is not None and len(new_blocks.blocks[0]) == 1
500
+ assert manager.coordinator.single_type_managers[0].req_to_blocks[
501
req0.request_id][-2].block_hash is not None
335
- assert manager.single_type_manager.req_to_blocks[
502
+ assert manager.coordinator.single_type_managers[0].req_to_blocks[
503
req0.request_id][-1].block_hash is None
504
505
@@ -346,23 +513,23 @@ def test_evict():
513
last_token_id = 5 * 16 + 7
514
req0 = make_request("0", list(range(last_token_id)))
515
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
349
- assert not computed_blocks.blocks
516
+ assert not computed_blocks.blocks[0]
517
assert num_computed_tokens == 0
518
blocks = manager.allocate_slots(req0, 5 * 16 + 7,
352
- len(computed_blocks.blocks) * 16,
519
+ len(computed_blocks.blocks[0]) * 16,
520
computed_blocks)
354
- assert len(blocks.blocks) == 6 # 5 full + 1 partial
521
+ assert len(blocks.blocks[0]) == 6 # 5 full + 1 partial
522
523
# 3 blocks.
524
req1 = make_request("1", list(range(last_token_id,
525
last_token_id + 3 * 16)))
526
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
360
- assert not computed_blocks.blocks
527
+ assert not computed_blocks.blocks[0]
528
assert num_computed_tokens == 0
529
blocks = manager.allocate_slots(req1, 3 * 16,
363
- len(computed_blocks.blocks) * 16,
530
+ len(computed_blocks.blocks[0]) * 16,
531
computed_blocks)
365
- assert len(blocks.blocks) == 3 # 3 full blocks
532
+ assert len(blocks.blocks[0]) == 3 # 3 full blocks
533
last_token_id += 3 * 16
534
535
# 10 - (6 + 3) == 1
@@ -382,7 +549,7 @@ def test_evict():
549
assert computed_blocks.get_block_ids() == [[1, 2]]
550
assert num_computed_tokens == 2 * 16
551
blocks = manager.allocate_slots(req2, 3,
385
- len(computed_blocks.blocks) * 16,
552
+ len(computed_blocks.blocks[0]) * 16,
553
computed_blocks)
554
assert blocks.get_block_ids() == [[10]]
555
assert manager.block_pool.free_block_queue.num_free_blocks == 7
@@ -404,12 +571,12 @@ def test_hash_block_correct_reuse():
571
num_tokens = block_size * 1
572
req = make_request("0", list(range(num_tokens)))
573
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)
407
- assert not computed_blocks.blocks
574
+ assert not computed_blocks.blocks[0]
575
assert num_computed_tokens == 0
576
blocks = manager.allocate_slots(req, num_tokens,
410
- len(computed_blocks.blocks) * 16,
577
+ len(computed_blocks.blocks[0]) * 16,
578
computed_blocks)
412
- assert len(blocks.blocks) == 1
579
+ assert len(blocks.blocks[0]) == 1
580
581
# Deallocate the block.
582
manager.free(req)
@@ -418,15 +585,15 @@ def test_hash_block_correct_reuse():
585
# block is cleared.
586
req = make_request("1", list(range(num_tokens - 1)))
587
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)
421
- assert not computed_blocks.blocks
588
+ assert not computed_blocks.blocks[0]
589
assert num_computed_tokens == 0
590
blocks = manager.allocate_slots(req, num_tokens - 1,
424
- len(computed_blocks.blocks) * 16,
591
+ len(computed_blocks.blocks[0]) * 16,
592
computed_blocks)
426
- assert len(blocks.blocks) == 1
593
+ assert len(blocks.blocks[0]) == 1
594
428
- assert manager.block_pool.blocks[
429
- blocks.blocks[0].block_id].block_hash is None
595
+ assert manager.block_pool.blocks[blocks.blocks[0]
596
+ [0].block_id].block_hash is None
597
598
599
def test_computed_blocks_not_evicted():
@@ -445,24 +612,24 @@ def test_computed_blocks_not_evicted():
612
num_tokens = block_size * 1
613
req0 = make_request("0", list(range(num_tokens)))
614
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
448
- assert not computed_blocks.blocks
615
+ assert not computed_blocks.blocks[0]
616
assert num_computed_tokens == 0
617
blocks = manager.allocate_slots(req0, num_tokens,
451
- len(computed_blocks.blocks) * 16,
618
+ len(computed_blocks.blocks[0]) * 16,
619
computed_blocks)
453
- assert len(blocks.blocks) == 1
454
- assert blocks.blocks[0].block_id == 1
620
+ assert len(blocks.blocks[0]) == 1
621
+ assert blocks.blocks[0][0].block_id == 1
622
623
# Allocate another block.
624
req1 = make_request("1", list(range(num_tokens, num_tokens * 2)))
625
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
459
- assert not computed_blocks.blocks
626
+ assert not computed_blocks.blocks[0]
627
assert num_computed_tokens == 0
628
blocks = manager.allocate_slots(req1, num_tokens,
462
- len(computed_blocks.blocks) * 16,
629
+ len(computed_blocks.blocks[0]) * 16,
630
computed_blocks)
464
- assert len(blocks.blocks) == 1
465
- assert blocks.blocks[0].block_id == 2
631
+ assert len(blocks.blocks[0]) == 1
632
+ assert blocks.blocks[0][0].block_id == 2
633
634
# Free the blocks.
635
manager.free(req0)
@@ -472,15 +639,15 @@ def test_computed_blocks_not_evicted():
639
# cached block rather than the first one.
640
req2 = make_request("2", list(range(num_tokens * 2)))
641
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
475
- assert len(computed_blocks.blocks) == 1
476
- assert computed_blocks.blocks[0].block_id == 1
642
+ assert len(computed_blocks.blocks[0]) == 1
643
+ assert computed_blocks.blocks[0][0].block_id == 1
644
assert num_computed_tokens == block_size
645
646
blocks = manager.allocate_slots(req2, num_tokens * 2 - num_tokens,
480
- len(computed_blocks.blocks) * 16,
647
+ len(computed_blocks.blocks[0]) * 16,
648
computed_blocks)
482
- assert len(blocks.blocks) == 1
483
- assert blocks.blocks[0].block_id == 2
649
+ assert len(blocks.blocks[0]) == 1
650
+ assert blocks.blocks[0][0].block_id == 2
651
652
653
def test_basic_prefix_caching_disabled():
@@ -497,12 +664,12 @@ def test_basic_prefix_caching_disabled():
664
req1 = make_request("1", list(range(10))) # 2 blocks and some more
665
666
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
500
- assert not computed_blocks.blocks
667
+ assert not computed_blocks.blocks[0]
668
assert num_computed_tokens == 0
669
blocks = manager.allocate_slots(req1, 10,
503
- len(computed_blocks.blocks) * 16,
670
+ len(computed_blocks.blocks[0]) * 16,
671
computed_blocks)
505
- assert len(blocks.blocks) == 3
672
+ assert len(blocks.blocks[0]) == 3
673
674
# Free the blocks.
675
manager.free(req1)
@@ -510,20 +677,20 @@ def test_basic_prefix_caching_disabled():
677
# No caching.
678
req2 = make_request("2", list(range(16))) # shared prefix
679
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
513
- assert not computed_blocks.blocks
680
+ assert not computed_blocks.blocks[0]
681
assert num_computed_tokens == 0
682
blocks = manager.allocate_slots(req2, 16,
516
- len(computed_blocks.blocks) * 16,
683
+ len(computed_blocks.blocks[0]) * 16,
684
computed_blocks)
518
- assert len(blocks.blocks) == 4
685
+ assert len(blocks.blocks[0]) == 4
686
687
# New requests should not have any blocks.
688
req3 = make_request("3", list(range(4)))
689
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req3)
523
- assert not computed_blocks.blocks
690
+ assert not computed_blocks.blocks[0]
691
assert num_computed_tokens == 0
692
blocks = manager.allocate_slots(req3, 4,
526
- len(computed_blocks.blocks) * 16,
693
+ len(computed_blocks.blocks[0]) * 16,
694
computed_blocks)
695
assert not blocks
696
@@ -558,6 +725,7 @@ def test_cache_blocks(hash_fn):
725
num_full_blocks=2,
726
block_size=block_size,
727
hash_fn=hash_fn,
728
+ kv_cache_group_id=0,
729
)
730
731
assert len(block_pool.cached_block_hash_to_block) == 2
@@ -573,11 +741,83 @@ def test_cache_blocks(hash_fn):
741
num_full_blocks=3,
742
block_size=block_size,
743
hash_fn=hash_fn,
744
+ kv_cache_group_id=0,
745
)
746
assert len(block_pool.cached_block_hash_to_block) == 3
747
assert blocks[0].block_hash is not None
748
749
750
+def test_cache_blocks_multi_group():
751
+ """
752
+ This tests that blocks are cached correctly for different kv cache groups.
753
+ """
754
+ block_size = 4
755
+ block_pool = BlockPool(num_gpu_blocks=10, enable_caching=True)
756
+
757
+ # Req:
758
+ # Block 0/4: [0, 1, 2, 3]
759
+ # Block 1/5: [4, 5, 6, 7]
760
+ # Block 2/6: [8, 9, 10, 11]
761
+ # Block 3/7: [12, 13]
762
+ req = make_request("0", list(range(14)))
763
+
764
+ # Cache the blocks for group 0.
765
+ blocks = [KVCacheBlock(block_id=i) for i in range(2)]
766
+ block_hashes: list[BlockHash] = []
767
+ block_pool.cache_full_blocks(
768
+ request=req,
769
+ blocks=blocks,
770
+ block_hashes=block_hashes,
771
+ num_cached_blocks=0,
772
+ num_full_blocks=2,
773
+ block_size=block_size,
774
+ hash_fn=hash,
775
+ kv_cache_group_id=0,
776
+ )
777
+ assert len(block_pool.cached_block_hash_to_block) == 2
778
+ assert len(block_hashes) == 2
779
+ assert all([block.block_hash is not None for block in blocks])
780
+
781
+ # Cache the blocks for group 1.
782
+ blocks = [KVCacheBlock(block_id=i) for i in range(3)]
783
+ block_pool.cache_full_blocks(
784
+ request=req,
785
+ blocks=blocks,
786
+ block_hashes=block_hashes,
787
+ num_cached_blocks=0,
788
+ num_full_blocks=3,
789
+ block_size=block_size,
790
+ hash_fn=hash,
791
+ kv_cache_group_id=1,
792
+ )
793
+ assert len(block_pool.cached_block_hash_to_block) == 5
794
+ assert len(block_hashes) == 3
795
+ assert all([block.block_hash is not None for block in blocks])
796
+
797
+ # Block hash 0: hit for group 0 and 1
798
+ # Block hash 1: hit for group 0 and 1
799
+ # Block hash 2: hit for group 1
800
+
801
+ assert block_pool.get_cached_block(block_hashes[0],
802
+ kv_cache_group_ids=[0]) is not None
803
+ assert block_pool.get_cached_block(block_hashes[1],
804
+ kv_cache_group_ids=[0]) is not None
805
+ assert block_pool.get_cached_block(block_hashes[2],
806
+ kv_cache_group_ids=[0]) is None
807
+ assert block_pool.get_cached_block(block_hashes[0],
808
+ kv_cache_group_ids=[1]) is not None
809
+ assert block_pool.get_cached_block(block_hashes[1],
810
+ kv_cache_group_ids=[1]) is not None
811
+ assert block_pool.get_cached_block(block_hashes[2],
812
+ kv_cache_group_ids=[1]) is not None
813
+ assert block_pool.get_cached_block(block_hashes[0],
814
+ kv_cache_group_ids=[0, 1]) is not None
815
+ assert block_pool.get_cached_block(block_hashes[1],
816
+ kv_cache_group_ids=[0, 1]) is not None
817
+ assert block_pool.get_cached_block(block_hashes[2],
818
+ kv_cache_group_ids=[0, 1]) is None
819
+
820
+
821
def test_mm_prefix_caching():
822
"""
823
This tests that the multi-modal prefix caching is correct.
@@ -614,7 +854,7 @@ def test_mm_prefix_caching():
854
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
855
856
# Completed block should have hashes with extra keys.
617
- assert not computed_blocks.blocks
857
+ assert not computed_blocks.blocks[0]
858
assert num_computed_tokens == 0
859
block_hashes = manager.req_to_block_hashes[req0.request_id]
860
assert len(block_hashes) == 3
@@ -623,7 +863,7 @@ def test_mm_prefix_caching():
863
assert block_hashes[2].extra_keys == ("bbb", )
864
865
blocks = manager.allocate_slots(req0, 59,
626
- len(computed_blocks.blocks) * 16,
866
+ len(computed_blocks.blocks[0]) * 16,
867
computed_blocks)
868
assert blocks.get_block_ids() == [[1, 2, 3, 4]]
869
req0.num_computed_tokens = 59
@@ -632,9 +872,9 @@ def test_mm_prefix_caching():
872
for _ in range(5):
873
req0.append_output_token_ids(8)
874
new_blocks = manager.allocate_slots(req0, 5,
635
- len(computed_blocks.blocks) * 16,
875
+ len(computed_blocks.blocks[0]) * 16,
876
computed_blocks)
637
- assert new_blocks is not None and len(new_blocks.blocks) == 0
877
+ assert new_blocks is not None and len(new_blocks.blocks[0]) == 0
878
879
# The just completed block should have hashes with extra keys.
880
assert len(block_hashes) == 4
@@ -652,7 +892,7 @@ def test_mm_prefix_caching():
892
mm_positions=mm_positions,
893
mm_hashes=mm_hashes)
894
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
655
- assert len(computed_blocks.blocks) == 3
895
+ assert len(computed_blocks.blocks[0]) == 3
896
assert num_computed_tokens == 3 * 16
897
898
@@ -675,7 +915,7 @@ def test_cache_key_salting():
915
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
916
917
# Completed block should have hashes with extra keys.
678
- assert not computed_blocks.blocks
918
+ assert not computed_blocks.blocks[0]
919
assert num_computed_tokens == 0
920
block_hashes = manager.req_to_block_hashes[req0.request_id]
921
assert len(block_hashes) == 3
@@ -684,7 +924,7 @@ def test_cache_key_salting():
924
assert block_hashes[2].extra_keys is None
925
926
blocks = manager.allocate_slots(req0, 59,
687
- len(computed_blocks.blocks) * 16,
927
+ len(computed_blocks.blocks[0]) * 16,
928
computed_blocks)
929
assert blocks.get_block_ids() == [[1, 2, 3, 4]]
930
req0.num_computed_tokens = 59
@@ -693,9 +933,9 @@ def test_cache_key_salting():
933
for _ in range(5):
934
req0.append_output_token_ids(8)
935
new_blocks = manager.allocate_slots(req0, 5,
696
- len(computed_blocks.blocks) * 16,
936
+ len(computed_blocks.blocks[0]) * 16,
937
computed_blocks)
698
- assert new_blocks is not None and len(new_blocks.blocks) == 0
938
+ assert new_blocks is not None and len(new_blocks.blocks[0]) == 0
939
940
# Now one more block that should not have extra keys.
941
assert len(block_hashes) == 4
@@ -706,14 +946,14 @@ def test_cache_key_salting():
946
req1 = make_request("1", token_ids, cache_salt="salt1")
947
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
948
# Should match only a prefix of 3 blocks.
709
- assert len(computed_blocks.blocks) == 3
949
+ assert len(computed_blocks.blocks[0]) == 3
950
assert num_computed_tokens == 3 * block_size
951
952
# Test cache miss with same content but different salt.
953
token_ids = common_token_ids + [4] * 11
954
req2 = make_request("2", token_ids, cache_salt="salt2")
955
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
716
- assert len(computed_blocks.blocks) == 0
956
+ assert len(computed_blocks.blocks[0]) == 0
957
assert num_computed_tokens == 0
958
block_hashes = manager.req_to_block_hashes[req2.request_id]
959
assert len(block_hashes) == 3
@@ -738,20 +978,24 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
978
common_token_ids = [i for i in range(3) for _ in range(16)]
979
req0 = make_request("0", common_token_ids)
980
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req0)
741
- assert not computed_blocks.blocks
981
+ assert not computed_blocks.blocks[0]
982
assert num_computed_tokens == 0
983
manager.allocate_slots(req0, 48,
744
- len(computed_blocks.blocks) * 16, computed_blocks)
745
- block_part0 = manager.single_type_manager.req_to_blocks[req0.request_id]
984
+ len(computed_blocks.blocks[0]) * 16,
985
+ computed_blocks)
986
+ block_part0 = manager.coordinator.single_type_managers[0].req_to_blocks[
987
+ req0.request_id]
988
989
# | Common-0 | Common-1 | Common-2 | Req1-3 | Req1-4 | Req1-5 | ... |
990
req1 = make_request("1", common_token_ids * 2)
991
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req1)
750
- assert computed_blocks.blocks == block_part0
992
+ assert computed_blocks.blocks[0] == block_part0
993
assert num_computed_tokens == 3 * 16
994
manager.allocate_slots(req1, 48,
753
- len(computed_blocks.blocks) * 16, computed_blocks)
754
- block_part1 = manager.single_type_manager.req_to_blocks[req1.request_id]
995
+ len(computed_blocks.blocks[0]) * 16,
996
+ computed_blocks)
997
+ block_part1 = manager.coordinator.single_type_managers[0].req_to_blocks[
998
+ req1.request_id]
999
# | Common-0 | Common-1 | Common-2 | Req1-3 (F) | Req1-4 (F) |
1000
# | Req1-5(F)| ... |
1001
manager.free(req1)
@@ -762,10 +1006,11 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
1006
# | Req1-5(F)| Req2-0 | Req2-1 | ... |
1007
req2 = make_request("2", [7] * block_size * 2)
1008
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req2)
765
- assert not computed_blocks.blocks
1009
+ assert not computed_blocks.blocks[0]
1010
assert num_computed_tokens == 0
1011
manager.allocate_slots(req2, block_size * 2,
768
- len(computed_blocks.blocks) * 16, computed_blocks)
1012
+ len(computed_blocks.blocks[0]) * 16,
1013
+ computed_blocks)
1014
1015
# Req3 is Req2 + 3 new blocks, so the first 6 blocks are computed,
1016
# but it cannot be allocated due to insufficient free blocks (2).
@@ -773,11 +1018,11 @@ def test_prefill_not_enough_free_blocks_with_computed_blocks():
1018
assert manager.block_pool.free_block_queue.num_free_blocks == 5
1019
req3 = make_request("3", common_token_ids * 3)
1020
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req3)
776
- assert computed_blocks.blocks == block_part1
1021
+ assert computed_blocks.blocks[0] == block_part1
1022
assert num_computed_tokens == 6 * 16
1023
# Req3 cannot be allocated.
1024
assert manager.allocate_slots(req3, 48,
780
- len(computed_blocks.blocks) * 16,
1025
+ len(computed_blocks.blocks[0]) * 16,
1026
computed_blocks) is None
1027
# Block 0-2 are used by Req 1.
1028
assert {block.ref_cnt for block in block_part1[:3]} == {1}
@@ -804,9 +1049,9 @@ def test_reset_prefix_cache():
1049
req1 = make_request("1", all_token_ids)
1050
computed_blocks, _ = manager.get_computed_blocks(req1)
1051
assert len(manager.req_to_block_hashes[req1.request_id]) == 3
807
- assert len(computed_blocks.blocks) == 3
1052
+ assert len(computed_blocks.blocks[0]) == 3
1053
blocks = manager.allocate_slots(req1, 7,
809
- len(computed_blocks.blocks) * 16,
1054
+ len(computed_blocks.blocks[0]) * 16,
1055
computed_blocks)
1056
assert blocks.get_block_ids() == [[5]]
1057
@@ -836,10 +1081,11 @@ def test_prefix_cache_stats_disabled():
1081
# Call all functions that check whether log_stats is disabled.
1082
req = make_request("0", list(range(16)))
1083
computed_blocks, num_computed_tokens = manager.get_computed_blocks(req)
839
- assert not computed_blocks.blocks
1084
+ assert not computed_blocks.blocks[0]
1085
assert num_computed_tokens == 0
1086
manager.allocate_slots(req, 16,
842
- len(computed_blocks.blocks) * 16, computed_blocks)
1087
+ len(computed_blocks.blocks[0]) * 16,
1088
+ computed_blocks)
1089
manager.reset_prefix_cache()
1090
1091
# Ensure prefix_cache_stats remains None
@@ -918,7 +1164,8 @@ def test_eagle_enabled_removes_last_block():
1164
# Prime the cache
1165
computed_blocks, _ = manager.get_computed_blocks(req)
1166
manager.allocate_slots(req, len(token_ids),
921
- len(computed_blocks.blocks) * 16, computed_blocks)
1167
+ len(computed_blocks.blocks[0]) * 16,
1168
+ computed_blocks)
1169
manager.free(req)
1170
1171
# New request with same tokens + Eagle enabled
@@ -928,7 +1175,7 @@ def test_eagle_enabled_removes_last_block():
1175
# Should retain 1 block:
1176
# 1. Original 3 blocks → pop last hash → 2 matched blocks
1177
# 2. drop last matched block → 1 remaining block
931
- assert len(computed_blocks.blocks) == 1
1178
+ assert len(computed_blocks.blocks[0]) == 1
1179
assert num_tokens == 1 * block_size # 16 tokens
1180
1181
@@ -948,14 +1195,15 @@ def test_eagle_with_partial_blocks():
1195
# Prime the cache
1196
computed_blocks, _ = manager.get_computed_blocks(req)
1197
manager.allocate_slots(req, len(token_ids),
951
- len(computed_blocks.blocks) * 16, computed_blocks)
1198
+ len(computed_blocks.blocks[0]) * 16,
1199
+ computed_blocks)
1200
manager.free(req)
1201
1202
# New request with Eagle enabled
1203
req_eagle = make_request("partial_eagle", token_ids)
1204
computed_blocks, num_tokens = manager.get_computed_blocks(req_eagle)
1205
# Original match: 2 full blocks → Eagle removes 1 → 1 remaining
958
- assert len(computed_blocks.blocks) == 1
1206
+ assert len(computed_blocks.blocks[0]) == 1
1207
assert num_tokens == 1 * block_size
1208
1209
@@ -973,7 +1221,7 @@ def test_eagle_with_sliding_window():
1221
manager = KVCacheManager(
1222
KVCacheConfig(
1223
num_blocks=10,
976
- tensors={},
1224
+ kv_cache_tensors=[],
1225
kv_cache_groups=[KVCacheGroupSpec(['layer'], sliding_window_spec)],
1226
),
1227
max_model_len=8192,
@@ -988,7 +1236,8 @@ def test_eagle_with_sliding_window():
1236
# Prime the cache
1237
computed_blocks, _ = manager.get_computed_blocks(req)
1238
manager.allocate_slots(req, len(token_ids),
991
- len(computed_blocks.blocks) * 16, computed_blocks)
1239
+ len(computed_blocks.blocks[0]) * 16,
1240
+ computed_blocks)
1241
# record the block hash of the first block in the request for later use
1242
block_hash_first_block = manager.req_to_block_hashes[req.request_id][0]
1243
assert block_hash_first_block is not None
@@ -998,13 +1247,14 @@ def test_eagle_with_sliding_window():
1247
req_eagle = make_request("partial_eagle", token_ids)
1248
computed_blocks, num_tokens = manager.get_computed_blocks(req_eagle)
1249
# Original match: 2 full blocks → Eagle removes 1 → 1 remaining
1001
- assert len(computed_blocks.blocks) == 1
1250
+ assert len(computed_blocks.blocks[0]) == 1
1251
assert num_tokens == 1 * block_size
1252
1253
# Evict the first block in the request
1254
assert manager.block_pool.get_cached_block(
1006
- block_hash_first_block) is not None
1007
- manager.block_pool.cached_block_hash_to_block.pop(block_hash_first_block)
1255
+ block_hash_first_block, kv_cache_group_ids=[0]) is not None
1256
+ manager.block_pool.cached_block_hash_to_block.pop(
1257
+ BlockHashWithGroupId(block_hash_first_block, 0))
1258
1259
# New request
1260
req_after_evict = make_request("partial_eagle_after_evict", token_ids)
@@ -1012,5 +1262,5 @@ def test_eagle_with_sliding_window():
1262
# Cache miss. The only hit prefix is [NULL_BLOCK, BLOCK_2] if eagle is
1263
# not considered. But after dropping the last matched block due to eagle,
1264
# there will be no matched prefix.
1015
- assert len(computed_blocks.blocks) == 0
1265
+ assert len(computed_blocks.blocks[0]) == 0
1266
assert num_tokens == 0
tests/v1/core/test_scheduler.py
+8
-8
@@ -97,7 +97,7 @@ def create_scheduler(
97
)
98
kv_cache_config = KVCacheConfig(
99
num_blocks=num_blocks, # A large number of blocks to hold all requests
100
- tensors={},
100
+ kv_cache_tensors=[],
101
kv_cache_groups=[
102
KVCacheGroupSpec(['layer'],
103
FullAttentionSpec(block_size, 1, 1, torch.float32,
@@ -814,10 +814,10 @@ def _assert_right_kv_cache_manager(
814
# Make sure the request stats are right.
815
EXPECTED_TOTAL_BLOCKS = num_tokens // block_size
816
for req_id in req_ids:
817
- blocks = (scheduler.kv_cache_manager.single_type_manager.
818
- req_to_blocks[req_id])
817
+ blocks = (scheduler.kv_cache_manager.coordinator.
818
+ single_type_managers[0].req_to_blocks[req_id])
819
hashes = scheduler.kv_cache_manager.req_to_block_hashes[req_id]
820
- assert (scheduler.kv_cache_manager.single_type_manager.
820
+ assert (scheduler.kv_cache_manager.coordinator.single_type_managers[0].
821
num_cached_block[req_id] == EXPECTED_TOTAL_BLOCKS)
822
assert len(blocks) == EXPECTED_TOTAL_BLOCKS
823
assert len(hashes) == EXPECTED_TOTAL_BLOCKS
@@ -1198,11 +1198,11 @@ def assert_scheduler_empty(scheduler: Scheduler):
1198
assert len(scheduler.encoder_cache_manager.cached) == 0
1199
1200
# KVCache Manager.
1201
- assert len(
1202
- scheduler.kv_cache_manager.single_type_manager.req_to_blocks) == 0
1201
+ assert len(scheduler.kv_cache_manager.coordinator.single_type_managers[0].
1202
+ req_to_blocks) == 0
1203
assert len(scheduler.kv_cache_manager.req_to_block_hashes) == 0
1204
- assert len(
1205
- scheduler.kv_cache_manager.single_type_manager.num_cached_block) == 0
1204
+ assert len(scheduler.kv_cache_manager.coordinator.single_type_managers[0].
1205
+ num_cached_block) == 0
1206
num_free_blocks = (
1207
scheduler.kv_cache_manager.block_pool.free_block_queue.num_free_blocks)
1208
assert num_free_blocks == (
tests/v1/core/test_specialized_manager.py
+16
-11
@@ -4,7 +4,8 @@
4
import torch
5
6
from vllm.v1.core.block_pool import BlockPool
7
-from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock
7
+from vllm.v1.core.kv_cache_utils import (BlockHash, BlockHashWithGroupId,
8
+ KVCacheBlock)
9
from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager
10
from vllm.v1.kv_cache_interface import SlidingWindowSpec
11
@@ -12,9 +13,8 @@ from vllm.v1.kv_cache_interface import SlidingWindowSpec
13
def get_sliding_window_manager(sliding_window_spec, block_pool):
14
return SlidingWindowManager(sliding_window_spec,
15
block_pool,
15
- use_eagle=False,
16
- num_kv_cache_groups=1,
17
- caching_hash_fn=lambda x: x)
16
+ caching_hash_fn=lambda x: x,
17
+ kv_cache_group_id=0)
18
19
20
def test_sliding_window_possible_cached_prefix():
@@ -42,13 +42,18 @@ def test_sliding_window_possible_cached_prefix():
42
for i, (block_hash,
43
is_cached) in enumerate(zip(block_hash_list, block_is_cached)):
44
if is_cached:
45
- block_pool.cached_block_hash_to_block[block_hash] = {
46
- i: block_pool.blocks[i + 10]
47
- }
45
+ block_pool.cached_block_hash_to_block[BlockHashWithGroupId(
46
+ block_hash, 0)] = {
47
+ i: block_pool.blocks[i + 10],
48
+ }
49
50
computed_blocks = manager.find_longest_cache_hit(
50
- block_hash_list,
51
- len(block_hash_list) * block_size)
51
+ block_hashes=block_hash_list,
52
+ max_length=len(block_hash_list) * block_size,
53
+ kv_cache_group_ids=[0],
54
+ block_pool=block_pool,
55
+ kv_cache_spec=sliding_window_spec,
56
+ use_eagle=False)[0]
57
assert len(computed_blocks) == expect_length
58
59
assert all(block == block_pool.null_block
@@ -95,13 +100,13 @@ def test_sliding_window_remove_skipped_blocks():
100
101
null_block_id = block_pool.null_block.block_id
102
98
- def id_to_block_table(ids):
103
+ def id_to_block_table(ids) -> list[KVCacheBlock]:
104
return [
105
KVCacheBlock(id_)
106
if id_ != null_block_id else block_pool.null_block for id_ in ids
107
]
108
104
- def assert_block_id(block_table, ids):
109
+ def assert_block_id(block_table: list[KVCacheBlock], ids: list[int]):
110
for block, id_ in zip(block_table, ids):
111
if id_ == null_block_id:
112
assert block == block_pool.null_block
tests/v1/e2e/test_correctness_sliding_window.py
+2
-2
@@ -18,7 +18,7 @@ class TestConfig:
18
19
model_config = {
20
"bigcode/starcoder2-3b": TestConfig(4096, (800, 1100)),
21
- "google/gemma-2-2b-it": TestConfig(4096, (400, 800)),
21
+ "google/gemma-3-1b-it": TestConfig(4096, (400, 800)),
22
}
23
24
@@ -26,7 +26,7 @@ model_config = {
26
"model",
27
[
28
"bigcode/starcoder2-3b", # sliding window only
29
- "google/gemma-2-2b-it", # sliding window + full attention
29
+ "google/gemma-3-1b-it", # sliding window + full attention
30
])
31
@pytest.mark.parametrize("batch_size", [5])
32
@pytest.mark.parametrize("seed", [1])
tests/v1/kv_connector/unit/test_nixl_connector.py
+2
-2
@@ -36,8 +36,8 @@ def test_basic_inferface():
36
req_meta = kv_connector_metadata.requests[request_id]
37
38
for block_id, block in zip(
39
- req_meta.local_block_ids, scheduler.kv_cache_manager.
40
- single_type_manager.req_to_blocks[request_id]):
39
+ req_meta.local_block_ids, scheduler.kv_cache_manager.coordinator.
40
+ single_type_managers[0].req_to_blocks[request_id]):
41
assert block_id == block.block_id
42
43
tests/v1/kv_connector/unit/test_remote_decode_lifecycle.py
+2
-2
@@ -54,8 +54,8 @@ def test_basic_lifecycle():
54
assert len(scheduler.waiting) == 0
55
56
# ... but blocks should not be freed.
57
- blocks = scheduler.kv_cache_manager.single_type_manager.req_to_blocks[
58
- request_id]
57
+ blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[
58
+ 0].req_to_blocks[request_id]
59
for block in blocks:
60
assert block.ref_cnt == 1
61
tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py
+12
-12
@@ -51,8 +51,8 @@ def test_basic_lifecycle():
51
assert (block_pool.free_block_queue.num_free_blocks
52
< START_FREE_BLOCK_QUEUE_SIZE)
53
assert len(block_pool.cached_block_hash_to_block) == 0
54
- blocks = scheduler.kv_cache_manager.single_type_manager.req_to_blocks[
55
- request_id]
54
+ blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[
55
+ 0].req_to_blocks[request_id]
56
for block in blocks:
57
assert block._block_hash is None
58
@@ -87,8 +87,8 @@ def test_basic_lifecycle():
87
88
# Confirm the block are actually allocated.
89
num_hashed_blocks = 0
90
- blocks = scheduler.kv_cache_manager.single_type_manager.req_to_blocks[
91
- request_id]
90
+ blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[
91
+ 0].req_to_blocks[request_id]
92
for block in blocks:
93
assert block.ref_cnt == 1
94
num_hashed_blocks += (1 if block._block_hash is not None else 0)
@@ -261,10 +261,10 @@ def test_no_spurious_prefix_caching():
261
assert len(scheduler.running) == 1
262
assert len(scheduler.waiting) == 1
263
264
- local_blocks = scheduler.kv_cache_manager.single_type_manager.req_to_blocks[
265
- request_local.request_id]
266
- remote_blocks = scheduler.kv_cache_manager.single_type_manager.req_to_blocks[ # noqa: E501
267
- request_remote.request_id]
264
+ local_blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[
265
+ 0].req_to_blocks[request_local.request_id]
266
+ remote_blocks = scheduler.kv_cache_manager.coordinator.single_type_managers[
267
+ 0].req_to_blocks[request_remote.request_id]
268
269
# Local should have cached blocks (but not all due to preallocate).
270
num_hashed_blocks = 0
@@ -300,8 +300,8 @@ def test_full_block_prompt():
300
# STEP (1): Initialize a recv.
301
scheduler_output = scheduler.schedule()
302
# All blocks should be allocated.
303
- num_blocks = len(scheduler.kv_cache_manager.single_type_manager.
304
- req_to_blocks[request_id])
303
+ num_blocks = len(scheduler.kv_cache_manager.coordinator.
304
+ single_type_managers[0].req_to_blocks[request_id])
305
assert num_blocks == NUM_EXTERNAL_FULL_BLOCKS
306
model_runner_output = EMPTY_MODEL_RUNNER_OUTPUT
307
scheduler.update_from_output(scheduler_output, model_runner_output)
@@ -319,8 +319,8 @@ def test_full_block_prompt():
319
320
# We need to recompute the final token of the prompt to generate
321
# the first new token, so we should not have a new block.
322
- num_blocks = len(scheduler.kv_cache_manager.single_type_manager.
323
- req_to_blocks[request_id])
322
+ num_blocks = len(scheduler.kv_cache_manager.coordinator.
323
+ single_type_managers[0].req_to_blocks[request_id])
324
assert num_blocks == NUM_EXTERNAL_FULL_BLOCKS
325
assert (scheduler_output.scheduled_new_reqs[0].num_computed_tokens ==
326
NUM_TOKENS - 1)
tests/v1/kv_connector/unit/utils.py
+5
-5
@@ -32,11 +32,11 @@ def assert_scheduler_empty(scheduler: Scheduler):
32
assert len(scheduler.encoder_cache_manager.cached) == 0
33
34
# KVCache Manager.
35
- assert len(
36
- scheduler.kv_cache_manager.single_type_manager.req_to_blocks) == 0
35
+ assert len(scheduler.kv_cache_manager.coordinator.single_type_managers[0].
36
+ req_to_blocks) == 0
37
assert len(scheduler.kv_cache_manager.req_to_block_hashes) == 0
38
- assert len(
39
- scheduler.kv_cache_manager.single_type_manager.num_cached_block) == 0
38
+ assert len(scheduler.kv_cache_manager.coordinator.single_type_managers[0].
39
+ num_cached_block) == 0
40
num_free_blocks = (
41
scheduler.kv_cache_manager.block_pool.free_block_queue.num_free_blocks)
42
assert num_free_blocks == (
@@ -96,7 +96,7 @@ def create_scheduler(
96
block_size = vllm_config.cache_config.block_size
97
kv_cache_config = KVCacheConfig(
98
num_blocks=num_blocks, # A large number of blocks to hold all requests
99
- tensors={},
99
+ kv_cache_tensors=[],
100
kv_cache_groups=[
101
KVCacheGroupSpec(['layer'],
102
FullAttentionSpec(block_size, 1, 1, torch.float32,
tests/v1/worker/test_gpu_model_runner.py
+14
-13
@@ -40,12 +40,13 @@ def initialize_kv_cache(runner: GPUModelRunner):
40
tensor_size = attn_spec.page_size_bytes * NUM_BLOCKS
41
kv_cache_config = KVCacheConfig(
42
num_blocks=NUM_BLOCKS,
43
- tensors={
44
- "layer.0": KVCacheTensor(size=tensor_size),
45
- },
43
+ kv_cache_tensors=[
44
+ KVCacheTensor(size=tensor_size, shared_by=["layer.0"]),
45
+ ],
46
kv_cache_groups=[
47
KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=attn_spec)
48
- ])
48
+ ],
49
+ )
50
runner.kv_cache_config = kv_cache_config
51
runner.input_batch = InputBatch(
52
max_num_reqs=runner.max_num_reqs,
@@ -518,9 +519,9 @@ def test_init_kv_cache_without_kv_sharing():
519
kv_cache_config = get_kv_cache_config(vllm_config, kv_cache_spec,
520
available_memory)
521
assert kv_cache_config.num_blocks == num_expected_blocks
521
- assert len(kv_cache_config.tensors) == 2
522
- assert kv_cache_config.tensors[layer_0].size == available_memory // 2
523
- assert kv_cache_config.tensors[layer_1].size == available_memory // 2
522
+ assert len(kv_cache_config.kv_cache_tensors) == 2
523
+ assert kv_cache_config.kv_cache_tensors[0].size == available_memory // 2
524
+ assert kv_cache_config.kv_cache_tensors[1].size == available_memory // 2
525
526
max_context_len =\
527
estimate_max_model_len(vllm_config, kv_cache_spec, 5 * GiB_bytes)
@@ -530,9 +531,9 @@ def test_init_kv_cache_without_kv_sharing():
531
# important: override tensor size to prevent large mem alloc during test
532
# this will only allocate 2 block worth of memory (2 * 32kb)
533
kv_cache_config.num_blocks = 1
533
- for layer in kv_cache_config.tensors:
534
- kv_cache_config.tensors[layer].size =\
535
- kv_cache_spec[layer].page_size_bytes
534
+ for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
535
+ kv_cache_tensor.size = (
536
+ kv_cache_spec[kv_cache_tensor.shared_by[0]].page_size_bytes)
537
538
runner.initialize_kv_cache(kv_cache_config)
539
@@ -589,10 +590,10 @@ def test_init_kv_cache_with_kv_sharing_valid():
590
kv_cache_config = get_kv_cache_config(vllm_config, kv_cache_spec,
591
available_memory)
592
assert kv_cache_config.num_blocks == num_expected_blocks
592
- assert len(kv_cache_config.tensors) == 1
593
+ assert len(kv_cache_config.kv_cache_tensors) == 1
594
# Each layer now has twice the available memory for KV cache
595
# compared to no KV sharing
595
- assert kv_cache_config.tensors[layer_0].size == available_memory
596
+ assert kv_cache_config.kv_cache_tensors[0].size == available_memory
597
598
max_context_len =\
599
estimate_max_model_len(vllm_config, kv_cache_spec, 5 * GiB_bytes)
@@ -602,7 +603,7 @@ def test_init_kv_cache_with_kv_sharing_valid():
603
# important: override tensor size to prevent large mem alloc during test
604
# this will only allocate 1 block worth of memory (32kb)
605
kv_cache_config.num_blocks = 1
605
- kv_cache_config.tensors[layer_0].size =\
606
+ kv_cache_config.kv_cache_tensors[0].size =\
607
kv_cache_spec[layer_0].page_size_bytes
608
609
runner.initialize_kv_cache(kv_cache_config)
vllm/config.py
+21
@@ -2104,6 +2104,12 @@ class SchedulerConfig:
2104
default scheduler. Can be a class directly or the path to a class of form
2105
"mod.custom_class"."""
2106
2107
+ disable_hybrid_kv_cache_manager: bool = False
2108
+ """If set to True, KV cache manager will allocate the same size of KV cache
2109
+ for all attention layers even if there are multiple type of attention layers
2110
+ like full attention and sliding window attention.
2111
+ """
2112
+
2113
def compute_hash(self) -> str:
2114
"""
2115
WARNING: Whenever a new field is added to this config,
@@ -4465,6 +4471,21 @@ class VllmConfig:
4471
if not self.instance_id:
4472
self.instance_id = random_uuid()[:5]
4473
4474
+ if (envs.VLLM_USE_V1
4475
+ and not self.scheduler_config.disable_hybrid_kv_cache_manager):
4476
+ # logger should only print warning message for hybrid models. As we
4477
+ # can't know whether the model is hybrid or not now, so we don't log
4478
+ # warning message here and will log it later.
4479
+ if not (current_platform.is_cuda() or current_platform.is_rocm()):
4480
+ # Hybrid KV cache manager is not supported on non-GPU platforms.
4481
+ self.disable_hybrid_kv_cache_manager = True
4482
+ if self.kv_transfer_config is not None:
4483
+ # Hybrid KV cache manager is not compatible with KV transfer.
4484
+ self.disable_hybrid_kv_cache_manager = True
4485
+ if self.kv_events_config is not None:
4486
+ # Hybrid KV cache manager is not compatible with KV events.
4487
+ self.disable_hybrid_kv_cache_manager = True
4488
+
4489
def update_sizes_for_sequence_parallelism(self,
4490
possible_sizes: list) -> list:
4491
# remove the sizes that not multiple of tp_size when
vllm/engine/arg_utils.py
+8
@@ -387,6 +387,9 @@ class EngineArgs:
387
bool] = SchedulerConfig.enable_chunked_prefill
388
disable_chunked_mm_input: bool = SchedulerConfig.disable_chunked_mm_input
389
390
+ disable_hybrid_kv_cache_manager: bool = (
391
+ SchedulerConfig.disable_hybrid_kv_cache_manager)
392
+
393
guided_decoding_backend: GuidedDecodingBackend = DecodingConfig.backend
394
guided_decoding_disable_fallback: bool = DecodingConfig.disable_fallback
395
guided_decoding_disable_any_whitespace: bool = \
@@ -849,6 +852,9 @@ class EngineArgs:
852
**scheduler_kwargs["disable_chunked_mm_input"])
853
scheduler_group.add_argument("--scheduler-cls",
854
**scheduler_kwargs["scheduler_cls"])
855
+ scheduler_group.add_argument(
856
+ "--disable-hybrid-kv-cache-manager",
857
+ **scheduler_kwargs["disable_hybrid_kv_cache_manager"])
858
859
# vLLM arguments
860
vllm_kwargs = get_kwargs(VllmConfig)
@@ -1174,6 +1180,8 @@ class EngineArgs:
1180
max_num_partial_prefills=self.max_num_partial_prefills,
1181
max_long_partial_prefills=self.max_long_partial_prefills,
1182
long_prefill_token_threshold=self.long_prefill_token_threshold,
1183
+ disable_hybrid_kv_cache_manager=self.
1184
+ disable_hybrid_kv_cache_manager,
1185
)
1186
1187
lora_config = LoRAConfig(
vllm/v1/core/block_pool.py
+44
-25
@@ -7,8 +7,8 @@ from typing import Callable, Optional
7
from vllm.distributed.kv_events import (AllBlocksCleared, BlockRemoved,
8
BlockStored, KVCacheEvent)
9
from vllm.logger import init_logger
10
-from vllm.v1.core.kv_cache_utils import (BlockHash, FreeKVCacheBlockQueue,
11
- KVCacheBlock,
10
+from vllm.v1.core.kv_cache_utils import (BlockHash, BlockHashWithGroupId,
11
+ FreeKVCacheBlockQueue, KVCacheBlock,
12
generate_block_hash_extra_keys,
13
hash_block_tokens)
14
from vllm.v1.request import Request
@@ -27,6 +27,7 @@ class BlockPool:
27
Args:
28
num_gpu_blocks: The number of blocks in the pool.
29
enable_caching: Whether to enable prefix caching.
30
+ enable_kv_cache_events: Whether to enable kv cache events.
31
"""
32
33
def __init__(
@@ -56,7 +57,7 @@ class BlockPool:
57
# if there is already an identical block in the cache. This is because
58
# we want to make sure the allocated block IDs won't change so that
59
# block tables are append-only.
59
- self.cached_block_hash_to_block: dict[BlockHash, dict[
60
+ self.cached_block_hash_to_block: dict[BlockHashWithGroupId, dict[
61
int, KVCacheBlock]] = defaultdict(dict)
62
63
# To represent a placeholder block with block_id=0.
@@ -68,22 +69,29 @@ class BlockPool:
69
self.enable_kv_cache_events = enable_kv_cache_events
70
self.kv_event_queue: list[KVCacheEvent] = []
71
71
- def get_cached_block(self,
72
- block_hash: BlockHash) -> Optional[KVCacheBlock]:
73
- """Get a cached block by the block hash, or None if cache miss.
72
+ def get_cached_block(
73
+ self, block_hash: BlockHash,
74
+ kv_cache_group_ids: list[int]) -> Optional[list[KVCacheBlock]]:
75
+ """Get the cached block by the block hash for each group in
76
+ `kv_cache_group_ids`, or None if cache miss for any group.
77
If there are duplicated blocks, we return the first block in the cache.
78
79
Args:
80
block_hash: The hash value of the block.
81
+ kv_cache_group_ids: The ids of the KV cache groups.
82
83
Returns:
80
- The cached block if it exists, or None.
84
+ The cached blocks if exists, or None.
85
"""
82
- cached_blocks = self.cached_block_hash_to_block.get(block_hash)
83
- if not cached_blocks:
84
- return None
85
- first_block_id = next(iter(cached_blocks))
86
- return cached_blocks[first_block_id]
86
+ cached_blocks = []
87
+ for group_id in kv_cache_group_ids:
88
+ cached_blocks_one_group = self.cached_block_hash_to_block.get(
89
+ BlockHashWithGroupId(block_hash, group_id))
90
+ if not cached_blocks_one_group:
91
+ return None
92
+ first_block_id = next(iter(cached_blocks_one_group))
93
+ cached_blocks.append(cached_blocks_one_group[first_block_id])
94
+ return cached_blocks
95
96
def cache_full_blocks(
97
self,
@@ -93,6 +101,7 @@ class BlockPool:
101
num_cached_blocks: int,
102
num_full_blocks: int,
103
block_size: int,
104
+ kv_cache_group_id: int,
105
hash_fn: Callable,
106
) -> None:
107
"""Cache a list of full blocks for prefix caching.
@@ -112,6 +121,7 @@ class BlockPool:
121
num_full_blocks: The number of blocks that are full and should
122
be cached after this function.
123
block_size: Number of tokens in each block.
124
+ kv_cache_group_id: The id of the KV cache group.
125
hash_fn: The hash function to use for block hashes.
126
"""
127
if num_cached_blocks == num_full_blocks:
@@ -126,7 +136,7 @@ class BlockPool:
136
else:
137
prev_block = blocks[num_cached_blocks - 1]
138
assert prev_block.block_hash is not None
129
- prev_block_hash_value = prev_block.block_hash.hash_value
139
+ prev_block_hash_value = prev_block.block_hash.get_hash_value()
140
141
parent_block_hash = prev_block_hash_value
142
new_hashes: Optional[list[int]] = ([] if self.enable_kv_cache_events
@@ -138,8 +148,9 @@ class BlockPool:
148
# The block hash may already be computed in
149
# "get_computed_blocks" if the tokens are not generated by
150
# this request (either the prompt tokens or the previously
141
- # generated tokens with preemption). In this case we simply
142
- # reuse the block hash.
151
+ # generated tokens with preemption), or by other
152
+ # single_type_managers with the same block_size.
153
+ # In this case we simply reuse the block hash.
154
block_hash = new_block_hashes[i]
155
else:
156
# Otherwise compute the block hash and cache it in the request
@@ -166,8 +177,11 @@ class BlockPool:
177
block_hashes.append(block_hash)
178
179
# Update and added the full block to the cache.
169
- blk.block_hash = block_hash
170
- self.cached_block_hash_to_block[block_hash][blk.block_id] = blk
180
+ block_hash_with_group_id = BlockHashWithGroupId(
181
+ block_hash, kv_cache_group_id)
182
+ blk.block_hash = block_hash_with_group_id
183
+ self.cached_block_hash_to_block[block_hash_with_group_id][
184
+ blk.block_id] = blk
185
if new_hashes is not None:
186
new_hashes.append(block_hash.hash_value)
187
prev_block_hash_value = block_hash.hash_value
@@ -237,12 +251,16 @@ class BlockPool:
251
del self.cached_block_hash_to_block[block_hash]
252
253
if self.enable_kv_cache_events:
254
+ # FIXME (Chen): Not sure whether we should return `hash_value`
255
+ # or `(hash_value, group_id)` here. But it's fine now because
256
+ # we disable hybrid kv cache manager when kv cache event is
257
+ # enabled, so there is only one group.
258
self.kv_event_queue.append(
241
- BlockRemoved(block_hashes=[block_hash.hash_value]))
259
+ BlockRemoved(block_hashes=[block_hash.get_hash_value()]))
260
return True
261
return False
262
245
- def touch(self, blocks: list[KVCacheBlock]) -> None:
263
+ def touch(self, blocks: list[list[KVCacheBlock]]) -> None:
264
"""Touch a block increases its reference count by 1, and may remove
265
the block from the free queue. This is used when a block is hit by
266
another request with the same prefix.
@@ -250,12 +268,13 @@ class BlockPool:
268
Args:
269
blocks: A list of blocks to touch.
270
"""
253
- for block in blocks:
254
- # ref_cnt=0 means this block is in the free list (i.e. eviction
255
- # candidate), so remove it.
256
- if block.ref_cnt == 0 and not block.is_null:
257
- self.free_block_queue.remove(block)
258
- block.incr_ref()
271
+ for blocks_per_group in blocks:
272
+ for block in blocks_per_group:
273
+ # ref_cnt=0 means this block is in the free list (i.e. eviction
274
+ # candidate), so remove it.
275
+ if block.ref_cnt == 0 and not block.is_null:
276
+ self.free_block_queue.remove(block)
277
+ block.incr_ref()
278
279
def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None:
280
"""Free a list of blocks. The blocks should be ordered by their
vllm/v1/core/kv_cache_coordinator.py
new
+358
@@ -0,0 +1,358 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from abc import ABC, abstractmethod
3
+from typing import Callable, Optional
4
+
5
+from vllm.v1.core.block_pool import BlockPool
6
+from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock
7
+from vllm.v1.core.single_type_kv_cache_manager import (
8
+ FullAttentionManager, SingleTypeKVCacheManager,
9
+ get_manager_for_kv_cache_spec)
10
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig
11
+from vllm.v1.request import Request
12
+
13
+
14
+class KVCacheCoordinator(ABC):
15
+ """
16
+ Coordinate the KV cache of different KV cache groups.
17
+ """
18
+
19
+ def __init__(
20
+ self,
21
+ kv_cache_config: KVCacheConfig,
22
+ max_model_len: int,
23
+ use_eagle: bool,
24
+ enable_caching: bool,
25
+ caching_hash_fn: Callable,
26
+ enable_kv_cache_events: bool,
27
+ ):
28
+ self.kv_cache_config = kv_cache_config
29
+ self.max_model_len = max_model_len
30
+
31
+ self.block_pool = BlockPool(kv_cache_config.num_blocks, enable_caching,
32
+ enable_kv_cache_events)
33
+ self.single_type_managers: list[SingleTypeKVCacheManager] = []
34
+
35
+ # Needs special handling for find_longest_cache_hit if eagle is enabled
36
+ self.use_eagle = use_eagle
37
+
38
+ for i in range(len(self.kv_cache_config.kv_cache_groups)):
39
+ kv_cache_spec = self.kv_cache_config.kv_cache_groups[
40
+ i].kv_cache_spec
41
+ self.single_type_managers.append(
42
+ get_manager_for_kv_cache_spec(
43
+ kv_cache_spec=kv_cache_spec,
44
+ block_pool=self.block_pool,
45
+ kv_cache_group_id=i,
46
+ caching_hash_fn=caching_hash_fn,
47
+ ))
48
+
49
+ def get_num_blocks_to_allocate(
50
+ self, request_id: str, num_tokens: int,
51
+ new_computed_blocks: list[list[KVCacheBlock]]) -> int:
52
+ """
53
+ Get the number of blocks needed to be allocated for the request.
54
+
55
+ Args:
56
+ request_id: The request ID.
57
+ num_tokens: The total number of tokens that need a slot (including
58
+ tokens that are already allocated).
59
+ new_computed_blocks: The new computed blocks just hitting the
60
+ prefix caching.
61
+
62
+ Returns:
63
+ The number of blocks.
64
+ """
65
+ num_blocks_to_allocate = 0
66
+ for i, manager in enumerate(self.single_type_managers):
67
+ num_blocks_to_allocate += manager.get_num_blocks_to_allocate(
68
+ request_id, num_tokens, new_computed_blocks[i])
69
+ return num_blocks_to_allocate
70
+
71
+ def save_new_computed_blocks(
72
+ self, request_id: str,
73
+ new_computed_blocks: list[list[KVCacheBlock]]) -> None:
74
+ """
75
+ Add the new computed blocks to the request.
76
+
77
+ Args:
78
+ request_id: The request ID.
79
+ new_computed_blocks: The new computed blocks just hitting the
80
+ prefix cache.
81
+ """
82
+ for i, manager in enumerate(self.single_type_managers):
83
+ manager.save_new_computed_blocks(request_id,
84
+ new_computed_blocks[i])
85
+
86
+ def allocate_new_blocks(self, request_id: str,
87
+ num_tokens: int) -> list[list[KVCacheBlock]]:
88
+ """
89
+ Allocate new blocks for the request to give it at least `num_tokens`
90
+ token slots.
91
+
92
+ Args:
93
+ request_id: The request ID.
94
+ num_tokens: The total number of tokens that need a slot (including
95
+ tokens that are already allocated).
96
+
97
+ Returns:
98
+ The new allocated blocks.
99
+ """
100
+ new_blocks = []
101
+ for manager in self.single_type_managers:
102
+ new_blocks.append(
103
+ manager.allocate_new_blocks(request_id, num_tokens))
104
+ return new_blocks
105
+
106
+ def cache_blocks(self, request: Request, block_hashes: list[BlockHash],
107
+ num_computed_tokens: int) -> None:
108
+ """
109
+ Cache the blocks for the request.
110
+
111
+ Args:
112
+ request: The request.
113
+ block_hashes: The block hashes of the request.
114
+ num_tokens: The total number of tokens that need to be cached
115
+ (including tokens that are already cached).
116
+ """
117
+ for manager in self.single_type_managers:
118
+ manager.cache_blocks(request, block_hashes, num_computed_tokens)
119
+
120
+ def free(self, request_id: str) -> None:
121
+ """
122
+ Free the blocks for the request.
123
+
124
+ Args:
125
+ request_id: The request ID.
126
+ """
127
+ for manager in self.single_type_managers:
128
+ manager.free(request_id)
129
+
130
+ def get_num_common_prefix_blocks(self, request_id: str,
131
+ num_running_requests: int) -> list[int]:
132
+ """
133
+ Get the number of common prefix blocks for a request.
134
+
135
+ Args:
136
+ request_id: The request ID.
137
+ block_hashes: The block hashes of the request.
138
+
139
+ Returns:
140
+ The number of common prefix blocks.
141
+ """
142
+ num_blocks_per_group = [
143
+ manager.get_num_common_prefix_blocks(request_id,
144
+ num_running_requests)
145
+ for manager in self.single_type_managers
146
+ ]
147
+ return num_blocks_per_group
148
+
149
+ def remove_skipped_blocks(self, request_id: str,
150
+ num_computed_tokens: int) -> None:
151
+ """
152
+ Remove the blocks that are no longer needed from `blocks` and replace
153
+ the removed blocks with null_block.
154
+
155
+ Args:
156
+ request_id: The request ID.
157
+ num_computed_tokens: The number of tokens that have been computed.
158
+ """
159
+ for manager in self.single_type_managers:
160
+ manager.remove_skipped_blocks(request_id, num_computed_tokens)
161
+
162
+ def get_blocks(self, request_id: str) -> list[list[KVCacheBlock]]:
163
+ """
164
+ Get the blocks for the request.
165
+ """
166
+ return [
167
+ manager.req_to_blocks[request_id]
168
+ for manager in self.single_type_managers
169
+ ]
170
+
171
+ @abstractmethod
172
+ def find_longest_cache_hit(
173
+ self, block_hashes: list[BlockHash],
174
+ max_cache_hit_length: int) -> tuple[list[list[KVCacheBlock]], int]:
175
+ pass
176
+
177
+
178
+class UnitaryKVCacheCoordinator(KVCacheCoordinator):
179
+ """
180
+ KV cache coordinator for models with only one KV cache group. This is the
181
+ case for models with only one KV cache type, e.g., all attention layers use
182
+ full attention or all attention layers use sliding window attention.
183
+ """
184
+
185
+ def __init__(self, kv_cache_config: KVCacheConfig, max_model_len: int,
186
+ use_eagle: bool, enable_caching: bool,
187
+ caching_hash_fn: Callable, enable_kv_cache_events: bool):
188
+ super().__init__(kv_cache_config, max_model_len, use_eagle,
189
+ enable_caching, caching_hash_fn,
190
+ enable_kv_cache_events)
191
+ self.kv_cache_spec = self.kv_cache_config.kv_cache_groups[
192
+ 0].kv_cache_spec
193
+ self.block_size = self.kv_cache_spec.block_size
194
+ assert len(self.kv_cache_config.kv_cache_groups) == 1, (
195
+ "UnitaryKVCacheCoordinator assumes only one kv cache group")
196
+
197
+ def find_longest_cache_hit(
198
+ self, block_hashes: list[BlockHash],
199
+ max_cache_hit_length: int) -> tuple[list[list[KVCacheBlock]], int]:
200
+ hit_blocks = self.single_type_managers[0].find_longest_cache_hit(
201
+ block_hashes=block_hashes,
202
+ max_length=max_cache_hit_length,
203
+ kv_cache_group_ids=[0],
204
+ block_pool=self.block_pool,
205
+ kv_cache_spec=self.kv_cache_spec,
206
+ use_eagle=self.use_eagle,
207
+ )
208
+ return hit_blocks, len(hit_blocks[0]) * self.block_size
209
+
210
+
211
+class HybridKVCacheCoordinator(KVCacheCoordinator):
212
+ """
213
+ KV cache coordinator for hybrid models with multiple KV cache types, and
214
+ thus multiple kv cache groups.
215
+ To simplify `find_longest_cache_hit`, it only supports the combination of
216
+ two types of KV cache groups, and one of them must be full attention.
217
+ May extend to more general cases in the future.
218
+ """
219
+
220
+ def __init__(self, kv_cache_config: KVCacheConfig, max_model_len: int,
221
+ use_eagle: bool, enable_caching: bool,
222
+ caching_hash_fn: Callable, enable_kv_cache_events: bool):
223
+ super().__init__(kv_cache_config, max_model_len, use_eagle,
224
+ enable_caching, caching_hash_fn,
225
+ enable_kv_cache_events)
226
+ self.verify_and_split_kv_cache_groups()
227
+
228
+ def verify_and_split_kv_cache_groups(self) -> None:
229
+ """
230
+ Verifies that the model has exactly two types of KV cache groups, and
231
+ one of them is full attention. Then, split the kv cache groups into full
232
+ attention groups and other groups.
233
+ """
234
+ full_attention_type_id: Optional[str] = None
235
+ other_type_id: Optional[str] = None
236
+ self.full_attention_group_ids: list[int] = []
237
+ self.other_group_ids: list[int] = []
238
+ for i, g in enumerate(self.kv_cache_config.kv_cache_groups):
239
+ if isinstance(g.kv_cache_spec, FullAttentionSpec):
240
+ if full_attention_type_id is None:
241
+ full_attention_type_id = g.kv_cache_spec.type_id
242
+ else:
243
+ assert full_attention_type_id == g.kv_cache_spec.type_id, (
244
+ "HybridKVCacheCoordinator assumes exactly one type of "
245
+ "full attention groups now.")
246
+ self.full_attention_group_ids.append(i)
247
+ else:
248
+ if other_type_id is None:
249
+ other_type_id = g.kv_cache_spec.type_id
250
+ else:
251
+ assert other_type_id == g.kv_cache_spec.type_id, (
252
+ "HybridKVCacheCoordinator assumes "
253
+ "exactly one other type of groups now.")
254
+ self.other_group_ids.append(i)
255
+
256
+ assert full_attention_type_id is not None, (
257
+ "HybridKVCacheCoordinator assumes exactly one type of full "
258
+ "attention groups now.")
259
+ assert other_type_id is not None, (
260
+ "HybridKVCacheCoordinator assumes exactly one type of other "
261
+ "groups now.")
262
+
263
+ self.full_attention_manager_cls = FullAttentionManager
264
+ self.other_attention_cls = self.single_type_managers[
265
+ self.other_group_ids[0]].__class__
266
+
267
+ self.full_attention_spec = self.kv_cache_config.kv_cache_groups[
268
+ self.full_attention_group_ids[0]].kv_cache_spec
269
+ self.other_spec = self.kv_cache_config.kv_cache_groups[
270
+ self.other_group_ids[0]].kv_cache_spec
271
+
272
+ self.full_attention_block_size = self.full_attention_spec.block_size
273
+ self.other_block_size = self.other_spec.block_size
274
+ assert self.other_block_size % self.full_attention_block_size == 0, (
275
+ "KVCacheCoordinator assumes the block_size of full attention "
276
+ "layers is divisible by other layers now.")
277
+
278
+ def find_longest_cache_hit(
279
+ self,
280
+ block_hashes: list[BlockHash],
281
+ max_cache_hit_length: int,
282
+ ) -> tuple[list[list[KVCacheBlock]], int]:
283
+ """
284
+ Find the longest cache hit for the request.
285
+
286
+ Args:
287
+ block_hashes: The block hashes of the request.
288
+ max_cache_hit_length: The maximum length of the cache hit.
289
+
290
+ Returns:
291
+ A tuple containing:
292
+ - A list of the cache hit blocks for each single type manager.
293
+ - The number of tokens of the longest cache hit.
294
+ """
295
+ # First, find the longest cache hit for full attention.
296
+ hit_blocks_full_attn = (
297
+ self.full_attention_manager_cls.find_longest_cache_hit(
298
+ block_hashes=block_hashes,
299
+ max_length=max_cache_hit_length,
300
+ kv_cache_group_ids=self.full_attention_group_ids,
301
+ block_pool=self.block_pool,
302
+ kv_cache_spec=self.full_attention_spec,
303
+ use_eagle=self.use_eagle,
304
+ ))
305
+ hit_length = len(
306
+ hit_blocks_full_attn[0]) * self.full_attention_block_size
307
+
308
+ # Next, find the cache hit for the other attention WITHIN
309
+ # the cache hit of full attention.
310
+ hit_blocks_other_attn = (
311
+ self.other_attention_cls.find_longest_cache_hit(
312
+ block_hashes=block_hashes,
313
+ max_length=hit_length,
314
+ kv_cache_group_ids=self.other_group_ids,
315
+ block_pool=self.block_pool,
316
+ kv_cache_spec=self.other_spec,
317
+ use_eagle=self.use_eagle,
318
+ ))
319
+ hit_length = len(hit_blocks_other_attn[0]) * self.other_block_size
320
+
321
+ # NOTE: the prefix cache hit length must be a multiply of block_size as
322
+ # we don't support partial block cache hit yet. The cache hit length
323
+ # of other attention is ensured to be a multiply of the block size of
324
+ # full attention layers in current implementation, because hit_length is
325
+ # a multiply of other attention's block size, and other attention's
326
+ # block size is a multiply of full attention's block size (verified in
327
+ # `verify_and_split_kv_cache_groups`).
328
+ assert hit_length % self.full_attention_block_size == 0
329
+
330
+ # Truncate the full attention cache hit to the length of the
331
+ # cache hit of the other attention.
332
+ for i in range(len(hit_blocks_full_attn)):
333
+ del hit_blocks_full_attn[i][hit_length //
334
+ self.full_attention_block_size:]
335
+
336
+ # Merge the hit blocks of full attention and other attention.
337
+ hit_blocks = hit_blocks_other_attn
338
+ for group_id, blocks in enumerate(hit_blocks_full_attn):
339
+ # NOTE: there is only one full attention group in most cases. So
340
+ # the time complexity of insert is fine.
341
+ hit_blocks.insert(group_id, blocks)
342
+ return hit_blocks, hit_length
343
+
344
+
345
+def get_kv_cache_coordinator(
346
+ kv_cache_config: KVCacheConfig, max_model_len: int, use_eagle: bool,
347
+ enable_caching: bool, caching_hash_fn: Callable,
348
+ enable_kv_cache_events: bool) -> KVCacheCoordinator:
349
+ if len(kv_cache_config.kv_cache_groups) == 1:
350
+ return UnitaryKVCacheCoordinator(kv_cache_config, max_model_len,
351
+ use_eagle, enable_caching,
352
+ caching_hash_fn,
353
+ enable_kv_cache_events)
354
+ else:
355
+ return HybridKVCacheCoordinator(kv_cache_config, max_model_len,
356
+ use_eagle, enable_caching,
357
+ caching_hash_fn,
358
+ enable_kv_cache_events)
vllm/v1/core/kv_cache_manager.py
+76
-59
@@ -8,11 +8,9 @@ from typing import Optional
8
from vllm.distributed.kv_events import KVCacheEvent
9
from vllm.logger import init_logger
10
from vllm.utils import sha256
11
-from vllm.v1.core.block_pool import BlockPool
11
+from vllm.v1.core.kv_cache_coordinator import get_kv_cache_coordinator
12
from vllm.v1.core.kv_cache_utils import (BlockHash, KVCacheBlock,
13
hash_request_tokens)
14
-from vllm.v1.core.single_type_kv_cache_manager import (
15
- get_manager_for_kv_cache_spec)
14
from vllm.v1.kv_cache_interface import KVCacheConfig
15
from vllm.v1.metrics.stats import PrefixCacheStats
16
from vllm.v1.request import Request, RequestStatus
@@ -22,16 +20,24 @@ logger = init_logger(__name__)
20
21
@dataclass
22
class KVCacheBlocks:
25
- blocks: list[KVCacheBlock]
23
+ """
24
+ The allocation result of KVCacheManager, work as the interface between
25
+ Scheduler and KVCacheManager, to hide KVCacheManager's internal data
26
+ structure from the Scheduler.
27
+ """
28
+ blocks: list[list[KVCacheBlock]]
29
+ """
30
+ blocks[i][j] refers to the i-th kv_cache_group and the j-th block of tokens.
31
+ We don't use block of tokens as the outer dimension because it assumes all
32
+ kv_cache_groups have the same number of blocks, which is true for now but
33
+ will be broken if we want to give different block_size to different
34
+ kv_cache_groups in the future.
35
+ """
36
37
def __add__(self, other: "KVCacheBlocks") -> "KVCacheBlocks":
38
"""Adds two KVCacheBlocks instances."""
29
- return KVCacheBlocks(self.blocks + other.blocks)
30
-
31
- @classmethod
32
- def create_empty(cls) -> "KVCacheBlocks":
33
- """Creates a new KVCacheBlocks instance with no blocks."""
34
- return cls([])
39
+ return KVCacheBlocks(
40
+ [blk1 + blk2 for blk1, blk2 in zip(self.blocks, other.blocks)])
41
42
def get_block_ids(self) -> list[list[int]]:
43
"""
@@ -39,15 +45,20 @@ class KVCacheBlocks:
45
46
Returns:
47
list[list[int]]: A two-level list where
42
- * the outer list corresponds to KV cache groups (only 1 group now)
48
+ * the outer list corresponds to KV cache groups
49
* each inner list contains the block_ids of the blocks in that group
50
"""
45
- return [[block.block_id for block in self.blocks]]
51
+ block_ids = []
52
+ for group in self.blocks:
53
+ block_ids.append([blk.block_id for blk in group])
54
+ return block_ids
55
56
def get_unhashed_block_ids(self) -> list[int]:
57
"""Get block_ids of unhashed blocks from KVCacheBlocks instance."""
58
+ assert len(self.blocks) == 1, "Only one group is supported"
59
return [
50
- block.block_id for block in self.blocks if block.block_hash is None
60
+ block.block_id for block in self.blocks[0]
61
+ if block.block_hash is None
62
]
63
64
@@ -63,12 +74,6 @@ class KVCacheManager:
74
log_stats: bool = False,
75
enable_kv_cache_events: bool = False,
76
) -> None:
66
- assert len(kv_cache_config.kv_cache_groups) == 1, (
67
- "KVCacheManager does not support hybrid models with more than 1 "
68
- "kv cache group")
69
- kv_cache_spec = kv_cache_config.kv_cache_groups[0].kv_cache_spec
70
- self.block_size = kv_cache_spec.block_size
71
- self.num_gpu_blocks = kv_cache_config.num_blocks
77
self.max_model_len = max_model_len
78
79
self.enable_caching = enable_caching
@@ -77,17 +82,24 @@ class KVCacheManager:
82
self.log_stats = log_stats
83
# FIXME: make prefix cache stats conditional on log_stats
84
self.prefix_cache_stats = PrefixCacheStats() if log_stats else None
80
-
81
- self.block_pool = BlockPool(self.num_gpu_blocks, enable_caching,
82
- enable_kv_cache_events)
83
-
84
- self.single_type_manager = get_manager_for_kv_cache_spec(
85
- kv_cache_spec=kv_cache_spec,
86
- block_pool=self.block_pool,
85
+ assert len(
86
+ set(g.kv_cache_spec.block_size
87
+ for g in kv_cache_config.kv_cache_groups)
88
+ ) == 1, "Only one block size is supported for now"
89
+ self.block_size = kv_cache_config.kv_cache_groups[
90
+ 0].kv_cache_spec.block_size
91
+
92
+ self.coordinator = get_kv_cache_coordinator(
93
+ kv_cache_config=kv_cache_config,
94
+ max_model_len=self.max_model_len,
95
use_eagle=self.use_eagle,
88
- num_kv_cache_groups=1,
96
+ enable_caching=enable_caching,
97
caching_hash_fn=self.caching_hash_fn,
98
+ enable_kv_cache_events=enable_kv_cache_events,
99
)
100
+ self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)
101
+ self.block_pool = self.coordinator.block_pool
102
+ self.kv_cache_config = kv_cache_config
103
104
# Mapping from request ID to kv block hashes.
105
# This is to avoid recomputing the block hashes for each call of
@@ -133,7 +145,7 @@ class KVCacheManager:
145
# When the request requires prompt logprobs, we skip prefix caching.
146
if (not self.enable_caching
147
or request.sampling_params.prompt_logprobs is not None):
136
- return KVCacheBlocks.create_empty(), 0
148
+ return self.create_empty_block_list(), 0
149
150
# The block hashes for the request may already be computed
151
# if the scheduler has tried to schedule the request before.
@@ -154,20 +166,16 @@ class KVCacheManager:
166
# num_computed_tokens to be block-size aligned. Removing this limitation
167
# could slightly improve performance in the future.
168
max_cache_hit_length = request.num_tokens - 1
157
-
158
- computed_blocks = self.single_type_manager.find_longest_cache_hit(
159
- block_hashes, max_cache_hit_length)
160
- # NOTE(woosuk): Since incomplete blocks are not eligible for
161
- # sharing, `num_computed_tokens` is always a multiple of
162
- # `block_size`.
163
- num_computed_tokens = len(computed_blocks) * self.block_size
169
+ computed_blocks, num_new_computed_tokens = (
170
+ self.coordinator.find_longest_cache_hit(block_hashes,
171
+ max_cache_hit_length))
172
173
if self.log_stats:
174
assert self.prefix_cache_stats is not None
175
self.prefix_cache_stats.queries += request.num_tokens
168
- self.prefix_cache_stats.hits += num_computed_tokens
176
+ self.prefix_cache_stats.hits += num_new_computed_tokens
177
170
- return KVCacheBlocks(computed_blocks), num_computed_tokens
178
+ return KVCacheBlocks(computed_blocks), num_new_computed_tokens
179
180
def allocate_slots(
181
self,
@@ -220,7 +228,9 @@ class KVCacheManager:
228
if new_computed_blocks is not None:
229
new_computed_block_list = new_computed_blocks.blocks
230
else:
223
- new_computed_block_list = []
231
+ new_computed_block_list = [
232
+ [] for _ in range(len(self.kv_cache_config.kv_cache_groups))
233
+ ]
234
235
# Free the blocks that are skipped during the attention computation
236
# (e.g., tokens outside the sliding window).
@@ -228,8 +238,8 @@ class KVCacheManager:
238
# insufficient free blocks.
239
# Should call this function before allocating new blocks to reduce
240
# the number of evicted blocks.
231
- self.single_type_manager.remove_skipped_blocks(
232
- request.request_id, request.num_computed_tokens)
241
+ self.coordinator.remove_skipped_blocks(request.request_id,
242
+ request.num_computed_tokens)
243
244
# The number of computed tokens is the number of computed tokens plus
245
# the new prefix caching hits
@@ -238,12 +248,12 @@ class KVCacheManager:
248
num_tokens_need_slot = min(
249
num_computed_tokens + num_new_tokens + num_lookahead_tokens,
250
self.max_model_len)
241
- num_blocks_to_allocate = (
242
- self.single_type_manager.get_num_blocks_to_allocate(
243
- request_id=request.request_id,
244
- num_tokens=num_tokens_need_slot,
245
- new_computed_blocks=new_computed_block_list,
246
- ))
251
+
252
+ num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate(
253
+ request_id=request.request_id,
254
+ num_tokens=num_tokens_need_slot,
255
+ new_computed_blocks=new_computed_block_list,
256
+ )
257
258
if num_blocks_to_allocate > self.block_pool.get_num_free_blocks():
259
# Cannot allocate new blocks
@@ -253,16 +263,16 @@ class KVCacheManager:
263
if self.enable_caching:
264
self.block_pool.touch(new_computed_block_list)
265
else:
256
- assert not new_computed_block_list, (
266
+ assert all(not blocks for blocks in new_computed_block_list), (
267
"Computed blocks should be empty when "
268
"prefix caching is disabled")
269
270
# Append the new computed blocks to the request blocks until now to
271
# avoid the case where the new blocks cannot be allocated.
262
- self.single_type_manager.save_new_computed_blocks(
263
- request.request_id, new_computed_block_list)
272
+ self.coordinator.save_new_computed_blocks(request.request_id,
273
+ new_computed_block_list)
274
265
- new_blocks = self.single_type_manager.allocate_new_blocks(
275
+ new_blocks = self.coordinator.allocate_new_blocks(
276
request.request_id, num_tokens_need_slot)
277
278
# P/D: delay caching blocks if we have to recv from
@@ -273,7 +283,7 @@ class KVCacheManager:
283
# Speculated tokens might be rejected in the future, so we does
284
# not cache any speculated tokens. We only cache blocks with
285
# generated (accepted) tokens.
276
- self.single_type_manager.cache_blocks(
286
+ self.coordinator.cache_blocks(
287
request, self.req_to_block_hashes[request.request_id],
288
num_computed_tokens + num_new_tokens - num_draft_tokens)
289
@@ -287,7 +297,7 @@ class KVCacheManager:
297
Args:
298
request: The request to free the blocks.
299
"""
290
- self.single_type_manager.free(request.request_id)
300
+ self.coordinator.free(request.request_id)
301
302
def reset_prefix_cache(self) -> bool:
303
"""Reset prefix cache. This function may be used in RLHF
@@ -345,10 +355,8 @@ class KVCacheManager:
355
group.
356
"""
357
assert request.status == RequestStatus.RUNNING
348
- return [
349
- self.single_type_manager.get_num_common_prefix_blocks(
350
- request.request_id, num_running_requests)
351
- ]
358
+ return self.coordinator.get_num_common_prefix_blocks(
359
+ request.request_id, num_running_requests)
360
361
def free_block_hashes(self, request: Request) -> None:
362
"""Discard the block hashes for the request.
@@ -368,6 +376,15 @@ class KVCacheManager:
376
377
def get_block_ids(self, request_id: str) -> list[list[int]]:
378
"""Get the block ids of a request."""
371
- assert request_id in self.single_type_manager.req_to_blocks
372
- return KVCacheBlocks(self.single_type_manager.req_to_blocks[request_id]
373
- ).get_block_ids()
379
+ return KVCacheBlocks(
380
+ self.coordinator.get_blocks(request_id)).get_block_ids()
381
+
382
+ def cache_blocks(self, request: Request, block_hashes: list[BlockHash],
383
+ num_computed_tokens: int) -> None:
384
+ """Cache the blocks for the request."""
385
+ self.coordinator.cache_blocks(request, block_hashes,
386
+ num_computed_tokens)
387
+
388
+ def create_empty_block_list(self) -> KVCacheBlocks:
389
+ """Creates a new KVCacheBlocks instance with no blocks."""
390
+ return KVCacheBlocks([[] for _ in range(self.num_kv_cache_groups)])
vllm/v1/core/kv_cache_utils.py
+271
-51
@@ -1,8 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""KV-Cache Utilities."""
4
+
5
import os
5
-from collections import deque
6
+from collections import defaultdict, deque
7
from collections.abc import Iterable, Sequence
8
from dataclasses import dataclass
9
from typing import Any, Callable, NamedTuple, Optional
@@ -33,6 +34,18 @@ class BlockHash(NamedTuple):
34
extra_keys: Optional[Any] = None
35
36
37
+class BlockHashWithGroupId(NamedTuple):
38
+ # The hash value for the contents (e.g., token_ids) of a block without group
39
+ # ID. The value is the same for blocks representing the same tokens but for
40
+ # different groups.
41
+ block_hash: BlockHash
42
+ # The KV cache group ID.
43
+ group_id: int
44
+
45
+ def get_hash_value(self) -> int:
46
+ return self.block_hash.hash_value
47
+
48
+
49
# The hash seed for the first block of the prefix block sequence.
50
#
51
# Even if the hash function is the builtin hash(), we use sha256 to generate
@@ -44,7 +57,7 @@ class BlockHash(NamedTuple):
57
# This aligns with the behavior of Python's hash() function, which also uses
58
# a random seed if PYTHONHASHSEED is not set.
59
NONE_HASH = int.from_bytes(os.urandom(32), byteorder="big") if os.getenv(
47
- 'PYTHONHASHSEED') is None else sha256(os.getenv('PYTHONHASHSEED'))
60
+ "PYTHONHASHSEED") is None else sha256(os.getenv("PYTHONHASHSEED"))
61
62
63
class PrefixCachingMetrics:
@@ -118,7 +131,7 @@ class KVCacheBlock:
131
ref_cnt: int = 0
132
# The hash of the block composed of (block hash, tuple of token IDs).
133
# It is only available when the block is full.
121
- _block_hash: Optional[BlockHash] = None
134
+ _block_hash: Optional[BlockHashWithGroupId] = None
135
136
# Used to construct a doubly linked list for free blocks.
137
# These two attributes should only be manipulated by FreeKVCacheBlockQueue.
@@ -135,11 +148,11 @@ class KVCacheBlock:
148
self.ref_cnt -= 1
149
150
@property
138
- def block_hash(self) -> Optional[BlockHash]:
151
+ def block_hash(self) -> Optional[BlockHashWithGroupId]:
152
return self._block_hash
153
154
@block_hash.setter
142
- def block_hash(self, block_hash: BlockHash):
155
+ def block_hash(self, block_hash: BlockHashWithGroupId):
156
assert self.block_hash is None, (
157
"The block already has a hash. This should not happen.")
158
self._block_hash = block_hash
@@ -151,10 +164,10 @@ class KVCacheBlock:
164
def __repr__(self) -> str:
165
# Use block_id instead of KVCacheBlock object to avoid calling __repr__
166
# on KVCacheBlock object recursively.
154
- prev_block_id = self.prev_free_block.block_id \
155
- if self.prev_free_block else None
156
- next_block_id = self.next_free_block.block_id \
157
- if self.next_free_block else None
167
+ prev_block_id = (self.prev_free_block.block_id
168
+ if self.prev_free_block else None)
169
+ next_block_id = (self.next_free_block.block_id
170
+ if self.next_free_block else None)
171
return (f"KVCacheBlock(block_id={self.block_id}, "
172
f"ref_cnt={self.ref_cnt}, "
173
f"_block_hash={self._block_hash}, "
@@ -570,20 +583,20 @@ def create_kv_cache_group_specs(
583
kv_cache_spec: dict[str, KVCacheSpec],
584
grouped_layer_names: list[list[str]]) -> list[KVCacheGroupSpec]:
585
"""
573
- Create KVCacheGroupSpec object for each kv cache group layer.
574
- The layers in the same group should share the same
575
- KVCacheSpec.
576
-
577
- Args:
578
- kv_cache_spec:
579
- A mapping from each layer name to its corresponding KVCacheSpec.
580
- grouped_layer_names:
581
- A list of kv cache groups, where each element is a list of layer
582
- names that belong to the same group and should share the same
583
- KVCacheSpec.
584
- Returns:
585
- A list of KVCacheGroupSpec objects, one for each group.
586
- """
586
+ Create KVCacheGroupSpec object for each kv cache group layer.
587
+ The layers in the same group should share the same
588
+ KVCacheSpec.
589
+
590
+ Args:
591
+ kv_cache_spec:
592
+ A mapping from each layer name to its corresponding KVCacheSpec.
593
+ grouped_layer_names:
594
+ A list of kv cache groups, where each element is a list of layer
595
+ names that belong to the same group and should share the same
596
+ KVCacheSpec.
597
+ Returns:
598
+ A list of KVCacheGroupSpec objects, one for each group.
599
+ """
600
kv_cache_groups = []
601
for layer_names_one_group in grouped_layer_names:
602
layer_specs = [
@@ -628,6 +641,37 @@ def get_max_concurrency_for_kv_cache_config(
641
return max_concurrency
642
643
644
+def get_num_blocks(vllm_config: VllmConfig, num_layers: int,
645
+ available_memory: int, page_size: int) -> int:
646
+ """
647
+ Get the number of kv cache blocks.
648
+
649
+ Args:
650
+ vllm_config: The global VllmConfig
651
+ num_layers: The number of layers
652
+ available_memory: Memory available for KV cache in bytes.
653
+ page_size: The page size of the KV cache.
654
+ """
655
+ num_blocks = int(available_memory // page_size // num_layers)
656
+ num_blocks = max(num_blocks, 0)
657
+ if vllm_config.cache_config.num_gpu_blocks_override is not None:
658
+ num_gpu_blocks_override = \
659
+ vllm_config.cache_config.num_gpu_blocks_override
660
+ logger.info(
661
+ "Overriding num_gpu_blocks=%d with "
662
+ "num_gpu_blocks_override=%d", num_blocks, num_gpu_blocks_override)
663
+ return num_blocks
664
+
665
+
666
+def get_uniform_page_size(kv_cache_spec: dict[str, KVCacheSpec]) -> int:
667
+ """
668
+ Get the page size of the KV cache.
669
+ """
670
+ page_sizes = set(layer.page_size_bytes for layer in kv_cache_spec.values())
671
+ assert len(page_sizes) == 1
672
+ return page_sizes.pop()
673
+
674
+
675
def _get_kv_cache_config_uniform_type(vllm_config: VllmConfig,
676
kv_cache_spec: dict[str, KVCacheSpec],
677
available_memory: int) -> KVCacheConfig:
@@ -644,32 +688,24 @@ def _get_kv_cache_config_uniform_type(vllm_config: VllmConfig,
688
The generated KVCacheConfig
689
"""
690
647
- page_sizes = {layer.page_size_bytes for layer in kv_cache_spec.values()}
648
- assert len(page_sizes) == 1
649
- page_size = page_sizes.pop()
650
-
651
- num_blocks = int(available_memory // page_size // len(kv_cache_spec))
652
- num_blocks = max(num_blocks, 0)
653
-
654
- if vllm_config.cache_config.num_gpu_blocks_override is not None:
655
- num_gpu_blocks_override = \
656
- vllm_config.cache_config.num_gpu_blocks_override
657
- logger.info(
658
- "Overriding num_gpu_blocks=%d with "
659
- "num_gpu_blocks_override=%d", num_blocks, num_gpu_blocks_override)
660
- num_blocks = num_gpu_blocks_override
691
+ page_size = get_uniform_page_size(kv_cache_spec)
692
+ num_blocks = get_num_blocks(vllm_config, len(kv_cache_spec),
693
+ available_memory, page_size)
694
695
per_layer_size = page_size * num_blocks
696
# All layers have the same KV cache spec, so we create one kv cache group
697
# for all layers.
698
grouped_layer_names = [list(kv_cache_spec.keys())]
699
700
+ # Each layer uses a separate Tensor to store its KV cache.
701
+ kv_cache_tensors = [
702
+ KVCacheTensor(size=per_layer_size, shared_by=[layer_name])
703
+ for layer_name in kv_cache_spec
704
+ ]
705
+
706
kv_cache_config = KVCacheConfig(
707
num_blocks=num_blocks,
669
- tensors={
670
- layer_name: KVCacheTensor(size=per_layer_size)
671
- for layer_name in kv_cache_spec
672
- },
708
+ kv_cache_tensors=kv_cache_tensors,
709
kv_cache_groups=create_kv_cache_group_specs(kv_cache_spec,
710
grouped_layer_names),
711
)
@@ -685,17 +721,185 @@ def _get_kv_cache_config_uniform_type(vllm_config: VllmConfig,
721
return kv_cache_config
722
723
724
+def is_kv_cache_page_size_uniform(
725
+ kv_cache_spec: dict[str, KVCacheSpec]) -> bool:
726
+ """
727
+ Whether all layers in the given KVCacheSpec have the same page size.
728
+ Args:
729
+ kv_cache_spec: The KVCacheSpec of each attention layer in the model
730
+
731
+ Returns:
732
+ True if all layers have the same page size, False otherwise.
733
+ """
734
+
735
+ page_sizes = {layer.page_size_bytes for layer in kv_cache_spec.values()}
736
+ return len(page_sizes) == 1
737
+
738
+
739
+def _get_kv_cache_config_uniform_page_size(
740
+ vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec],
741
+ available_memory: int) -> KVCacheConfig:
742
+ """
743
+ Generates the KV cache configuration for hybrid models with multiple
744
+ attention types but still with a uniform page size (physical memory per
745
+ block per layer) for all layers.
746
+
747
+ Detailed explanation about kv cache management of hybrid models:
748
+ The layers in the models are repeated with some patterns, e.g., a model
749
+ with 10 full attention layers and 20 sliding window attention layers can be
750
+ regarded as repeating the pattern (1 * full, 2 * sw) 10 times.
751
+ The KVCacheManager allocates different block tables for each of the 3 layers
752
+ in the pattern, and repeats each of them 10 times to generate the
753
+ block_table for the 30 layers in the model.
754
+ Therefore, we can group the layers in the model into 3 kv_cache_groups, each
755
+ of which contains 10 layers in the model.
756
+ The KVCacheManager allocates the block_table for each group based on its
757
+ kv_cache spec, and the model runner applies the block table to each layer
758
+ in the group.
759
+ For example:
760
+ 1. A model only uses full attention. The pattern is
761
+ (num_hidden_layers * full), so there is only one group and the block table
762
+ is shared by all layers. It is already handled by
763
+ `_get_kv_cache_config_uniform_type`.
764
+ 2. A model with 10 full attention layers and 20 sliding window
765
+ attention layers. There are 3 layers in the pattern (1 * full, 2 * sw), so
766
+ there are 3 kv_cache_groups, each of which represents 10 layers.
767
+
768
+ To simplify the implementation, we make the following assumptions:
769
+ 1. Physical memory per block: Must be the same across all KV cache groups.
770
+ Breaking this assumption is non-trivial due to memory fragmentation concerns
771
+ when allocating blocks of different sizes.
772
+ 2. Tokens per block (block_size): Currently, we directly use
773
+ `CacheConfig.block_size` for all layers. It can be extended to vary by KV
774
+ cache group, but within each KV cache group, all layers must share the same
775
+ block size.
776
+ 3. Physical memory per token per layer: This property is decided by model
777
+ config. Currently we only support models that have the same physical memory
778
+ per token per layer for all layers. Can be relaxed with a simple extension,
779
+ but still need to keep physical memory per block the same for all groups.
780
+ 4. Number of layers per group: Currently assumed the same for all layers.
781
+ Can be relaxed with a simple extension, but still need to keep physical
782
+ memory per block the same for all groups.
783
+ 5. Attention type within groups: All layers in a group must share the same
784
+ attention type. One exception is that, when
785
+ `--disable-hybrid-kv-cache-manager` is true, the single group for full
786
+ attention layers may also include attention layers using sliding window or
787
+ LLaMA 4 local attention. See `unify_hybrid_kv_cache_specs` for more details.
788
+ 6. Support for multiple attention types: The design for most components is
789
+ general to an arbitrary number of attention types. But
790
+ `find_longest_cache_hit` only supports one attention type or two
791
+ types of full-attention plus exactly one another type. The general
792
+ implementation of this function is feasible but we don't know how to
793
+ implement it cleanly yet.
794
+
795
+ As we assume tokens per block, physical memory per token per layer, and
796
+ number of layers per group are the same now, we can ensure that physical
797
+ memory per block is the same for all groups.
798
+
799
+ Args:
800
+ vllm_config: The global VllmConfig
801
+ kv_cache_spec: The KVCacheSpec of each attention layer in the model
802
+ available_memory: Memory available for KV cache in bytes.
803
+ Returns:
804
+ The generated KVCacheConfig
805
+ """
806
+ # Group all layers by type_id.
807
+ # E.g., 2 full attention layers and 3 sliding window attention layers,
808
+ # -> (full.0, full.1), (sw.0, sw.1, sw.2).
809
+ same_type_layers: dict[str, list[str]] = defaultdict(list)
810
+ for layer_name, layer_spec in kv_cache_spec.items():
811
+ same_type_layers[layer_spec.type_id].append(layer_name)
812
+
813
+ # Split each group into smaller groups, to make the number of layers in each
814
+ # group identical. Add padding to the last group of each type if necessary.
815
+ # E.g., (full.0, full.1), (sw.0, sw.1, sw.2)
816
+ # split to 3 groups with 2 layers each:
817
+ # (full.0, full.1), (sw.0, sw.1), (sw.2, padding).
818
+ # FIXME(Chen): At the moment of writing this code (2025-06-02), all
819
+ # open-source hybrid model follows a n:1 pattern between different attention
820
+ # types (e.g., Gemma3 5:1 between sw and full, LLaMA4 3:1 between local and
821
+ # full), so we can use the "1" in the n:1 pattern as the group size, which
822
+ # is the minimum number of layers among all attention types. Need a better
823
+ # strategy if we want to support more complex patterns (e.g., 20 full + 30
824
+ # sw, where the group size should be 10).
825
+ group_size = min([len(layers) for layers in same_type_layers.values()])
826
+ grouped_layers = []
827
+ for layers in same_type_layers.values():
828
+ num_padding_layers = group_size - len(layers) % group_size
829
+ if num_padding_layers != group_size:
830
+ logger.warning(
831
+ "Add %d padding layers, may waste at most %.2f%% KV cache memory", # noqa
832
+ num_padding_layers,
833
+ num_padding_layers / len(layers) * 100,
834
+ )
835
+ for i in range(0, len(layers), group_size):
836
+ grouped_layers.append(layers[i:i + group_size])
837
+ kv_cache_groups = create_kv_cache_group_specs(kv_cache_spec,
838
+ grouped_layers)
839
+
840
+ # Determine how model runners should initialize the KV cache tensors.
841
+ # We will have group_size memory pools, each is shared by one layer from
842
+ # each group. As layers of different groups have different block table,
843
+ # they will use different parts of the shared Tensor.
844
+ # The memory layout in the example will be:
845
+ # full.0, sw.0, sw.2: share a Tensor with size=available_memory//2
846
+ # full.1, sw.1: share another Tensor with size=available_memory//2
847
+ page_size = get_uniform_page_size(kv_cache_spec)
848
+ num_blocks = get_num_blocks(vllm_config, group_size, available_memory,
849
+ page_size)
850
+ per_memory_pool_size = page_size * num_blocks
851
+ kv_cache_tensors = []
852
+ for i in range(group_size):
853
+ shared_by = []
854
+ for j in range(len(kv_cache_groups)):
855
+ if i < len(grouped_layers[j]):
856
+ shared_by.append(grouped_layers[j][i])
857
+ kv_cache_tensors.append(
858
+ KVCacheTensor(size=per_memory_pool_size, shared_by=shared_by))
859
+
860
+ kv_cache_config = KVCacheConfig(
861
+ num_blocks=num_blocks,
862
+ kv_cache_tensors=kv_cache_tensors,
863
+ kv_cache_groups=kv_cache_groups,
864
+ )
865
+
866
+ # Print the KV cache size and maximum concurrency.
867
+ num_tokens = num_blocks // len(
868
+ grouped_layers) * vllm_config.cache_config.block_size
869
+ num_tokens_str = f"{num_tokens:,}"
870
+ logger.info("GPU KV cache size: %s tokens", num_tokens_str)
871
+ max_model_len_str = f"{vllm_config.model_config.max_model_len:,}"
872
+ max_concurrency = get_max_concurrency_for_kv_cache_config(
873
+ vllm_config, kv_cache_config)
874
+ logger.info("Maximum concurrency for %s tokens per request: %.2fx",
875
+ max_model_len_str, max_concurrency)
876
+ return kv_cache_config
877
+
878
+
879
def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
880
"""
690
- Only models with one type of KV cache are supported yet. This function tries
691
- to convert the KV cache specs to one type if the model is a hybrid model
692
- with multiple type of KV cache. It will convert all SlidingWindowSpec to
693
- FullAttentionSpec if both types are present.
881
+ This function tries to convert the KV cache specs to one type if the model
882
+ is a hybrid model with multiple type of KV cache. It will convert all
883
+ SlidingWindowSpec to FullAttentionSpec if both types are present.
884
885
Args:
886
kv_cache_spec: The kv cache spec of each attention layer in the model
887
"""
888
889
+ def is_hybrid(kv_cache_spec: dict[str, KVCacheSpec]) -> bool:
890
+ type_ids = set(layer_spec.type_id
891
+ for layer_spec in kv_cache_spec.values())
892
+ return len(type_ids) > 1
893
+
894
+ if not is_hybrid(kv_cache_spec):
895
+ return
896
+
897
+ logger.warning(
898
+ "Hybrid KV cache manager is disabled for this hybrid model, "
899
+ "This means we do not enable any optimizations for saving KV cache "
900
+ "memory (e.g., dropping the KV cache outside the sliding window). "
901
+ "The compute of layers like sliding window is still saved.")
902
+
903
has_full_attention = any(
904
isinstance(spec, FullAttentionSpec) for spec in kv_cache_spec.values())
905
has_sliding_window = any(
@@ -712,13 +916,18 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
916
sliding_window=spec.sliding_window,
917
)
918
919
+ if is_hybrid(kv_cache_spec):
920
+ raise ValueError("Hybrid KV cache manager is disabled but failed to "
921
+ "convert the KV cache specs to one unified type.")
922
+
923
716
-def get_kv_cache_config(vllm_config: VllmConfig,
717
- kv_cache_spec: dict[str, KVCacheSpec],
718
- available_memory: int) -> KVCacheConfig:
924
+def get_kv_cache_config(
925
+ vllm_config: VllmConfig,
926
+ kv_cache_spec: dict[str, KVCacheSpec],
927
+ available_memory: int,
928
+) -> KVCacheConfig:
929
"""
720
- Generates the KV cache configuration for a model
721
- TODO: support hybrid models with more than one type of KV cache.
930
+ Generates the KV cache configuration for a model.
931
932
Args:
933
vllm_config: The global VllmConfig
@@ -728,14 +937,25 @@ def get_kv_cache_config(vllm_config: VllmConfig,
937
Returns:
938
The generated KVCacheConfigs
939
"""
731
- unify_hybrid_kv_cache_specs(kv_cache_spec)
940
check_enough_kv_cache_memory(vllm_config, kv_cache_spec, available_memory)
941
+
942
+ if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
943
+ unify_hybrid_kv_cache_specs(kv_cache_spec)
944
+
945
if is_kv_cache_type_uniform(kv_cache_spec):
946
# KV cache of all layers are the same, which is true for
947
# most models. Allocate the same amount of memory for
948
# each layer.
949
return _get_kv_cache_config_uniform_type(vllm_config, kv_cache_spec,
950
available_memory)
951
+ elif is_kv_cache_page_size_uniform(kv_cache_spec):
952
+ # Model contains multiple attention types, but KV cache of all layers
953
+ # have the same physical memory per block per layer. Split the layers
954
+ # into groups with the same number of layers, and thus same total page
955
+ # size.
956
+ return _get_kv_cache_config_uniform_page_size(vllm_config,
957
+ kv_cache_spec,
958
+ available_memory)
959
960
raise NotImplementedError
961
vllm/v1/core/sched/scheduler.py
+4
-3
@@ -18,7 +18,7 @@ from vllm.logger import init_logger
18
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
19
from vllm.v1.core.encoder_cache_manager import (EncoderCacheManager,
20
compute_encoder_budget)
21
-from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
21
+from vllm.v1.core.kv_cache_manager import KVCacheManager
22
from vllm.v1.core.sched.interface import SchedulerInterface
23
from vllm.v1.core.sched.output import (CachedRequestData, NewRequestData,
24
SchedulerOutput)
@@ -377,7 +377,8 @@ class Scheduler(SchedulerInterface):
377
# KVTransfer: WAITING reqs have num_computed_tokens > 0
378
# after async KV recvs are completed.
379
else:
380
- new_computed_blocks = KVCacheBlocks.create_empty()
380
+ new_computed_blocks = (
381
+ self.kv_cache_manager.create_empty_block_list())
382
num_new_local_computed_tokens = 0
383
num_computed_tokens = request.num_computed_tokens
384
@@ -1010,7 +1011,7 @@ class Scheduler(SchedulerInterface):
1011
num_computed_tokens = len(block_ids) * self.block_size
1012
if num_computed_tokens == request.num_tokens:
1013
num_computed_tokens -= 1
1013
- self.kv_cache_manager.single_type_manager.cache_blocks(
1014
+ self.kv_cache_manager.cache_blocks(
1015
request,
1016
self.kv_cache_manager.req_to_block_hashes[request.request_id],
1017
num_computed_tokens,
vllm/v1/core/single_type_kv_cache_manager.py
+98
-53
@@ -22,8 +22,7 @@ class SingleTypeKVCacheManager(ABC):
22
self,
23
kv_cache_spec: KVCacheSpec,
24
block_pool: BlockPool,
25
- use_eagle: bool,
26
- num_kv_cache_groups: int,
25
+ kv_cache_group_id: int,
26
caching_hash_fn: Callable,
27
) -> None:
28
"""
@@ -31,9 +30,7 @@ class SingleTypeKVCacheManager(ABC):
30
Args:
31
kv_cache_spec: The kv_cache_spec for this manager.
32
block_pool: The block pool.
34
- use_eagle: Whether to use eagle.
35
- num_kv_cache_groups: The number of kv cache groups managed by this
36
- manager.
33
+ kv_cache_group_id: The id of the kv cache group of this manager.
34
caching_hash_fn: The caching hash function.
35
"""
36
@@ -41,9 +38,6 @@ class SingleTypeKVCacheManager(ABC):
38
self.kv_cache_spec = kv_cache_spec
39
self.block_pool = block_pool
40
44
- # Needs special handling for find_longest_cache_hit if eagle is enabled
45
- self.use_eagle = use_eagle
46
-
41
# Mapping from request ID to blocks to track the blocks allocated
42
# for each request, so that we can free the blocks when the request
43
# is finished.
@@ -56,8 +50,8 @@ class SingleTypeKVCacheManager(ABC):
50
# data for reempted ones.
51
self.num_cached_block: dict[str, int] = {}
52
59
- self.num_kv_cache_groups = num_kv_cache_groups
53
self.caching_hash_fn = caching_hash_fn
54
+ self.kv_cache_group_id = kv_cache_group_id
55
56
def get_num_blocks_to_allocate(
57
self, request_id: str, num_tokens: int,
@@ -86,8 +80,7 @@ class SingleTypeKVCacheManager(ABC):
80
num_evictable_computed_blocks = sum(
81
blk.ref_cnt == 0 and not blk.is_null
82
for blk in new_computed_blocks)
89
- return ((num_new_blocks + num_evictable_computed_blocks) *
90
- self.num_kv_cache_groups)
83
+ return num_new_blocks + num_evictable_computed_blocks
84
85
def save_new_computed_blocks(
86
self, request_id: str,
@@ -130,8 +123,7 @@ class SingleTypeKVCacheManager(ABC):
123
if num_new_blocks <= 0:
124
return []
125
else:
133
- new_blocks = self.block_pool.get_new_blocks(
134
- num_new_blocks * self.num_kv_cache_groups)
126
+ new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
127
req_blocks.extend(new_blocks)
128
return new_blocks
129
@@ -156,12 +148,19 @@ class SingleTypeKVCacheManager(ABC):
148
num_cached_blocks=num_cached_blocks,
149
num_full_blocks=num_full_blocks,
150
block_size=self.block_size,
151
+ kv_cache_group_id=self.kv_cache_group_id,
152
hash_fn=self.caching_hash_fn,
153
)
154
155
self.num_cached_block[request.request_id] = num_full_blocks
156
157
def free(self, request_id: str) -> None:
158
+ """
159
+ Free the blocks for the request.
160
+
161
+ Args:
162
+ request_id: The request ID.
163
+ """
164
# Default to [] in case a request is freed (aborted) before alloc.
165
req_blocks = self.req_to_blocks.pop(request_id, [])
166
@@ -188,12 +187,22 @@ class SingleTypeKVCacheManager(ABC):
187
188
raise NotImplementedError
189
190
+ @classmethod
191
@abstractmethod
192
- def find_longest_cache_hit(self, block_hashes: list[BlockHash],
193
- max_length: int) -> list[KVCacheBlock]:
192
+ def find_longest_cache_hit(
193
+ cls,
194
+ block_hashes: list[BlockHash],
195
+ max_length: int,
196
+ kv_cache_group_ids: list[int],
197
+ block_pool: BlockPool,
198
+ kv_cache_spec: KVCacheSpec,
199
+ use_eagle: bool,
200
+ ) -> list[list[KVCacheBlock]]:
201
"""
202
Get the longest cache hit prefix of the blocks that is not longer than
196
- `max_length`. If no cache hit is found, return an empty list.
203
+ `max_length`. The prefix should be a common prefix hit for all the
204
+ kv cache groups in `kv_cache_group_ids`. If no cache hit is found,
205
+ return an empty list.
206
If eagle is enabled, drop the last matched block to force recompute the
207
last block to get the required hidden states for eagle drafting head.
208
Need to be customized for each attention type.
@@ -201,12 +210,20 @@ class SingleTypeKVCacheManager(ABC):
210
Args:
211
block_hashes: The block hashes of the request.
212
max_length: The maximum length of the cache hit prefix.
213
+ kv_cache_group_ids: The ids of the kv cache groups.
214
+ block_pool: The block pool.
215
+ kv_cache_spec: The kv cache spec.
216
+ use_eagle: Whether to use eagle.
217
218
Returns:
206
- A list of cached blocks with skipped blocks replaced by null block.
219
+ A list of cached blocks with skipped blocks replaced by null block
220
+ for each kv cache group in `kv_cache_group_ids`.
221
+ Return a list of length `len(kv_cache_group_ids)`, where the i-th
222
+ element is a list of cached blocks for the i-th kv cache group
223
+ in `kv_cache_group_ids`.
224
For example, sliding window manager should return a list like
208
- [NULL, NULL, KVCacheBlock(7), KVCacheBlock(8)] for block size 4 and
209
- sliding window 8.
225
+ [[NULL, NULL, KVCacheBlock(7), KVCacheBlock(8)]] for block size 4
226
+ and sliding window 8 and len(kv_cache_group_ids) = 1.
227
"""
228
229
raise NotImplementedError
@@ -215,11 +232,9 @@ class SingleTypeKVCacheManager(ABC):
232
def remove_skipped_blocks(self, request_id: str,
233
num_computed_tokens: int) -> None:
234
"""
218
- Remove the blocks that are no longer needed from `blocks`. The removed
219
- blocks should be replaced by null_block. Return the removed blocks in
220
- eviction order, where the first returned block should be evicted first.
221
- Don't free the removed blocks in this function. Need to be customized
222
- for each attention type.
235
+ Remove the blocks that are no longer needed from `blocks` and free the
236
+ blocks. The removed blocks should be replaced by null_block.
237
+ Need to be customized for each attention type.
238
239
Args:
240
request_id: The request ID.
@@ -230,21 +245,36 @@ class SingleTypeKVCacheManager(ABC):
245
246
class FullAttentionManager(SingleTypeKVCacheManager):
247
233
- def find_longest_cache_hit(self, block_hashes: list[BlockHash],
234
- max_length: int) -> list[KVCacheBlock]:
235
- computed_blocks: list[KVCacheBlock] = []
236
- max_num_blocks = max_length // self.block_size
248
+ @classmethod
249
+ def find_longest_cache_hit(
250
+ cls,
251
+ block_hashes: list[BlockHash],
252
+ max_length: int,
253
+ kv_cache_group_ids: list[int],
254
+ block_pool: BlockPool,
255
+ kv_cache_spec: KVCacheSpec,
256
+ use_eagle: bool,
257
+ ) -> list[list[KVCacheBlock]]:
258
+ assert isinstance(kv_cache_spec, FullAttentionSpec), (
259
+ "FullAttentionManager can only be used for full attention groups")
260
+ computed_blocks: list[list[KVCacheBlock]] = [
261
+ [] for _ in range(len(kv_cache_group_ids))
262
+ ]
263
+ max_num_blocks = max_length // kv_cache_spec.block_size
264
for i in range(max_num_blocks):
265
block_hash = block_hashes[i]
266
# block_hashes is a chain of block hashes. If a block hash is not
267
# in the cached_block_hash_to_id, the following block hashes are
268
# not computed yet for sure.
242
- if cached_block := self.block_pool.get_cached_block(block_hash):
243
- computed_blocks.append(cached_block)
269
+ if cached_block := block_pool.get_cached_block(
270
+ block_hash, kv_cache_group_ids):
271
+ for j in range(len(kv_cache_group_ids)):
272
+ computed_blocks[j].append(cached_block[j])
273
else:
274
break
246
- if self.use_eagle and len(computed_blocks) > 0:
247
- computed_blocks.pop()
275
+ if use_eagle and len(computed_blocks[0]) > 0:
276
+ for j in range(len(kv_cache_group_ids)):
277
+ computed_blocks[j].pop()
278
return computed_blocks
279
280
def remove_skipped_blocks(self, request_id: str,
@@ -267,45 +297,58 @@ class FullAttentionManager(SingleTypeKVCacheManager):
297
class SlidingWindowManager(SingleTypeKVCacheManager):
298
299
def __init__(self, kv_cache_spec: SlidingWindowSpec, block_pool: BlockPool,
270
- use_eagle: bool, **kwargs) -> None:
271
- super().__init__(kv_cache_spec, block_pool, use_eagle, **kwargs)
300
+ **kwargs) -> None:
301
+ super().__init__(kv_cache_spec, block_pool, **kwargs)
302
self.sliding_window = kv_cache_spec.sliding_window
303
+ self._null_block = block_pool.null_block
304
+
305
+ @classmethod
306
+ def find_longest_cache_hit(
307
+ cls,
308
+ block_hashes: list[BlockHash],
309
+ max_length: int,
310
+ kv_cache_group_ids: list[int],
311
+ block_pool: BlockPool,
312
+ kv_cache_spec: KVCacheSpec,
313
+ use_eagle: bool,
314
+ ) -> list[list[KVCacheBlock]]:
315
+ assert isinstance(kv_cache_spec, SlidingWindowSpec), (
316
+ "SlidingWindowManager can only be used for sliding window groups")
317
+
318
# The number of contiguous blocks needed for prefix cache hit.
319
# -1 since the input token itself is also included in the window
275
- self.sliding_window_contiguous_blocks = cdiv(
276
- (kv_cache_spec.sliding_window - 1), self.block_size)
277
- if self.use_eagle:
320
+ sliding_window_contiguous_blocks = cdiv(
321
+ kv_cache_spec.sliding_window - 1, kv_cache_spec.block_size)
322
+ if use_eagle:
323
# Need to drop the last matched block if eagle is enabled. For
324
# sliding window layer, we achieve this by increasing the number of
325
# contiguous blocks needed for prefix cache hit by one and dropping
326
# the last matched block.
282
- self.sliding_window_contiguous_blocks += 1
283
- self._null_block = block_pool.null_block
327
+ sliding_window_contiguous_blocks += 1
328
285
- def find_longest_cache_hit(self, block_hashes: list[BlockHash],
286
- max_length: int) -> list[KVCacheBlock]:
329
# TODO: reduce i by sliding_window_contiguous_blocks when cache miss, to
330
# optimize the time complexity from O(max_num_blocks) to
331
# O(max_num_blocks / sliding_window_contiguous_blocks +
332
# sliding_window_contiguous_blocks),
333
# which is good for low cache hit rate scenarios.
292
- max_num_blocks = max_length // self.block_size
293
- computed_blocks = [self._null_block] * max_num_blocks
334
+ max_num_blocks = max_length // kv_cache_spec.block_size
335
+ computed_blocks = [[block_pool.null_block] * max_num_blocks
336
+ for _ in range(len(kv_cache_group_ids))]
337
num_contiguous_blocks = 0
295
-
338
match_found = False
339
# Search from right to left and early stop when a match is found.
340
for i in range(max_num_blocks - 1, -1, -1):
299
- if cached_block := self.block_pool.get_cached_block(
300
- block_hashes[i]):
301
- computed_blocks[i] = cached_block
341
+ if cached_block := block_pool.get_cached_block(
342
+ block_hashes[i], kv_cache_group_ids):
343
+ for j in range(len(kv_cache_group_ids)):
344
+ computed_blocks[j][i] = cached_block[j]
345
num_contiguous_blocks += 1
303
- if (num_contiguous_blocks
304
- >= self.sliding_window_contiguous_blocks):
346
+ if (num_contiguous_blocks >= sliding_window_contiguous_blocks):
347
# Trim the trailing blocks.
348
# E.g., [NULL, NULL, 8, 3, NULL, 9] -> [NULL, NULL, 8, 3]
349
# when sliding_window_contiguous_blocks=2.
308
- del computed_blocks[i + num_contiguous_blocks:]
350
+ for j in range(len(kv_cache_group_ids)):
351
+ del computed_blocks[j][i + num_contiguous_blocks:]
352
match_found = True
353
break
354
else:
@@ -313,9 +356,11 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
356
if not match_found:
357
# The first `num_contiguous_blocks` is a cache hit even if
358
# `num_contiguous_blocks < sliding_window_contiguous_blocks`.
316
- del computed_blocks[num_contiguous_blocks:]
317
- if self.use_eagle and len(computed_blocks) > 0:
318
- computed_blocks.pop()
359
+ for j in range(len(kv_cache_group_ids)):
360
+ del computed_blocks[j][num_contiguous_blocks:]
361
+ if use_eagle and len(computed_blocks[0]) > 0:
362
+ for j in range(len(kv_cache_group_ids)):
363
+ computed_blocks[j].pop()
364
return computed_blocks
365
366
def remove_skipped_blocks(self, request_id: str,
vllm/v1/kv_cache_interface.py
+9
-24
@@ -157,11 +157,10 @@ class SlidingWindowSpec(AttentionSpec):
157
@dataclass
158
class KVCacheTensor:
159
"""
160
- A dataclass for specifying how the workers should initialize the KV cache
161
- for a layer. Only contains the size of KV cache for that layer for now. Will
162
- be extended to support multiple layers sharing the same memory pool.
160
+ A class for specifying how the workers should initialize the KV cache.
161
"""
164
- size: int # The size of KV cache Tensor in bytes
162
+ size: int # size of the KV cache tensor in bytes
163
+ shared_by: list[str] # layer names that share the same KV cache tensor
164
165
166
@dataclass
@@ -183,27 +182,13 @@ class KVCacheConfig:
182
"""
183
"""The number of KV cache blocks"""
184
num_blocks: int
186
- """layer_name -> how to initialize KV cache for that layer"""
187
- tensors: dict[str, KVCacheTensor]
185
+ """How should model runner initialize the KV cache tensors for each layer"""
186
+ kv_cache_tensors: list[KVCacheTensor]
187
"""
188
The kv cache groups of the model.
190
- The layers in the models are repeated with some patterns, e.g., a model
191
- with 10 full attention layers and 20 sliding window attention layers can be
192
- regarded as repeating the pattern (1 * full, 2 * sw) 10 times.
193
- The KVCacheManager allocates different block tables for each of the 3 layers
194
- in the pattern, and repeats each of them 10 times to generate the
195
- block_table for the 30 layers in the model.
196
- Therefore, we can group the layers in the model into 3 groups, each of which
197
- contains 10 layers in the model.
198
- The KVCacheManager allocates the block_table for each group based on its
199
- kv_cache spec, and the model runner applies the block table to each layer
200
- in the group.
201
- For example:
202
- 1. A model only uses full attention. The pattern is
203
- (num_hidden_layers * full), so there is only one group and the block table
204
- is shared by all layers.
205
- 2. (WIP) A model with 10 full attention layers and 20 sliding window
206
- attention layers. There are 3 layers in the pattern (1 * full, 2 * sw), so
207
- there are 3 groups, each of which represents 10 layers in the model.
189
+ For models with only one type of attention, there is only one group that
190
+ contains all layers.
191
+ For models with multiple types of attention, there will be multiple groups,
192
+ see `_get_kv_cache_config_uniform_page_size` for more details.
193
"""
194
kv_cache_groups: list[KVCacheGroupSpec]
vllm/v1/worker/gpu_model_runner.py
+87
-33
@@ -2088,33 +2088,58 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2088
block_sizes=block_sizes,
2089
)
2090
2091
- def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
2091
+ def _allocate_kv_cache_tensors(
2092
+ self, kv_cache_config: KVCacheConfig) -> dict[str, torch.Tensor]:
2093
"""
2093
- Initialize KV cache based on `kv_cache_config`.
2094
+ Initializes the KV cache buffer with the correct size. The buffer needs
2095
+ to be reshaped to the desired shape before being used by the models.
2096
+
2097
Args:
2095
- kv_cache_config: Configuration for the KV cache, including the KV
2096
- cache size of each layer
2098
+ kv_cache_config: The KV cache config
2099
+ Returns:
2100
+ dict[str, torch.Tensor]: A map between layer names to their
2101
+ corresponding memory buffer for KV cache.
2102
+ """
2103
+ kv_cache_raw_tensors: dict[str, torch.Tensor] = {}
2104
+ for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
2105
+ tensor = torch.zeros(kv_cache_tensor.size,
2106
+ dtype=torch.int8,
2107
+ device=self.device)
2108
+ for layer_name in kv_cache_tensor.shared_by:
2109
+ kv_cache_raw_tensors[layer_name] = tensor
2110
+
2111
+ layer_names = set()
2112
+ for group in kv_cache_config.kv_cache_groups:
2113
+ layer_names.update(group.layer_names)
2114
+ assert layer_names == set(kv_cache_raw_tensors.keys(
2115
+ )), "Some layers are not correctly initialized"
2116
+ return kv_cache_raw_tensors
2117
+
2118
+ def _reshape_kv_cache_tensors(
2119
+ self,
2120
+ kv_cache_config: KVCacheConfig,
2121
+ kv_cache_raw_tensors: dict[str, torch.Tensor],
2122
+ ) -> dict[str, torch.Tensor]:
2123
"""
2098
- self.kv_cache_config = kv_cache_config
2099
- self.may_reinitialize_input_batch(kv_cache_config)
2100
- self.initialize_attn_backend(kv_cache_config)
2124
+ Reshape the KV cache tensors to the desired shape and dtype.
2125
2126
+ Args:
2127
+ kv_cache_config: The KV cache config
2128
+ kv_cache_raw_tensors: The KV cache buffer of each layer, with
2129
+ correct size but uninitialized shape.
2130
+ Returns:
2131
+ Dict[str, torch.Tensor]: A map between layer names to their
2132
+ corresponding memory buffer for KV cache.
2133
+ """
2134
kv_caches: dict[str, torch.Tensor] = {}
2103
-
2104
- for i, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):
2105
- kv_cache_spec = kv_cache_group.kv_cache_spec
2106
- for layer_name in kv_cache_group.layer_names:
2107
- tensor_config = kv_cache_config.tensors[layer_name]
2108
- assert tensor_config.size % kv_cache_spec.page_size_bytes == 0
2109
- num_blocks = tensor_config.size // kv_cache_spec.page_size_bytes
2110
- # `num_blocks` is the number of blocks the model runner can use.
2111
- # `kv_cache_config.num_blocks` is the number of blocks that
2112
- # KVCacheManager may allocate.
2113
- # Since different GPUs may have different number of layers and
2114
- # different memory capacities, `num_blocks` can be different on
2115
- # different GPUs, and `kv_cache_config.num_blocks` is set to
2116
- # the min of all `num_blocks`. Verify it here.
2117
- assert num_blocks >= kv_cache_config.num_blocks
2135
+ for i, kv_cache_group_spec in enumerate(
2136
+ kv_cache_config.kv_cache_groups):
2137
+ kv_cache_spec = kv_cache_group_spec.kv_cache_spec
2138
+ for layer_name in kv_cache_group_spec.layer_names:
2139
+ raw_tensor = kv_cache_raw_tensors[layer_name]
2140
+ assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
2141
+ num_blocks = (raw_tensor.numel() //
2142
+ kv_cache_spec.page_size_bytes)
2143
if isinstance(kv_cache_spec, AttentionSpec):
2144
kv_cache_shape = self.attn_backends[i].get_kv_cache_shape(
2145
num_blocks, kv_cache_spec.block_size,
@@ -2140,13 +2165,29 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2165
kv_cache_stride_order.index(i)
2166
for i in range(len(kv_cache_stride_order))
2167
]
2143
- kv_caches[layer_name] = torch.zeros(
2144
- kv_cache_shape, dtype=dtype,
2145
- device=self.device).permute(*inv_order)
2168
+ kv_caches[layer_name] = kv_cache_raw_tensors[
2169
+ layer_name].view(dtype).view(kv_cache_shape).permute(
2170
+ *inv_order)
2171
else:
2147
- # TODO: add new branches when introducing more types of
2148
- # KV cache specs.
2149
- raise ValueError("Unknown KV cache spec type.")
2172
+ raise NotImplementedError
2173
+ return kv_caches
2174
+
2175
+ def initialize_kv_cache_tensors(
2176
+ self, kv_cache_config: KVCacheConfig) -> dict[str, torch.Tensor]:
2177
+ """
2178
+ Initialize the memory buffer for KV cache.
2179
+
2180
+ Args:
2181
+ kv_cache_config: The KV cache config
2182
+ Returns:
2183
+ Dict[str, torch.Tensor]: A map between layer names to their
2184
+ corresponding memory buffer for KV cache.
2185
+ """
2186
+ # Initialize the memory buffer for KV cache
2187
+ kv_cache_raw_tensors = self._allocate_kv_cache_tensors(kv_cache_config)
2188
+ # Change the memory buffer to the desired shape
2189
+ kv_caches = self._reshape_kv_cache_tensors(kv_cache_config,
2190
+ kv_cache_raw_tensors)
2191
2192
# Setup `kv_cache_config` and `kv_caches` for models
2193
# with cross-layer KV sharing
@@ -2157,17 +2198,30 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2198
kv_caches,
2199
)
2200
2201
+ bind_kv_cache(
2202
+ kv_caches,
2203
+ self.vllm_config.compilation_config.static_forward_context,
2204
+ self.kv_caches)
2205
+ return kv_caches
2206
+
2207
+ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
2208
+ """
2209
+ Initialize KV cache based on `kv_cache_config`.
2210
+ Args:
2211
+ kv_cache_config: Configuration for the KV cache, including the KV
2212
+ cache size of each layer
2213
+ """
2214
+ self.kv_cache_config = kv_cache_config
2215
+ self.may_reinitialize_input_batch(kv_cache_config)
2216
+ self.initialize_attn_backend(kv_cache_config)
2217
+ kv_caches = self.initialize_kv_cache_tensors(kv_cache_config)
2218
+
2219
if self.speculative_config and self.speculative_config.use_eagle():
2220
assert isinstance(self.drafter, EagleProposer)
2221
# validate all draft model layers belong to the same kv cache
2222
# group
2223
self.drafter.validate_same_kv_cache_group(kv_cache_config)
2224
2166
- bind_kv_cache(
2167
- kv_caches,
2168
- self.vllm_config.compilation_config.static_forward_context,
2169
- self.kv_caches)
2170
-
2225
if has_kv_transfer_group():
2226
get_kv_transfer_group().register_kv_caches(kv_caches)
2227
vllm/v1/worker/tpu_model_runner.py
+10
-4
@@ -1365,14 +1365,20 @@ class TPUModelRunner(LoRAModelRunnerMixin):
1365
assert self.block_table_cpu.dtype == self.input_batch.block_table[
1366
0].get_cpu_tensor().dtype
1367
1368
- kv_caches: dict[str, torch.Tensor] = {}
1368
+ kv_cache_sizes = {}
1369
+ for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
1370
+ assert len(kv_cache_tensor.shared_by) == 1, (
1371
+ "KV cache tensor shared by multiple layers is not supported in "
1372
+ "TPU.")
1373
+ kv_cache_sizes[kv_cache_tensor.shared_by[0]] = kv_cache_tensor.size
1374
1375
+ kv_caches: dict[str, torch.Tensor] = {}
1376
for kv_cache_group in kv_cache_config.kv_cache_groups:
1377
kv_cache_spec = kv_cache_group.kv_cache_spec
1378
for layer_name in kv_cache_group.layer_names:
1373
- tensor_config = kv_cache_config.tensors[layer_name]
1374
- assert tensor_config.size % kv_cache_spec.page_size_bytes == 0
1375
- num_blocks = tensor_config.size // kv_cache_spec.page_size_bytes
1379
+ tensor_size = kv_cache_sizes[layer_name]
1380
+ assert tensor_size % kv_cache_spec.page_size_bytes == 0
1381
+ num_blocks = tensor_size // kv_cache_spec.page_size_bytes # noqa
1382
if isinstance(kv_cache_spec, AttentionSpec):
1383
if self.use_spmd:
1384
num_kv_heads = kv_cache_spec.num_kv_heads