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 14, 05:30 UTC
[Hybrid Allocator] Support Pipeline Parallel (#23974)
Signed-off-by: Chen Zhang <zhangch99@outlook.com>
Chen Zhang committed
Sep 14, 2025 at 15:55 UTC
8e5cdcda4e5a55ff49d92d37139042dda44b6b3c
7 files changed
+457
-220
tests/distributed/test_pipeline_parallel.py
+1
-3
@@ -215,9 +215,7 @@ TEXT_GENERATION_MODELS = {
215
EMBEDDING_MODELS = { # type: ignore[var-annotated]
216
# [Text-only]
217
"intfloat/e5-mistral-7b-instruct": PPTestSettings.fast(runner="pooling"),
218
- # TODO: re-enable when https://github.com/vllm-project/vllm/issues/23883
219
- # is fixed
220
- #"BAAI/bge-multilingual-gemma2": PPTestSettings.fast(runner="pooling"),
218
+ "BAAI/bge-multilingual-gemma2": PPTestSettings.fast(runner="pooling"),
219
"Qwen/Qwen2.5-Math-RM-72B": PPTestSettings.fast(
220
load_format="dummy", runner="pooling"
221
),
tests/models/test_initialization.py
+5
-5
@@ -10,7 +10,7 @@ from vllm import LLM
10
from vllm.config import ModelImpl
11
from vllm.engine.llm_engine import LLMEngine as V0LLMEngine
12
from vllm.utils import GiB_bytes
13
-from vllm.v1.core.kv_cache_utils import get_kv_cache_config
13
+from vllm.v1.core.kv_cache_utils import get_kv_cache_configs
14
from vllm.v1.engine.core import EngineCore as V1EngineCore
15
16
from ..utils import create_new_process_for_each_test
@@ -68,11 +68,11 @@ def can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch,
68
69
def _initialize_kv_caches_v1(self, vllm_config):
70
kv_cache_specs = self.model_executor.get_kv_cache_specs()
71
- scheduler_kv_cache_config = get_kv_cache_config(
71
+ scheduler_kv_cache_config = get_kv_cache_configs(
72
vllm_config,
73
- kv_cache_specs[0],
74
- 10 * GiB_bytes,
75
- )
73
+ kv_cache_specs,
74
+ [10 * GiB_bytes],
75
+ )[0]
76
77
# gpu_blocks (> 0), cpu_blocks, scheduler_kv_cache_config
78
return 1, 0, scheduler_kv_cache_config
tests/v1/core/test_kv_cache_utils.py
+289
-81
@@ -18,13 +18,12 @@ from vllm.v1.core.kv_cache_manager import KVCacheManager
18
from vllm.v1.core.kv_cache_utils import (
19
BlockHash, FreeKVCacheBlockQueue, KVCacheBlock, PrefixCachingMetrics,
20
estimate_max_model_len, generate_block_hash_extra_keys,
21
- get_kv_cache_config, get_max_concurrency_for_kv_cache_config,
21
+ get_kv_cache_configs, get_max_concurrency_for_kv_cache_config,
22
get_request_block_hasher, hash_block_tokens, init_none_hash,
23
- is_kv_cache_type_uniform, make_block_hash_with_group_id,
24
- unify_kv_cache_configs)
23
+ is_kv_cache_type_uniform, make_block_hash_with_group_id)
24
from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
26
- KVCacheGroupSpec, KVCacheTensor,
27
- SlidingWindowSpec)
25
+ KVCacheGroupSpec, KVCacheSpec,
26
+ KVCacheTensor, SlidingWindowSpec)
27
from vllm.v1.metrics.stats import PrefixCacheStats
28
from vllm.v1.request import Request
29
@@ -531,102 +530,288 @@ def test_metrics():
530
assert not metrics.query_queue
531
532
534
-def test_unify_kv_cache_configs():
535
- same_kv_cache_config = [
533
+def test_get_kv_cache_configs_multiple_workers():
534
+ model_config = ModelConfig(max_model_len=16)
535
+ vllm_config = VllmConfig(model_config=model_config)
536
+
537
+ ref_kv_cache_spec = new_kv_cache_spec()
538
+ same_kv_cache_specs = [{
539
+ "layer1": new_kv_cache_spec(),
540
+ "layer2": new_kv_cache_spec(),
541
+ }, {
542
+ "layer1": new_kv_cache_spec(),
543
+ "layer2": new_kv_cache_spec(),
544
+ }]
545
+
546
+ # Basic case. All things are the same.
547
+ kv_cache_configs = get_kv_cache_configs(vllm_config, same_kv_cache_specs, [
548
+ ref_kv_cache_spec.page_size_bytes * 2 * 10,
549
+ ref_kv_cache_spec.page_size_bytes * 2 * 10
550
+ ])
551
+ assert kv_cache_configs == [
552
KVCacheConfig(
553
num_blocks=10,
554
kv_cache_tensors=[
539
- KVCacheTensor(size=100, shared_by=["layer1"]),
540
- KVCacheTensor(size=100, shared_by=["layer2"]),
555
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 10,
556
+ shared_by=["layer1"]),
557
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 10,
558
+ shared_by=["layer2"]),
559
],
560
kv_cache_groups=[
543
- KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
544
- KVCacheGroupSpec(["layer2"],
545
- new_kv_cache_spec(num_kv_heads=4)),
561
+ KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
562
],
563
),
564
KVCacheConfig(
549
- num_blocks=20,
565
+ num_blocks=10,
566
kv_cache_tensors=[
551
- KVCacheTensor(size=100, shared_by=["layer1"]),
552
- KVCacheTensor(size=100, shared_by=["layer2"]),
567
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 10,
568
+ shared_by=["layer1"]),
569
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 10,
570
+ shared_by=["layer2"]),
571
],
572
kv_cache_groups=[
555
- KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
556
- KVCacheGroupSpec(["layer2"],
557
- new_kv_cache_spec(num_kv_heads=4)),
573
+ KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
574
],
575
),
576
]
561
- unify_kv_cache_configs(same_kv_cache_config)
562
- assert same_kv_cache_config[0].num_blocks == 10
563
- assert same_kv_cache_config[1].num_blocks == 10
577
565
- need_sort_kv_cache_config = [
578
+ # Different available memory. This is the case for TP.
579
+ # Use the smallest memory available.
580
+ kv_cache_configs = get_kv_cache_configs(vllm_config, same_kv_cache_specs, [
581
+ ref_kv_cache_spec.page_size_bytes * 2 * 10,
582
+ ref_kv_cache_spec.page_size_bytes * 2 * 20
583
+ ])
584
+ assert kv_cache_configs == [
585
KVCacheConfig(
586
num_blocks=10,
587
kv_cache_tensors=[
569
- KVCacheTensor(size=100, shared_by=["layer1"]),
570
- KVCacheTensor(size=100, shared_by=["layer2"]),
588
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 10,
589
+ shared_by=["layer1"]),
590
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 10,
591
+ shared_by=["layer2"]),
592
],
593
kv_cache_groups=[
573
- KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
574
- KVCacheGroupSpec(["layer2"],
575
- new_kv_cache_spec(num_kv_heads=4)),
594
+ KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
595
+ ],
596
+ ),
597
+ KVCacheConfig(
598
+ num_blocks=10,
599
+ kv_cache_tensors=[
600
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 20,
601
+ shared_by=["layer1"]),
602
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 20,
603
+ shared_by=["layer2"]),
604
+ ],
605
+ kv_cache_groups=[
606
+ KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
607
],
608
),
609
+ ]
610
+
611
+ # Different KV cache specs. This is the case for PP.
612
+ different_layer_specs = [{
613
+ "layer1": new_kv_cache_spec(),
614
+ }, {
615
+ "layer2": new_kv_cache_spec(),
616
+ "layer3": new_kv_cache_spec(),
617
+ }]
618
+
619
+ # Different workers have different layers.
620
+ kv_cache_configs = get_kv_cache_configs(
621
+ vllm_config, different_layer_specs, [
622
+ ref_kv_cache_spec.page_size_bytes * 2 * 10,
623
+ ref_kv_cache_spec.page_size_bytes * 2 * 10
624
+ ])
625
+ assert kv_cache_configs == [
626
KVCacheConfig(
579
- num_blocks=20,
627
+ num_blocks=10,
628
kv_cache_tensors=[
581
- KVCacheTensor(size=100, shared_by=["layer1"]),
582
- KVCacheTensor(size=100, shared_by=["layer2"]),
629
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 20,
630
+ shared_by=["layer1"]),
631
],
632
kv_cache_groups=[
585
- KVCacheGroupSpec(["layer2"],
586
- new_kv_cache_spec(num_kv_heads=4)),
633
KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
634
],
635
),
636
+ KVCacheConfig(
637
+ num_blocks=10,
638
+ kv_cache_tensors=[
639
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 10,
640
+ shared_by=["layer2"]),
641
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 10,
642
+ shared_by=["layer3"]),
643
+ ],
644
+ kv_cache_groups=[
645
+ KVCacheGroupSpec(["layer2", "layer3"], new_kv_cache_spec()),
646
+ ],
647
+ ),
648
+ ]
649
+
650
+ # Some layers are the same, some are different. This is the case for TP+PP
651
+ tp_pp_kv_cache_specs = [{
652
+ "layer1": new_kv_cache_spec(),
653
+ "layer2": new_kv_cache_spec(),
654
+ }, {
655
+ "layer1": new_kv_cache_spec(),
656
+ "layer2": new_kv_cache_spec(),
657
+ }, {
658
+ "layer3": new_kv_cache_spec(),
659
+ }, {
660
+ "layer3": new_kv_cache_spec(),
661
+ }]
662
+
663
+ kv_cache_configs = get_kv_cache_configs(
664
+ vllm_config, tp_pp_kv_cache_specs, [
665
+ ref_kv_cache_spec.page_size_bytes * 2 * 10,
666
+ ref_kv_cache_spec.page_size_bytes * 2 * 10,
667
+ ref_kv_cache_spec.page_size_bytes * 2 * 10,
668
+ ref_kv_cache_spec.page_size_bytes * 2 * 10,
669
+ ])
670
+ assert kv_cache_configs == [
671
+ KVCacheConfig(
672
+ num_blocks=10,
673
+ kv_cache_tensors=[
674
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 10,
675
+ shared_by=["layer1"]),
676
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 10,
677
+ shared_by=["layer2"]),
678
+ ],
679
+ kv_cache_groups=[
680
+ KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
681
+ ],
682
+ ),
683
+ KVCacheConfig(
684
+ num_blocks=10,
685
+ kv_cache_tensors=[
686
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 10,
687
+ shared_by=["layer1"]),
688
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 10,
689
+ shared_by=["layer2"]),
690
+ ],
691
+ kv_cache_groups=[
692
+ KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
693
+ ],
694
+ ),
695
+ KVCacheConfig(
696
+ num_blocks=10,
697
+ kv_cache_tensors=[
698
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 20,
699
+ shared_by=["layer3"]),
700
+ ],
701
+ kv_cache_groups=[
702
+ KVCacheGroupSpec(["layer3"], ref_kv_cache_spec),
703
+ ],
704
+ ),
705
+ KVCacheConfig(
706
+ num_blocks=10,
707
+ kv_cache_tensors=[
708
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 20,
709
+ shared_by=["layer3"]),
710
+ ],
711
+ kv_cache_groups=[
712
+ KVCacheGroupSpec(["layer3"], ref_kv_cache_spec),
713
+ ],
714
+ ),
715
]
716
592
- unify_kv_cache_configs(need_sort_kv_cache_config)
593
- sorted_kv_cache_groups = [
594
- KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
595
- KVCacheGroupSpec(["layer2"], new_kv_cache_spec(num_kv_heads=4)),
717
+ # Different workers have different types of layers. This is the case for
718
+ # hybrid models + PP.
719
+ different_type_layer_specs = [{
720
+ "layer1": new_kv_cache_spec(),
721
+ "layer2": new_kv_cache_spec(),
722
+ }, {
723
+ "layer3": new_sliding_window_spec(),
724
+ "layer4": new_sliding_window_spec(),
725
+ }]
726
+ kv_cache_configs = get_kv_cache_configs(
727
+ vllm_config, different_type_layer_specs, [
728
+ ref_kv_cache_spec.page_size_bytes * 2 * 10,
729
+ ref_kv_cache_spec.page_size_bytes * 2 * 10,
730
+ ])
731
+ assert kv_cache_configs == [
732
+ KVCacheConfig(
733
+ num_blocks=10,
734
+ kv_cache_tensors=[
735
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 10,
736
+ shared_by=["layer1"]),
737
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 10,
738
+ shared_by=["layer2"]),
739
+ ],
740
+ kv_cache_groups=[
741
+ KVCacheGroupSpec(["layer1", "layer2"], ref_kv_cache_spec),
742
+ KVCacheGroupSpec([], new_sliding_window_spec()),
743
+ ],
744
+ ),
745
+ KVCacheConfig(
746
+ num_blocks=10,
747
+ kv_cache_tensors=[
748
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 10,
749
+ shared_by=["layer3"]),
750
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 10,
751
+ shared_by=["layer4"]),
752
+ ],
753
+ kv_cache_groups=[
754
+ KVCacheGroupSpec([], ref_kv_cache_spec),
755
+ KVCacheGroupSpec(["layer3", "layer4"],
756
+ new_sliding_window_spec()),
757
+ ],
758
+ ),
759
]
597
- assert (
598
- need_sort_kv_cache_config[0].kv_cache_groups == sorted_kv_cache_groups)
599
- assert (
600
- need_sort_kv_cache_config[1].kv_cache_groups == sorted_kv_cache_groups)
760
602
- diff_kv_cache_config = [
761
+ # When divided into multiple KVCacheGroups, need to ensure the number of
762
+ # layers per group is similar.
763
+ different_type_layer_specs = [{
764
+ "layer1": new_kv_cache_spec(),
765
+ "layer2": new_sliding_window_spec(),
766
+ "layer3": new_sliding_window_spec(),
767
+ }, {
768
+ "layer4": new_kv_cache_spec(),
769
+ "layer5": new_sliding_window_spec(),
770
+ "layer6": new_sliding_window_spec(),
771
+ }]
772
+ kv_cache_configs = get_kv_cache_configs(
773
+ vllm_config, different_type_layer_specs, [
774
+ ref_kv_cache_spec.page_size_bytes * 10,
775
+ ref_kv_cache_spec.page_size_bytes * 10,
776
+ ])
777
+ assert kv_cache_configs == [
778
KVCacheConfig(
779
num_blocks=10,
780
kv_cache_tensors=[
606
- KVCacheTensor(size=100, shared_by=["layer1"]),
607
- KVCacheTensor(size=100, shared_by=["layer2"]),
781
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 10,
782
+ shared_by=["layer1", "layer2", "layer3"]),
783
],
784
kv_cache_groups=[
610
- KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
611
- KVCacheGroupSpec(["layer2"],
612
- new_kv_cache_spec(num_kv_heads=4)),
785
+ KVCacheGroupSpec(["layer1"], ref_kv_cache_spec),
786
+ KVCacheGroupSpec(["layer2"], new_sliding_window_spec()),
787
+ KVCacheGroupSpec(["layer3"], new_sliding_window_spec()),
788
],
789
),
790
KVCacheConfig(
616
- num_blocks=20,
791
+ num_blocks=10,
792
kv_cache_tensors=[
618
- KVCacheTensor(size=100, shared_by=["layer1"]),
619
- KVCacheTensor(size=100, shared_by=["layer2"]),
793
+ KVCacheTensor(size=ref_kv_cache_spec.page_size_bytes * 10,
794
+ shared_by=["layer4", "layer5", "layer6"]),
795
],
796
kv_cache_groups=[
622
- KVCacheGroupSpec(["layer1"], new_kv_cache_spec()),
623
- KVCacheGroupSpec(["layer2"],
624
- new_kv_cache_spec(num_kv_heads=8)),
797
+ KVCacheGroupSpec(["layer4"], ref_kv_cache_spec),
798
+ KVCacheGroupSpec(["layer5"], new_sliding_window_spec()),
799
+ KVCacheGroupSpec(["layer6"], new_sliding_window_spec()),
800
],
801
),
802
]
803
+
804
+ # Have conflicting layers. Need to raise an error.
805
+ conflicting_layer_specs = [{
806
+ "layer1": new_kv_cache_spec(),
807
+ }, {
808
+ "layer1": new_sliding_window_spec(),
809
+ }]
810
with pytest.raises(AssertionError):
629
- unify_kv_cache_configs(diff_kv_cache_config)
811
+ get_kv_cache_configs(vllm_config, conflicting_layer_specs, [
812
+ ref_kv_cache_spec.page_size_bytes * 2 * 10,
813
+ ref_kv_cache_spec.page_size_bytes * 2 * 10,
814
+ ])
815
816
817
def test_merge_kv_cache_spec():
@@ -890,7 +1075,7 @@ def test_allocate_with_lookahead():
1075
assert len(blocks.get_block_ids()[0]) == 2
1076
1077
893
-def test_get_kv_cache_config():
1078
+def test_get_kv_cache_config_one_worker():
1079
# pass max_model_len to pass check_enough_kv_cache_memory
1080
model_config = ModelConfig(max_model_len=16)
1081
vllm_config = VllmConfig(model_config=model_config)
@@ -901,8 +1086,10 @@ def test_get_kv_cache_config():
1086
'layer_1': new_kv_cache_spec(),
1087
'layer_2': new_kv_cache_spec(),
1088
}
904
- kv_cache_config_full = get_kv_cache_config(
905
- vllm_config, kv_cache_specs_full, mem_per_block_per_layer * 2 * 32)
1089
+ kv_cache_config_full = get_kv_cache_configs(
1090
+ vllm_config, [kv_cache_specs_full],
1091
+ [mem_per_block_per_layer * 2 * 32])[0]
1092
+ print(kv_cache_config_full)
1093
assert kv_cache_config_full == KVCacheConfig(
1094
num_blocks=32,
1095
kv_cache_tensors=[
@@ -920,8 +1107,9 @@ def test_get_kv_cache_config():
1107
'layer_1': new_sliding_window_spec(),
1108
'layer_2': new_sliding_window_spec(),
1109
}
923
- kv_cache_config_sliding = get_kv_cache_config(
924
- vllm_config, kv_cache_specs_sliding, mem_per_block_per_layer * 2 * 32)
1110
+ kv_cache_config_sliding = get_kv_cache_configs(
1111
+ vllm_config, [kv_cache_specs_sliding],
1112
+ [mem_per_block_per_layer * 2 * 32])[0]
1113
assert kv_cache_config_sliding == KVCacheConfig(
1114
num_blocks=32,
1115
kv_cache_tensors=[
@@ -940,8 +1128,9 @@ def test_get_kv_cache_config():
1128
'layer_1': new_kv_cache_spec(),
1129
'layer_2': new_sliding_window_spec(),
1130
}
943
- kv_cache_config_hybrid = get_kv_cache_config(
944
- vllm_config, kv_cache_specs_hybrid, mem_per_block_per_layer * 2 * 32)
1131
+ kv_cache_config_hybrid = get_kv_cache_configs(
1132
+ vllm_config, [kv_cache_specs_hybrid],
1133
+ [mem_per_block_per_layer * 2 * 32])[0]
1134
assert kv_cache_config_hybrid == KVCacheConfig(
1135
num_blocks=32,
1136
kv_cache_tensors=[
@@ -962,8 +1151,9 @@ def test_get_kv_cache_config():
1151
'layer_1': new_kv_cache_spec(),
1152
'layer_2': new_sliding_window_spec(),
1153
}
965
- kv_cache_config_hybrid = get_kv_cache_config(
966
- vllm_config, kv_cache_specs_hybrid, mem_per_block_per_layer * 2 * 32)
1154
+ kv_cache_config_hybrid = get_kv_cache_configs(
1155
+ vllm_config, [kv_cache_specs_hybrid],
1156
+ [mem_per_block_per_layer * 2 * 32])[0]
1157
assert kv_cache_config_hybrid == KVCacheConfig(
1158
num_blocks=64,
1159
kv_cache_tensors=[
@@ -985,21 +1175,22 @@ def test_get_kv_cache_config():
1175
'layer_5': new_sliding_window_spec(),
1176
'layer_6': new_sliding_window_spec(),
1177
}
988
- kv_cache_config_hybrid = get_kv_cache_config(
989
- vllm_config, kv_cache_specs_hybrid, mem_per_block_per_layer * 2 * 32)
1178
+ kv_cache_config_hybrid = get_kv_cache_configs(
1179
+ vllm_config, [kv_cache_specs_hybrid],
1180
+ [mem_per_block_per_layer * 2 * 32])[0]
1181
assert kv_cache_config_hybrid == KVCacheConfig(
1182
num_blocks=32,
1183
kv_cache_tensors=[
1184
KVCacheTensor(size=mem_per_block_per_layer * 32,
994
- shared_by=["layer_1", "layer_3", "layer_5"]),
1185
+ shared_by=["layer_1", "layer_3", "layer_4"]),
1186
KVCacheTensor(size=mem_per_block_per_layer * 32,
996
- shared_by=["layer_2", "layer_4", "layer_6"]),
1187
+ shared_by=["layer_2", "layer_5", "layer_6"]),
1188
],
1189
kv_cache_groups=[
1190
KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec()),
1000
- KVCacheGroupSpec(["layer_3", "layer_4"],
1191
+ KVCacheGroupSpec(["layer_3", "layer_5"],
1192
new_sliding_window_spec()),
1002
- KVCacheGroupSpec(["layer_5", "layer_6"],
1193
+ KVCacheGroupSpec(["layer_4", "layer_6"],
1194
new_sliding_window_spec()),
1195
],
1196
)
@@ -1017,27 +1208,30 @@ def test_get_kv_cache_config():
1208
'layer_9': new_sliding_window_spec(),
1209
'layer_10': new_sliding_window_spec(),
1210
}
1020
- kv_cache_config_hybrid = get_kv_cache_config(
1021
- vllm_config, kv_cache_specs_hybrid, mem_per_block_per_layer * 3 * 32)
1211
+ kv_cache_config_hybrid = get_kv_cache_configs(
1212
+ vllm_config, [kv_cache_specs_hybrid],
1213
+ [mem_per_block_per_layer * 3 * 32])[0]
1214
assert kv_cache_config_hybrid == KVCacheConfig(
1215
num_blocks=32,
1216
kv_cache_tensors=[
1217
KVCacheTensor(
1218
size=mem_per_block_per_layer * 32,
1027
- shared_by=["layer_1", "layer_4", "layer_7", "layer_10"]),
1028
- KVCacheTensor(size=mem_per_block_per_layer * 32,
1029
- shared_by=["layer_2", "layer_5", "layer_8"]),
1219
+ shared_by=["layer_1", "layer_4", "layer_5", "layer_6"]),
1220
+ KVCacheTensor(
1221
+ size=mem_per_block_per_layer * 32,
1222
+ shared_by=["layer_2", "layer_7", "layer_8", "layer_9"]),
1223
KVCacheTensor(size=mem_per_block_per_layer * 32,
1031
- shared_by=["layer_3", "layer_6", "layer_9"]),
1224
+ shared_by=["layer_3", "layer_10"]),
1225
],
1226
kv_cache_groups=[
1227
KVCacheGroupSpec(["layer_1", "layer_2", "layer_3"],
1228
new_kv_cache_spec()),
1036
- KVCacheGroupSpec(["layer_4", "layer_5", "layer_6"],
1229
+ KVCacheGroupSpec(["layer_4", "layer_7", "layer_10"],
1230
+ new_sliding_window_spec()),
1231
+ KVCacheGroupSpec(["layer_5", "layer_8"],
1232
new_sliding_window_spec()),
1038
- KVCacheGroupSpec(["layer_7", "layer_8", "layer_9"],
1233
+ KVCacheGroupSpec(["layer_6", "layer_9"],
1234
new_sliding_window_spec()),
1040
- KVCacheGroupSpec(["layer_10"], new_sliding_window_spec()),
1235
],
1236
)
1237
@@ -1047,13 +1241,14 @@ def test_get_kv_cache_config():
1241
'layer_2': new_kv_cache_spec(),
1242
}
1243
with pytest.raises(NotImplementedError):
1050
- get_kv_cache_config(vllm_config, kv_cache_specs_hybrid,
1051
- mem_per_block_per_layer * 2 * 32)
1244
+ get_kv_cache_configs(vllm_config, [kv_cache_specs_hybrid],
1245
+ [mem_per_block_per_layer * 2 * 32])[0]
1246
1247
# Test num_gpu_blocks_override
1248
vllm_config.cache_config.num_gpu_blocks_override = 16
1055
- kv_cache_config_override_blocks = get_kv_cache_config(
1056
- vllm_config, kv_cache_specs_full, mem_per_block_per_layer * 2 * 32)
1249
+ kv_cache_config_override_blocks = get_kv_cache_configs(
1250
+ vllm_config, [kv_cache_specs_full],
1251
+ [mem_per_block_per_layer * 2 * 32])[0]
1252
assert kv_cache_config_override_blocks == KVCacheConfig(
1253
num_blocks=16,
1254
kv_cache_tensors=[
@@ -1065,3 +1260,16 @@ def test_get_kv_cache_config():
1260
kv_cache_groups=[
1261
KVCacheGroupSpec(["layer_1", "layer_2"], new_kv_cache_spec())
1262
])
1263
+
1264
+
1265
+def test_get_kv_cache_configs_attention_free():
1266
+ kv_cache_specs: dict[str, KVCacheSpec] = {}
1267
+ vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=16))
1268
+ kv_cache_configs = get_kv_cache_configs(vllm_config, [kv_cache_specs], [0])
1269
+ assert kv_cache_configs == [
1270
+ KVCacheConfig(
1271
+ num_blocks=1,
1272
+ kv_cache_tensors=[],
1273
+ kv_cache_groups=[],
1274
+ )
1275
+ ]
tests/v1/tpu/worker/test_tpu_model_runner.py
+5
-5
@@ -10,7 +10,7 @@ from vllm.pooling_params import PoolingParams
10
from vllm.sampling_params import SamplingParams
11
from vllm.utils import GiB_bytes
12
from vllm.v1.core.kv_cache_utils import (estimate_max_model_len,
13
- get_kv_cache_config)
13
+ get_kv_cache_configs)
14
from vllm.v1.core.sched.output import (CachedRequestData, NewRequestData,
15
SchedulerOutput)
16
from vllm.v1.worker.tpu_model_runner import (
@@ -477,8 +477,8 @@ def test_init_kv_cache_without_kv_sharing():
477
# 2 (non-MLA) * 8 (num_heads) * 128 (head_dim)
478
# * 2 (bfloat16, kv_cache dtype) * 128 (block_size) = 512KB
479
num_expected_blocks = 20480 # 20GB / 512KB / 2 (num layers)
480
- kv_cache_config = get_kv_cache_config(vllm_config, kv_cache_spec,
481
- available_memory)
480
+ kv_cache_config = get_kv_cache_configs(vllm_config, [kv_cache_spec],
481
+ [available_memory])[0]
482
assert kv_cache_config.num_blocks == num_expected_blocks
483
assert len(kv_cache_config.kv_cache_tensors) == 2
484
assert kv_cache_config.kv_cache_tensors[0].size == available_memory // 2
@@ -550,8 +550,8 @@ def test_init_kv_cache_with_kv_sharing_valid():
550
# with KV sharing, we can allocate (available_mem//page_size//1) blocks
551
# which is twice as many as without KV sharing
552
num_expected_blocks = 2 * 20480 # 20GB / 512KB
553
- kv_cache_config = get_kv_cache_config(vllm_config, kv_cache_spec,
554
- available_memory)
553
+ kv_cache_config = get_kv_cache_configs(vllm_config, [kv_cache_spec],
554
+ [available_memory])[0]
555
assert kv_cache_config.num_blocks == num_expected_blocks
556
assert len(kv_cache_config.kv_cache_tensors) == 1
557
# Each layer now has twice the available memory for KV cache
tests/v1/worker/test_gpu_model_runner.py
+7
-7
@@ -15,7 +15,7 @@ from vllm.platforms import current_platform
15
from vllm.sampling_params import SamplingParams
16
from vllm.utils import GiB_bytes, update_environment_variables
17
from vllm.v1.core.kv_cache_utils import (estimate_max_model_len,
18
- get_kv_cache_config)
18
+ get_kv_cache_configs)
19
from vllm.v1.core.sched.output import (CachedRequestData, NewRequestData,
20
SchedulerOutput)
21
from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
@@ -585,8 +585,8 @@ def test_init_kv_cache_without_kv_sharing():
585
available_memory = 20 * GiB_bytes
586
# page size for layer 0's kv_cache_spec is 32KB
587
num_expected_blocks = 327680 # 20GB / 32KB / 2 (num layers)
588
- kv_cache_config = get_kv_cache_config(vllm_config, kv_cache_spec,
589
- available_memory)
588
+ kv_cache_config = get_kv_cache_configs(vllm_config, [kv_cache_spec],
589
+ [available_memory])[0]
590
assert kv_cache_config.num_blocks == num_expected_blocks
591
assert len(kv_cache_config.kv_cache_tensors) == 2
592
assert kv_cache_config.kv_cache_tensors[0].size == available_memory // 2
@@ -657,8 +657,8 @@ def test_init_kv_cache_with_kv_sharing_valid():
657
# with KV sharing, we can allocate (available_mem//page_size//1) blocks
658
# which is twice as many as without KV sharing
659
num_expected_blocks = 655360 # 20GB / 32KB
660
- kv_cache_config = get_kv_cache_config(vllm_config, kv_cache_spec,
661
- available_memory)
660
+ kv_cache_config = get_kv_cache_configs(vllm_config, [kv_cache_spec],
661
+ [available_memory])[0]
662
assert kv_cache_config.num_blocks == num_expected_blocks
663
assert len(kv_cache_config.kv_cache_tensors) == 1
664
# Each layer now has twice the available memory for KV cache
@@ -788,8 +788,8 @@ def test_hybrid_attention_mamba_tensor_shapes(monkeypatch):
788
kv_cache_spec = runner.get_kv_cache_spec()
789
790
available_memory = 5 * GiB_bytes
791
- kv_cache_config = get_kv_cache_config(vllm_config, kv_cache_spec,
792
- available_memory)
791
+ kv_cache_config = get_kv_cache_configs(vllm_config, [kv_cache_spec],
792
+ [available_memory])[0]
793
runner.initialize_kv_cache(kv_cache_config)
794
795
# random partition of blocks
vllm/v1/core/kv_cache_utils.py
+145
-104
@@ -5,7 +5,7 @@
5
import os
6
from collections import defaultdict, deque
7
from collections.abc import Iterable, Sequence
8
-from dataclasses import astuple, dataclass
8
+from dataclasses import dataclass
9
from typing import Any, Callable, NewType, Optional, Union
10
11
from vllm import envs
@@ -811,59 +811,21 @@ def get_uniform_page_size(kv_cache_spec: dict[str, KVCacheSpec]) -> int:
811
return page_sizes.pop()
812
813
814
-def _get_kv_cache_config_uniform_type(vllm_config: VllmConfig,
815
- kv_cache_spec: dict[str, KVCacheSpec],
816
- available_memory: int) -> KVCacheConfig:
814
+def _get_kv_cache_groups_uniform_type(
815
+ kv_cache_specs: dict[str, KVCacheSpec]) -> list[KVCacheGroupSpec]:
816
"""
817
Generates the KV cache configuration for a model with one type of KV cache.
818
Divide the available memory equally among all layers.
819
820
Args:
822
- vllm_config: The global VllmConfig
823
- kv_cache_spec: The kv cache spec of each attention layer in the model
824
- available_memory: Memory available for KV cache in bytes.
821
+ kv_cache_specs: The kv cache spec of each attention layer in the model
822
823
Returns:
827
- The generated KVCacheConfig
824
+ The generated KVCacheGroupSpecs
825
"""
826
830
- page_size = get_uniform_page_size(kv_cache_spec)
831
- num_blocks = get_num_blocks(vllm_config, len(kv_cache_spec),
832
- available_memory, page_size)
833
-
834
- per_layer_size = page_size * num_blocks
835
- # All layers have the same KV cache spec, so we create one kv cache group
836
- # for all layers.
837
- grouped_layer_names = [list(kv_cache_spec.keys())]
838
-
839
- # Each layer uses a separate Tensor to store its KV cache.
840
- kv_cache_tensors = [
841
- KVCacheTensor(size=per_layer_size, shared_by=[layer_name])
842
- for layer_name in kv_cache_spec
843
- ]
844
-
845
- kv_cache_config = KVCacheConfig(
846
- num_blocks=num_blocks,
847
- kv_cache_tensors=kv_cache_tensors,
848
- kv_cache_groups=create_kv_cache_group_specs(kv_cache_spec,
849
- grouped_layer_names),
850
- )
851
-
852
- num_tokens = num_blocks * vllm_config.cache_config.block_size
853
- if vllm_config.parallel_config.decode_context_parallel_size > 1:
854
- num_tokens *= vllm_config.parallel_config.decode_context_parallel_size
855
- logger.info(
856
- "Multiplying the GPU KV cache size by the dcp_world_size %d.",
857
- vllm_config.parallel_config.decode_context_parallel_size)
858
-
859
- num_tokens_str = f"{num_tokens:,}"
860
- logger.info("GPU KV cache size: %s tokens", num_tokens_str)
861
- max_model_len_str = f"{vllm_config.model_config.max_model_len:,}"
862
- max_concurrency = get_max_concurrency_for_kv_cache_config(
863
- vllm_config, kv_cache_config)
864
- logger.info("Maximum concurrency for %s tokens per request: %.2fx",
865
- max_model_len_str, max_concurrency)
866
- return kv_cache_config
827
+ return create_kv_cache_group_specs(kv_cache_specs,
828
+ [list(kv_cache_specs.keys())])
829
830
831
def is_kv_cache_page_size_uniform(
@@ -888,11 +850,10 @@ def is_kv_cache_type_attention_free(
850
return not kv_cache_spec
851
852
891
-def _get_kv_cache_config_uniform_page_size(
892
- vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec],
893
- available_memory: int) -> KVCacheConfig:
853
+def _get_kv_cache_groups_uniform_page_size(
854
+ kv_cache_spec: dict[str, KVCacheSpec]) -> list[KVCacheGroupSpec]:
855
"""
895
- Generates the KV cache configuration for hybrid models with multiple
856
+ Generates the KV cache groups for hybrid models with multiple
857
attention types but still with a uniform page size (physical memory per
858
block per layer) for all layers.
859
@@ -949,11 +910,9 @@ def _get_kv_cache_config_uniform_page_size(
910
memory per block is the same for all groups.
911
912
Args:
952
- vllm_config: The global VllmConfig
913
kv_cache_spec: The KVCacheSpec of each attention layer in the model
954
- available_memory: Memory available for KV cache in bytes.
914
Returns:
956
- The generated KVCacheConfig
915
+ The generated KVCacheGroupSpecs
916
"""
917
# Group all layers by kv_cache_spec.
918
# E.g., 2 full attention layers and 3 sliding window attention layers,
@@ -966,7 +925,7 @@ def _get_kv_cache_config_uniform_page_size(
925
# group identical. Add padding to the last group of each type if necessary.
926
# E.g., (full.0, full.1), (sw.0, sw.1, sw.2)
927
# split to 3 groups with 2 layers each:
969
- # (full.0, full.1), (sw.0, sw.1), (sw.2, padding).
928
+ # (full.0, full.1), (sw.0, sw.2), (sw.1, padding).
929
# FIXME(Chen): At the moment of writing this code (2025-06-02), all
930
# open-source hybrid model follows a n:1 pattern between different attention
931
# types (e.g., Gemma3 5:1 between sw and full, LLaMA4 3:1 between local and
@@ -984,19 +943,60 @@ def _get_kv_cache_config_uniform_page_size(
943
num_padding_layers,
944
num_padding_layers / len(layers) * 100,
945
)
987
- for i in range(0, len(layers), group_size):
988
- grouped_layers.append(layers[i:i + group_size])
989
- kv_cache_groups = create_kv_cache_group_specs(kv_cache_spec,
990
- grouped_layers)
946
+ num_groups = cdiv(len(layers), group_size)
947
+ # In PP case, say if we have
948
+ # - stage 0: full.0, sw.0, sw.1
949
+ # - stage 1: full.1, sw.2, sw.3
950
+ # We should have 3 groups: (full.0, full.1), (sw.0, sw.2), (sw.1, sw.3)
951
+ # It can't be (full.0, full.1), (sw.0, sw.1), (sw.2, sw.3) because
952
+ # the 3 groups in stage 0 will be (full.0), (sw.0, sw.1), (empty group)
953
+ # and it will be padded to (full.0, padding), (sw.0, sw.1),
954
+ # (padding, padding) to ensure the number of layers in each group is
955
+ # the same and will cause memory waste.
956
+ # To avoid this, we assign layers[i::num_groups] to the i-th group
957
+ # instead of layers[i * group_size: (i + 1) * group_size]
958
+ for i in range(num_groups):
959
+ grouped_layers.append(layers[i::num_groups])
960
+ return create_kv_cache_group_specs(kv_cache_spec, grouped_layers)
961
+
962
+
963
+def get_kv_cache_config_from_groups(vllm_config: VllmConfig,
964
+ kv_cache_groups: list[KVCacheGroupSpec],
965
+ kv_cache_specs: dict[str, KVCacheSpec],
966
+ available_memory: int) -> KVCacheConfig:
967
+ """
968
+ Generate the KV cache configuration from the KV cache groups and spec
969
+ of each layer.
970
+
971
+ Args:
972
+ vllm_config: The global VllmConfig
973
+ kv_cache_groups: The KV cache groups
974
+ kv_cache_specs: The KV cache spec of each attention layer in the model
975
+ available_memory: Memory available for KV cache in bytes
976
+ Returns:
977
+ The generated KVCacheConfig
978
+ """
979
+ if len(kv_cache_groups) == 0:
980
+ # Attention free models do not have KV cache.
981
+ # Return num_blocks=1 as BlockPool always needs a null_block.
982
+ return KVCacheConfig(
983
+ num_blocks=1,
984
+ kv_cache_tensors=[],
985
+ kv_cache_groups=kv_cache_groups,
986
+ )
987
988
# Determine how model runners should initialize the KV cache tensors.
989
# We will have group_size memory pools, each is shared by one layer from
990
# each group. As layers of different groups have different block table,
991
# they will use different parts of the shared Tensor.
996
- # The memory layout in the example will be:
997
- # full.0, sw.0, sw.2: share a Tensor with size=available_memory//2
998
- # full.1, sw.1: share another Tensor with size=available_memory//2
999
- page_size = get_uniform_page_size(kv_cache_spec)
992
+ # The memory layout for 3 groups (full.0, full.1), (sw.0, sw.2),
993
+ # (sw.1, padding) will be: (group_size = 2)
994
+ # full.0, sw.0, sw.1: share a Tensor with size=available_memory//2
995
+ # full.1, sw.2: share another Tensor with size=available_memory//2
996
+ group_size = max(len(group.layer_names) for group in kv_cache_groups)
997
+
998
+ page_size = get_uniform_page_size(kv_cache_specs)
999
+ assert group_size > 0, "group_size must be greater than 0"
1000
num_blocks = get_num_blocks(vllm_config, group_size, available_memory,
1001
page_size)
1002
per_memory_pool_size = page_size * num_blocks
@@ -1004,8 +1004,8 @@ def _get_kv_cache_config_uniform_page_size(
1004
for i in range(group_size):
1005
shared_by = []
1006
for j in range(len(kv_cache_groups)):
1007
- if i < len(grouped_layers[j]):
1008
- shared_by.append(grouped_layers[j][i])
1007
+ if i < len(kv_cache_groups[j].layer_names):
1008
+ shared_by.append(kv_cache_groups[j].layer_names[i])
1009
kv_cache_tensors.append(
1010
KVCacheTensor(size=per_memory_pool_size, shared_by=shared_by))
1011
@@ -1019,7 +1019,12 @@ def _get_kv_cache_config_uniform_page_size(
1019
[group.kv_cache_spec.block_size for group in kv_cache_groups])
1020
1021
# Print the KV cache size and maximum concurrency.
1022
- num_tokens = num_blocks // len(grouped_layers) * min_block_size
1022
+ num_tokens = num_blocks // len(kv_cache_groups) * min_block_size
1023
+ if vllm_config.parallel_config.decode_context_parallel_size > 1:
1024
+ num_tokens *= vllm_config.parallel_config.decode_context_parallel_size
1025
+ logger.info(
1026
+ "Multiplying the GPU KV cache size by the dcp_world_size %d.",
1027
+ vllm_config.parallel_config.decode_context_parallel_size)
1028
num_tokens_str = f"{num_tokens:,}"
1029
logger.info("GPU KV cache size: %s tokens", num_tokens_str)
1030
max_model_len_str = f"{vllm_config.model_config.max_model_len:,}"
@@ -1030,10 +1035,6 @@ def _get_kv_cache_config_uniform_page_size(
1035
return kv_cache_config
1036
1037
1033
-def _get_kv_cache_config_attention_free() -> KVCacheConfig:
1034
- return KVCacheConfig(num_blocks=1, kv_cache_tensors=[], kv_cache_groups=[])
1035
-
1036
-
1038
def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
1039
"""
1040
This function tries to convert the KV cache specs to one type if the model
@@ -1087,72 +1088,112 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
1088
"convert the KV cache specs to one unified type.")
1089
1090
1090
-def get_kv_cache_config(
1091
- vllm_config: VllmConfig,
1092
- kv_cache_spec: dict[str, KVCacheSpec],
1093
- available_memory: int,
1094
-) -> KVCacheConfig:
1091
+def get_kv_cache_groups(
1092
+ vllm_config: VllmConfig,
1093
+ kv_cache_spec: dict[str, KVCacheSpec]) -> list[KVCacheGroupSpec]:
1094
"""
1096
- Generates the KV cache configuration for a model.
1095
+ Split the layers in the model into groups with the same KV cache spec.
1096
1097
Args:
1098
vllm_config: The global VllmConfig
1099
kv_cache_spec: The kv cache spec of each attention layer in the model
1101
- available_memory: Memory available for KV cache in bytes.
1100
1101
Returns:
1104
- The generated KVCacheConfigs
1102
+ The generated KVCacheGroups
1103
"""
1106
- check_enough_kv_cache_memory(vllm_config, kv_cache_spec, available_memory)
1104
if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
1105
unify_hybrid_kv_cache_specs(kv_cache_spec)
1106
1107
if is_kv_cache_type_attention_free(kv_cache_spec):
1111
- # This returns a kv_cache config with 0 kv_cache groups and 1 block
1112
- # to allow for the KVCache manager to handle attention free models.
1113
- return _get_kv_cache_config_attention_free()
1108
+ # This returns an empty list to allow for the KVCacheManager to handle
1109
+ # attention free models.
1110
+ return []
1111
elif is_kv_cache_type_uniform(kv_cache_spec):
1112
# KV cache of all layers are the same, which is true for
1113
# most models. Allocate the same amount of memory for
1114
# each layer.
1118
- return _get_kv_cache_config_uniform_type(vllm_config, kv_cache_spec,
1119
- available_memory)
1115
+ return _get_kv_cache_groups_uniform_type(kv_cache_spec)
1116
elif is_kv_cache_page_size_uniform(kv_cache_spec):
1117
# Model contains multiple attention types, but KV cache of all layers
1118
# have the same physical memory per block per layer. Split the layers
1119
# into groups with the same number of layers, and thus same total page
1120
# size.
1125
- return _get_kv_cache_config_uniform_page_size(vllm_config,
1126
- kv_cache_spec,
1127
- available_memory)
1121
+ return _get_kv_cache_groups_uniform_page_size(kv_cache_spec)
1122
1123
raise NotImplementedError
1124
1125
1132
-def unify_kv_cache_configs(kv_cache_configs: list[KVCacheConfig]):
1126
+def get_kv_cache_configs(vllm_config: VllmConfig,
1127
+ kv_cache_specs: list[dict[str, KVCacheSpec]],
1128
+ available_memory: list[int]) -> list[KVCacheConfig]:
1129
"""
1134
- Make the KV cache configurations for each worker consistent, so that all
1135
- workers can be controlled by the same KVCacheManager.
1136
- This function verifies that the layer group of each worker are the same,
1137
- and changes the num_blocks of each worker to the smallest among all workers.
1130
+ Generates the KV cache configurations for a model.
1131
+ Since we use a shared centralized controller for all workers, we need the
1132
+ `kv_cache_config` to be consistent across all workers to make sure
1133
+ the KV cache allocation can be applied to all workers. However, different
1134
+ workers may have different memory available, and different type of layers
1135
+ (when pipeline parallel is enabled). To handle the difference between
1136
+ workers, the current implementation is:
1137
+ 1. Merge the KV cache specs of all workers to get the KVCacheSpecs for
1138
+ the whole model.
1139
+ 2. Generate the KV cache groups based on the layer ratio of the whole model.
1140
+ 3. Generate the KV cache configs for each worker based on the KV cache
1141
+ grouping strategy. (This is reasonable because the layer ratio of
1142
+ different PP stages are similar.)
1143
+ 4. Change the num_blocks of each worker to the smallest among all workers.
1144
1145
Args:
1140
- kv_cache_configs: The KV cache configurations for each worker. Will be
1141
- in-place modified to make them consistent.
1146
+ vllm_config: The global VllmConfig
1147
+ kv_cache_specs: List of dict[layer_name, KVCacheSpec] for each worker.
1148
+ available_memory: Memory available for KV cache in bytes for each
1149
+ worker.
1150
+
1151
+ Returns:
1152
+ The generated KVCacheConfigs for each worker.
1153
"""
1154
1144
- # Sort the kv cache groups by their KV cache spec.
1145
- # This can avoid the inconsistency caused by the order of groups.
1146
- for kv_cache_config in kv_cache_configs:
1147
- kv_cache_config.kv_cache_groups.sort(key=lambda x: (type(
1148
- x.kv_cache_spec).__name__, astuple(x.kv_cache_spec)))
1149
-
1150
- # Verify that the groups of each rank are the same.
1151
- for kv_cache_config in kv_cache_configs[1:]:
1152
- for group_rank_0, group_rank_i in zip(
1153
- kv_cache_configs[0].kv_cache_groups,
1154
- kv_cache_config.kv_cache_groups):
1155
- assert group_rank_0.kv_cache_spec == group_rank_i.kv_cache_spec
1155
+ # Check if the available memory is enough for each worker.
1156
+ for kv_cache_spec_one_worker, available_memory_one_worker in zip(
1157
+ kv_cache_specs, available_memory):
1158
+ check_enough_kv_cache_memory(vllm_config, kv_cache_spec_one_worker,
1159
+ available_memory_one_worker)
1160
+
1161
+ # Merge the KV cache specs of all workers. Different PP stages may have
1162
+ # different layer names, and different TP ranks of the same PP stage should
1163
+ # have the same KV cache spec.
1164
+ merged_kv_cache_specs: dict[str, KVCacheSpec] = {}
1165
+ for kv_cache_spec_one_worker in kv_cache_specs:
1166
+ for layer_name, layer_spec in kv_cache_spec_one_worker.items():
1167
+ if layer_name not in merged_kv_cache_specs:
1168
+ merged_kv_cache_specs[layer_name] = layer_spec
1169
+ else:
1170
+ assert merged_kv_cache_specs[layer_name] == layer_spec, (
1171
+ "The KV cache specs for the same layer are different "
1172
+ "across workers. This is not supported yet.")
1173
+ global_kv_cache_groups = get_kv_cache_groups(vllm_config,
1174
+ merged_kv_cache_specs)
1175
+
1176
+ kv_cache_configs: list[KVCacheConfig] = []
1177
+ for kv_cache_spec_one_worker, available_memory_one_worker in zip(
1178
+ kv_cache_specs, available_memory):
1179
+ kv_cache_groups_one_worker: list[KVCacheGroupSpec] = []
1180
+ for group in global_kv_cache_groups:
1181
+ group_layer_names_one_worker = [
1182
+ layer_name for layer_name in group.layer_names
1183
+ if layer_name in kv_cache_spec_one_worker
1184
+ ]
1185
+ kv_cache_groups_one_worker.append(
1186
+ KVCacheGroupSpec(group_layer_names_one_worker,
1187
+ group.kv_cache_spec))
1188
+ assert sum(
1189
+ len(group.layer_names) for group in
1190
+ kv_cache_groups_one_worker) == len(kv_cache_spec_one_worker), (
1191
+ "Some layers are not assigned to any group.")
1192
+ kv_cache_configs.append(
1193
+ get_kv_cache_config_from_groups(vllm_config,
1194
+ kv_cache_groups_one_worker,
1195
+ kv_cache_spec_one_worker,
1196
+ available_memory_one_worker))
1197
1198
# Change the num_blocks of each rank to the smallest among all ranks. We
1199
# do not need to shrink the tensor size because it is valid to only use the
vllm/v1/engine/core.py
+5
-15
@@ -29,10 +29,9 @@ from vllm.transformers_utils.config import (
29
maybe_register_config_serialize_by_value)
30
from vllm.utils import (decorate_logs, get_hash_fn_by_name, make_zmq_socket,
31
resolve_obj_by_qualname, set_process_title)
32
-from vllm.v1.core.kv_cache_utils import (BlockHash, get_kv_cache_config,
32
+from vllm.v1.core.kv_cache_utils import (BlockHash, get_kv_cache_configs,
33
get_request_block_hasher,
34
- init_none_hash,
35
- unify_kv_cache_configs)
34
+ init_none_hash)
35
from vllm.v1.core.sched.interface import SchedulerInterface
36
from vllm.v1.core.sched.output import SchedulerOutput
37
from vllm.v1.core.sched.scheduler import Scheduler as V1Scheduler
@@ -191,18 +190,9 @@ class EngineCore:
190
available_gpu_memory = [0] * len(kv_cache_specs)
191
192
assert len(kv_cache_specs) == len(available_gpu_memory)
194
- # Get the kv cache tensor size
195
- kv_cache_configs = [
196
- get_kv_cache_config(vllm_config, kv_cache_spec_one_worker,
197
- available_gpu_memory_one_worker)
198
- for kv_cache_spec_one_worker, available_gpu_memory_one_worker in
199
- zip(kv_cache_specs, available_gpu_memory)
200
- ]
201
-
202
- # Since we use a shared centralized controller, we need the
203
- # `kv_cache_config` to be consistent across all workers to make sure
204
- # all the memory operators can be applied to all workers.
205
- unify_kv_cache_configs(kv_cache_configs)
193
+
194
+ kv_cache_configs = get_kv_cache_configs(vllm_config, kv_cache_specs,
195
+ available_gpu_memory)
196
197
# All workers have the same kv_cache_config except layer names, so use
198
# an arbitrary one to initialize the scheduler.