[Bugfix] Shard UniformTypeKVCacheSpecs block table width under DCP (#50823)
Signed-off-by: Mikhail Kostryukov <mike@triptrack.net> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
drakosha committed
Aug 3, 2026 at 20:34 UTC
f0de1a604cad003379e5bb4dfc3cc5d2a1f25fa8
2 files changed
+42
tests/v1/core/test_kv_cache_utils.py
+28
@@ -1257,6 +1257,34 @@ def test_project_kv_cache_groups_to_worker():
1257
assert set(proj_spec.kv_cache_specs.keys()) == {"layer1", "layer3"}
1258
1259
1260
+@pytest.mark.parametrize(
1261
+ "layer_type,dcp_size,expected_width",
1262
+ [
1263
+ ("mla", 1, 64),
1264
+ ("mla", 2, 32),
1265
+ # Mamba state is replicated, not DCP-sharded, and its width is the
1266
+ # resident state block count rather than cdiv(max_len, block_size).
1267
+ ("mamba", 2, 3),
1268
+ ],
1269
+)
1270
+def test_uniform_type_spec_block_table_width_matches_layer_spec(
1271
+ layer_type, dcp_size, expected_width
1272
+):
1273
+ # The runner sizes the block table from the group spec while the metadata
1274
+ # builders are constructed from the per-layer spec, so the aggregate must
1275
+ # report the same width as the layers it wraps.
1276
+ vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
1277
+ vllm_config.parallel_config.decode_context_parallel_size = dcp_size
1278
+ layer_spec = new_mla_spec() if layer_type == "mla" else new_mamba_spec()
1279
+ uniform_spec = UniformTypeKVCacheSpecs(
1280
+ block_size=layer_spec.block_size,
1281
+ kv_cache_specs={"layer1": layer_spec, "layer2": layer_spec},
1282
+ )
1283
+
1284
+ assert layer_spec.max_num_blocks_per_req(vllm_config, 1024) == expected_width
1285
+ assert uniform_spec.max_num_blocks_per_req(vllm_config, 1024) == expected_width
1286
+
1287
+
1288
def test_merge_kv_cache_spec():
1289
same_layer_specs = [
1290
new_kv_cache_spec(num_kv_heads=32),
vllm/v1/kv_cache_interface.py
+14
@@ -855,6 +855,20 @@ class UniformTypeKVCacheSpecs(KVCacheSpec):
855
)
856
return max_num_pages * self.page_size_bytes
857
858
+ def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
859
+ # Metadata builders are constructed from the per-layer spec, so the base
860
+ # cdiv(max_len, block_size) would drop its DCP sharding and size the
861
+ # block table wider than those builders expect.
862
+ widths = {
863
+ spec.max_num_blocks_per_req(vllm_config, max_len)
864
+ for spec in self.kv_cache_specs.values()
865
+ }
866
+ assert len(widths) == 1, (
867
+ "All layers in the same KV cache group must need the same number "
868
+ f"of block table entries, got {sorted(widths)}."
869
+ )
870
+ return next(iter(widths))
871
+
872
@classmethod
873
def is_uniform_type(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> bool:
874
"""