[Bugfix] Count per-group blocks in get_max_concurrency_for_kv_cache_config (#48317)
Signed-off-by: David Orman <ormandj@corenode.com> Co-authored-by: Luke Alonso <lalonso@gmail.com> Co-authored-by: Martin Vit <martin@voipmonitor.org> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
ormandj committed
Jul 20, 2026 at 19:29 UTC
2e2e626b40e8321731cfa42f0402eb42f229e00a
2 files changed
+139
-12
tests/v1/core/test_kv_cache_utils.py
+123
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import copy
4
import hashlib
5
import importlib
6
from collections.abc import Callable
@@ -1477,6 +1478,128 @@ def test_get_max_concurrency_for_kv_cache_config():
1478
assert num_tokens == max_concurrency_hybrid_model * max_model_len
1479
assert max_concurrency == max_concurrency_hybrid_model
1480
1481
+ # Unequal group sizes in the standard layout: each group's pages cost
1482
+ # whole pool blocks, so a request needs 1024 + 129 = 1153 blocks — the
1483
+ # same as the equal-hybrid case above, regardless of the second group
1484
+ # holding only 2 layers.
1485
+ kv_cache_config_unequal_groups = KVCacheConfig(
1486
+ num_blocks=1153 * 3,
1487
+ kv_cache_tensors=[],
1488
+ kv_cache_groups=[
1489
+ KVCacheGroupSpec([f"layer_{i}" for i in range(32)], full_attention_spec),
1490
+ KVCacheGroupSpec(["layer_32", "layer_33"], sliding_window_spec),
1491
+ ],
1492
+ )
1493
+ assert (
1494
+ get_max_concurrency_for_kv_cache_config(
1495
+ vllm_config, kv_cache_config_unequal_groups
1496
+ )
1497
+ == 3
1498
+ )
1499
+
1500
+ # UniformTypeKVCacheSpecs group (worker config shape): the aggregated
1501
+ # spec's memory/page ratio equals a single layer's page count, so the
1502
+ # group needs 1024 blocks and the request 1153 in total. The previous
1503
+ # formula normalized both groups' memory by the first group's page size,
1504
+ # reporting 3459/1057 = 3.27 here instead of 3 — and a different value
1505
+ # again for the scheduler-config shape below.
1506
+ uniform_full_spec = UniformTypeKVCacheSpecs(
1507
+ block_size=full_attention_spec.block_size,
1508
+ kv_cache_specs={f"layer_{i}": full_attention_spec for i in range(4)},
1509
+ )
1510
+ kv_cache_config_uniform_group = KVCacheConfig(
1511
+ num_blocks=1153 * 3,
1512
+ kv_cache_tensors=[],
1513
+ kv_cache_groups=[
1514
+ KVCacheGroupSpec([f"layer_{i}" for i in range(4)], uniform_full_spec),
1515
+ KVCacheGroupSpec(["layer_4", "layer_5"], sliding_window_spec),
1516
+ ],
1517
+ )
1518
+ assert (
1519
+ get_max_concurrency_for_kv_cache_config(
1520
+ vllm_config, kv_cache_config_uniform_group
1521
+ )
1522
+ == 3
1523
+ )
1524
+
1525
+ # Scheduler-config shape: generate_scheduler_kv_cache_config replaces the
1526
+ # uniform-type group's spec with a representative per-layer spec.
1527
+ # Capacity must not change between the two shapes (the engine computes
1528
+ # on the scheduler config, the worker loop on the worker config).
1529
+ kv_cache_config_scheduler_shape = generate_scheduler_kv_cache_config(
1530
+ [copy.deepcopy(kv_cache_config_uniform_group)]
1531
+ )
1532
+ assert get_max_concurrency_for_kv_cache_config(
1533
+ vllm_config, kv_cache_config_scheduler_shape
1534
+ ) == get_max_concurrency_for_kv_cache_config(
1535
+ vllm_config, kv_cache_config_uniform_group
1536
+ )
1537
+
1538
+
1539
+def test_get_max_concurrency_packed_kv_cache_config():
1540
+ from vllm.v1.core.kv_cache_utils import (
1541
+ _get_kv_cache_config_packed,
1542
+ _use_packed_kv_cache_config,
1543
+ )
1544
+
1545
+ model_config = ModelConfig(
1546
+ "Qwen/Qwen1.5-7B",
1547
+ runner="generate",
1548
+ dtype="float16",
1549
+ max_model_len=16384,
1550
+ )
1551
+ scheduler_config = SchedulerConfig(
1552
+ max_num_batched_tokens=1024,
1553
+ enable_chunked_prefill=True,
1554
+ max_model_len=model_config.max_model_len,
1555
+ is_encoder_decoder=model_config.is_encoder_decoder,
1556
+ async_scheduling=False,
1557
+ )
1558
+ vllm_config = VllmConfig(
1559
+ model_config=model_config,
1560
+ scheduler_config=scheduler_config,
1561
+ )
1562
+
1563
+ # All-UniformTypeKVCacheSpecs groups select the packed layout.
1564
+ mla_specs = {f"layer_{i}": new_mla_spec() for i in range(4)}
1565
+ swa_specs = {
1566
+ f"layer_{i}": SlidingWindowMLASpec(
1567
+ block_size=16,
1568
+ num_kv_heads=1,
1569
+ head_size=576,
1570
+ dtype=torch.float32,
1571
+ sliding_window=128,
1572
+ )
1573
+ for i in range(4, 6)
1574
+ }
1575
+ kv_cache_groups = [
1576
+ KVCacheGroupSpec(
1577
+ list(mla_specs),
1578
+ UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=mla_specs),
1579
+ ),
1580
+ KVCacheGroupSpec(
1581
+ list(swa_specs),
1582
+ UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=swa_specs),
1583
+ ),
1584
+ ]
1585
+ assert _use_packed_kv_cache_config(vllm_config, kv_cache_groups)
1586
+ num_blocks, kv_cache_tensors = _get_kv_cache_config_packed(
1587
+ vllm_config, kv_cache_groups, 2 * GiB_bytes
1588
+ )
1589
+ assert num_blocks > 0
1590
+ kv_cache_config_packed = KVCacheConfig(
1591
+ num_blocks=num_blocks,
1592
+ kv_cache_tensors=kv_cache_tensors,
1593
+ kv_cache_groups=kv_cache_groups,
1594
+ )
1595
+ # Per-request blocks: the MLA group needs cdiv(16384, 16) = 1024 pages;
1596
+ # the SWA group cdiv(min(128 - 1 + 1024, 16384), 16) + 1 = 73. The
1597
+ # previous formula normalized by the first group's page size and gave
1598
+ # 1061 blocks per request instead of 1097.
1599
+ assert get_max_concurrency_for_kv_cache_config(
1600
+ vllm_config, kv_cache_config_packed
1601
+ ) == num_blocks / (1024 + 73)
1602
+
1603
1604
def test_allocate_with_lookahead():
1605
"""Verify that lookahead tokens correctly affect block allocation"""
vllm/v1/core/kv_cache_utils.py
+16
-12
@@ -939,19 +939,23 @@ def get_max_concurrency_for_kv_cache_config(
939
) -> float:
940
"""
941
Get the maximum concurrency for the given KV cache configuration.
942
- """
943
- num_layer_per_group = max(
944
- len(group.layer_names) for group in kv_cache_config.kv_cache_groups
945
- )
946
- max_memory_usage_per_request = num_layer_per_group * max_memory_usage_bytes(
947
- vllm_config, (group.kv_cache_spec for group in kv_cache_config.kv_cache_groups)
948
- )
949
- memory_per_block = (
950
- kv_cache_config.kv_cache_groups[0].kv_cache_spec.page_size_bytes
951
- * num_layer_per_group
942
+
943
+ A request at max_model_len consumes whole blocks from each group's block
944
+ table — cdiv(per-request bytes, page bytes) of the group's spec — and all
945
+ groups draw those block ids from one shared pool, so the per-request
946
+ total is the sum over groups. The memory/page ratio is identical whether
947
+ a group carries an aggregated UniformTypeKVCacheSpecs (worker config) or
948
+ a representative per-layer spec (scheduler config), so both capacity
949
+ call sites agree.
950
+ """
951
+ num_blocks_per_request = sum(
952
+ cdiv(
953
+ group.kv_cache_spec.max_memory_usage_bytes(vllm_config),
954
+ group.kv_cache_spec.page_size_bytes,
955
+ )
956
+ for group in kv_cache_config.kv_cache_groups
957
)
953
- num_block_per_request = cdiv(max_memory_usage_per_request, memory_per_block)
954
- max_concurrency = kv_cache_config.num_blocks / num_block_per_request
958
+ max_concurrency = kv_cache_config.num_blocks / num_blocks_per_request
959
return max_concurrency
960
961