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 8, 14:13 UTC
[Core][BugFix] Fix PP KV cache sharding memory validation (#33698)
Signed-off-by: junuxyz <216036880+junuxyz@users.noreply.github.com>
junuxyz committed
Feb 11, 2026 at 00:46 UTC
c5a66d16970fbbc4633761d30f12ec1fc98a9523
2 files changed
+168
-39
tests/v1/core/test_kv_cache_utils.py
+93
@@ -1046,6 +1046,99 @@ def test_get_kv_cache_configs_multiple_workers():
1046
)
1047
1048
1049
+@pytest.mark.parametrize(
1050
+ "asymmetric_memory",
1051
+ [False, True],
1052
+ ids=["symmetric", "asymmetric"],
1053
+)
1054
+def test_get_kv_cache_configs_pp_sharding(asymmetric_memory):
1055
+ model_config = ModelConfig(max_model_len=512)
1056
+ vllm_config = VllmConfig(model_config=model_config)
1057
+
1058
+ ref_kv_cache_spec = new_kv_cache_spec()
1059
+ pp_kv_cache_specs = [
1060
+ {"layer1": ref_kv_cache_spec},
1061
+ {"layer2": ref_kv_cache_spec},
1062
+ ]
1063
+
1064
+ expected_num_blocks = model_config.max_model_len // ref_kv_cache_spec.block_size + 1
1065
+ avail_memory = ref_kv_cache_spec.page_size_bytes * expected_num_blocks
1066
+
1067
+ # With per-worker validation, each worker only needs memory for its own
1068
+ # layers. Worker 2 having more memory shouldn't affect worker 1's config.
1069
+ available_memory = (
1070
+ [avail_memory, avail_memory * 2] if asymmetric_memory else [avail_memory] * 2
1071
+ )
1072
+
1073
+ kv_cache_configs = get_kv_cache_configs(
1074
+ vllm_config,
1075
+ pp_kv_cache_specs,
1076
+ available_memory,
1077
+ )
1078
+
1079
+ assert kv_cache_configs == [
1080
+ KVCacheConfig(
1081
+ num_blocks=expected_num_blocks,
1082
+ kv_cache_tensors=[
1083
+ KVCacheTensor(
1084
+ size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks,
1085
+ shared_by=["layer1"],
1086
+ ),
1087
+ ],
1088
+ kv_cache_groups=[KVCacheGroupSpec(["layer1"], ref_kv_cache_spec)],
1089
+ ),
1090
+ KVCacheConfig(
1091
+ num_blocks=expected_num_blocks,
1092
+ kv_cache_tensors=[
1093
+ KVCacheTensor(
1094
+ size=ref_kv_cache_spec.page_size_bytes * expected_num_blocks,
1095
+ shared_by=["layer2"],
1096
+ ),
1097
+ ],
1098
+ kv_cache_groups=[KVCacheGroupSpec(["layer2"], ref_kv_cache_spec)],
1099
+ ),
1100
+ ]
1101
+
1102
+
1103
+def test_project_kv_cache_groups_to_worker():
1104
+ spec_a = new_kv_cache_spec()
1105
+ spec_b = new_kv_cache_spec(num_kv_heads=4)
1106
+
1107
+ global_groups = [
1108
+ KVCacheGroupSpec(["layer1", "layer2", "layer3"], spec_a),
1109
+ ]
1110
+ worker_spec = {"layer1": spec_a, "layer2": spec_a}
1111
+ projected = kv_cache_utils._project_kv_cache_groups_to_worker(
1112
+ global_groups, worker_spec
1113
+ )
1114
+ assert len(projected) == 1
1115
+ assert projected[0].layer_names == ["layer1", "layer2"]
1116
+ assert projected[0].kv_cache_spec is spec_a
1117
+
1118
+ projected = kv_cache_utils._project_kv_cache_groups_to_worker(
1119
+ global_groups, {"layer4": spec_a}
1120
+ )
1121
+ assert len(projected) == 1
1122
+ assert projected[0].layer_names == []
1123
+ assert projected[0].kv_cache_spec is spec_a
1124
+
1125
+ uniform_spec = UniformTypeKVCacheSpecs(
1126
+ block_size=16,
1127
+ kv_cache_specs={"layer1": spec_a, "layer2": spec_b, "layer3": spec_a},
1128
+ )
1129
+ global_groups_uniform = [
1130
+ KVCacheGroupSpec(["layer1", "layer2", "layer3"], uniform_spec),
1131
+ ]
1132
+ projected = kv_cache_utils._project_kv_cache_groups_to_worker(
1133
+ global_groups_uniform, {"layer1": spec_a, "layer3": spec_a}
1134
+ )
1135
+ assert len(projected) == 1
1136
+ assert projected[0].layer_names == ["layer1", "layer3"]
1137
+ proj_spec = projected[0].kv_cache_spec
1138
+ assert isinstance(proj_spec, UniformTypeKVCacheSpecs)
1139
+ assert set(proj_spec.kv_cache_specs.keys()) == {"layer1", "layer3"}
1140
+
1141
+
1142
def test_merge_kv_cache_spec():
1143
same_layer_specs = [
1144
new_kv_cache_spec(num_kv_heads=32),
vllm/v1/core/kv_cache_utils.py
+75
-39
@@ -7,6 +7,7 @@ import os
7
from collections import defaultdict
8
from collections.abc import Callable, Iterable, Iterator, Sequence
9
from dataclasses import dataclass, replace
10
+from functools import partial
11
from typing import Any, NewType, TypeAlias, overload
12
13
from vllm import envs
@@ -1390,7 +1391,7 @@ def _estimate_max_model_len_from_groups(
1391
1392
def _auto_fit_max_model_len(
1393
vllm_config: VllmConfig,
1393
- kv_cache_groups: list[KVCacheGroupSpec],
1394
+ projected_groups_per_worker: list[list[KVCacheGroupSpec]],
1395
available_memory: list[int],
1396
) -> None:
1397
"""
@@ -1401,14 +1402,13 @@ def _auto_fit_max_model_len(
1402
1403
Args:
1404
vllm_config: The global VllmConfig (will be modified in-place)
1404
- kv_cache_groups: The global KV cache groups (from get_kv_cache_groups).
1405
- This correctly accounts for padding in hybrid models.
1405
+ projected_groups_per_worker: KV cache groups projected to each worker.
1406
available_memory: Memory available for KV cache in bytes for each
1407
worker.
1408
"""
1409
original_max = vllm_config.model_config.max_model_len
1410
1411
- if not kv_cache_groups:
1411
+ if all(not groups for groups in projected_groups_per_worker):
1412
# All workers have empty specs (attention-free model)
1413
logger.info_once(
1414
"Auto-fit max_model_len: attention-free model, "
@@ -1418,11 +1418,16 @@ def _auto_fit_max_model_len(
1418
)
1419
return
1420
1421
- # Use minimum available memory across all workers
1422
- min_available_memory = min(available_memory)
1423
- auto_fit_max = _estimate_max_model_len_from_groups(
1424
- vllm_config, kv_cache_groups, min_available_memory
1425
- )
1421
+ # Find the max_model_len that fits across all workers.
1422
+ auto_fit_max = original_max
1423
+ limiting_worker_mem = available_memory[0]
1424
+ for groups, avail_mem in zip(projected_groups_per_worker, available_memory):
1425
+ if not groups:
1426
+ continue
1427
+ worker_max = _estimate_max_model_len_from_groups(vllm_config, groups, avail_mem)
1428
+ if worker_max < auto_fit_max:
1429
+ auto_fit_max = worker_max
1430
+ limiting_worker_mem = avail_mem
1431
1432
if auto_fit_max <= 0:
1433
raise ValueError(
@@ -1446,11 +1451,47 @@ def _auto_fit_max_model_len(
1451
"available GPU memory (%s GiB available for KV cache)",
1452
original_max,
1453
auto_fit_max,
1449
- format_gib(min_available_memory),
1454
+ format_gib(limiting_worker_mem),
1455
scope="local",
1456
)
1457
1458
1459
+def _project_kv_cache_groups_to_worker(
1460
+ global_kv_cache_groups: list[KVCacheGroupSpec],
1461
+ worker_spec: dict[str, KVCacheSpec],
1462
+) -> list[KVCacheGroupSpec]:
1463
+ """
1464
+ Projects global KV cache groups onto a single worker's assigned layers.
1465
+
1466
+ In pipeline parallelism, each worker only owns a subset of layers. This
1467
+ function filters the global groups to include only layers present on the
1468
+ given worker, adjusting UniformTypeKVCacheSpecs accordingly.
1469
+
1470
+ Args:
1471
+ global_kv_cache_groups: The global KV cache groups for the whole model.
1472
+ worker_spec: The KV cache spec of each layer on this worker.
1473
+
1474
+ Returns:
1475
+ The projected KV cache groups containing only this worker's layers.
1476
+ """
1477
+ projected_groups: list[KVCacheGroupSpec] = []
1478
+ for group in global_kv_cache_groups:
1479
+ worker_layer_names = [
1480
+ layer_name for layer_name in group.layer_names if layer_name in worker_spec
1481
+ ]
1482
+ group_spec = group.kv_cache_spec
1483
+ if worker_layer_names and isinstance(group_spec, UniformTypeKVCacheSpecs):
1484
+ group_spec = UniformTypeKVCacheSpecs(
1485
+ block_size=group_spec.block_size,
1486
+ kv_cache_specs={
1487
+ layer_name: group_spec.kv_cache_specs[layer_name]
1488
+ for layer_name in worker_layer_names
1489
+ },
1490
+ )
1491
+ projected_groups.append(KVCacheGroupSpec(worker_layer_names, group_spec))
1492
+ return projected_groups
1493
+
1494
+
1495
def get_kv_cache_configs(
1496
vllm_config: VllmConfig,
1497
kv_cache_specs: list[dict[str, KVCacheSpec]],
@@ -1468,7 +1509,8 @@ def get_kv_cache_configs(
1509
the whole model.
1510
2. Generate the KV cache groups based on the layer ratio of the whole model.
1511
This also handles spec unification for hybrid models.
1471
- 3. Handle auto-fit max_model_len and memory checks using the unified specs.
1512
+ 3. Handle auto-fit max_model_len and memory checks using per-worker
1513
+ projected groups to account for PP sharding.
1514
4. Generate the KV cache configs for each worker based on the KV cache
1515
grouping strategy. (This is reasonable because the layer ratio of
1516
different PP stages are similar.)
@@ -1506,44 +1548,38 @@ def get_kv_cache_configs(
1548
1549
# If original_max_model_len was -1, automatically
1550
# determine the maximum model length that fits in available GPU memory.
1509
- # We use the global groups here to correctly account for padding.
1551
+ # We use per-worker projected groups to account for PP sharding.
1552
+ projected_groups_per_worker = [
1553
+ _project_kv_cache_groups_to_worker(global_kv_cache_groups, worker_spec)
1554
+ for worker_spec in kv_cache_specs
1555
+ ]
1556
+
1557
if vllm_config.model_config.original_max_model_len == -1:
1511
- _auto_fit_max_model_len(vllm_config, global_kv_cache_groups, available_memory)
1558
+ _auto_fit_max_model_len(
1559
+ vllm_config, projected_groups_per_worker, available_memory
1560
+ )
1561
1513
- # Check if the available memory is enough (using min across all workers).
1514
- # We use the global groups to correctly account for padding.
1515
- if global_kv_cache_groups:
1562
+ # Check if the available memory is enough per worker.
1563
+ for groups, avail_mem in zip(projected_groups_per_worker, available_memory):
1564
+ if not groups:
1565
+ continue
1566
_check_enough_kv_cache_memory(
1517
- min(available_memory),
1518
- lambda: _max_memory_usage_bytes_from_groups(
1519
- vllm_config, global_kv_cache_groups
1520
- ),
1567
+ avail_mem,
1568
+ partial(_max_memory_usage_bytes_from_groups, vllm_config, groups),
1569
vllm_config.model_config.max_model_len,
1522
- lambda am: _estimate_max_model_len_from_groups(
1523
- vllm_config, global_kv_cache_groups, am
1524
- ),
1570
+ partial(_estimate_max_model_len_from_groups, vllm_config, groups),
1571
)
1572
1573
kv_cache_configs: list[KVCacheConfig] = []
1528
- for kv_cache_spec_one_worker, available_memory_one_worker in zip(
1529
- kv_cache_specs, available_memory
1574
+ for projected_groups, kv_cache_spec_one_worker, available_memory_one_worker in zip(
1575
+ projected_groups_per_worker, kv_cache_specs, available_memory
1576
):
1531
- kv_cache_groups_one_worker: list[KVCacheGroupSpec] = []
1532
- for group in global_kv_cache_groups:
1533
- group_layer_names_one_worker = [
1534
- layer_name
1535
- for layer_name in group.layer_names
1536
- if layer_name in kv_cache_spec_one_worker
1537
- ]
1538
- kv_cache_groups_one_worker.append(
1539
- KVCacheGroupSpec(group_layer_names_one_worker, group.kv_cache_spec)
1540
- )
1541
- assert sum(
1542
- len(group.layer_names) for group in kv_cache_groups_one_worker
1543
- ) == len(kv_cache_spec_one_worker), "Some layers are not assigned to any group."
1577
+ assert sum(len(group.layer_names) for group in projected_groups) == len(
1578
+ kv_cache_spec_one_worker
1579
+ ), "Some layers are not assigned to any group."
1580
kv_cache_configs.append(
1581
get_kv_cache_config_from_groups(
1546
- vllm_config, kv_cache_groups_one_worker, available_memory_one_worker
1582
+ vllm_config, projected_groups, available_memory_one_worker
1583
)
1584
)
1585