From f0de1a604cad003379e5bb4dfc3cc5d2a1f25fa8 Mon Sep 17 00:00:00 2001 From: drakosha Date: Mon, 3 Aug 2026 20:34:29 +0300 Subject: [PATCH] [Bugfix] Shard UniformTypeKVCacheSpecs block table width under DCP (#50823) Signed-off-by: Mikhail Kostryukov Signed-off-by: Matthew Bonanni Co-authored-by: Claude Co-authored-by: Matthew Bonanni --- tests/v1/core/test_kv_cache_utils.py | 28 ++++++++++++++++++++++++++++ vllm/v1/kv_cache_interface.py | 14 ++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 4ee2ea9ef4..857cbdabe4 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -1257,6 +1257,34 @@ def test_project_kv_cache_groups_to_worker(): assert set(proj_spec.kv_cache_specs.keys()) == {"layer1", "layer3"} +@pytest.mark.parametrize( + "layer_type,dcp_size,expected_width", + [ + ("mla", 1, 64), + ("mla", 2, 32), + # Mamba state is replicated, not DCP-sharded, and its width is the + # resident state block count rather than cdiv(max_len, block_size). + ("mamba", 2, 3), + ], +) +def test_uniform_type_spec_block_table_width_matches_layer_spec( + layer_type, dcp_size, expected_width +): + # The runner sizes the block table from the group spec while the metadata + # builders are constructed from the per-layer spec, so the aggregate must + # report the same width as the layers it wraps. + vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) + vllm_config.parallel_config.decode_context_parallel_size = dcp_size + layer_spec = new_mla_spec() if layer_type == "mla" else new_mamba_spec() + uniform_spec = UniformTypeKVCacheSpecs( + block_size=layer_spec.block_size, + kv_cache_specs={"layer1": layer_spec, "layer2": layer_spec}, + ) + + assert layer_spec.max_num_blocks_per_req(vllm_config, 1024) == expected_width + assert uniform_spec.max_num_blocks_per_req(vllm_config, 1024) == expected_width + + def test_merge_kv_cache_spec(): same_layer_specs = [ new_kv_cache_spec(num_kv_heads=32), diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index cc61f1f5ac..beb06b08f0 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -855,6 +855,20 @@ class UniformTypeKVCacheSpecs(KVCacheSpec): ) return max_num_pages * self.page_size_bytes + def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: + # Metadata builders are constructed from the per-layer spec, so the base + # cdiv(max_len, block_size) would drop its DCP sharding and size the + # block table wider than those builders expect. + widths = { + spec.max_num_blocks_per_req(vllm_config, max_len) + for spec in self.kv_cache_specs.values() + } + assert len(widths) == 1, ( + "All layers in the same KV cache group must need the same number " + f"of block table entries, got {sorted(widths)}." + ) + return next(iter(widths)) + @classmethod def is_uniform_type(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> bool: """