[Bugfix][KV Offload] Scale UniformTypeKVCacheSpecs groups by DCP (#50883)
Signed-off-by: Mikhail Kostryukov <mike@triptrack.net> Co-authored-by: Claude <noreply@anthropic.com>
drakosha committed
Sep 2, 2026 at 17:11 UTC
ba6c60e98fd7ae91852ad81fa880d3f9153639c3
2 files changed
+82
-8
tests/v1/kv_connector/unit/offloading_connector/test_config.py
+75
@@ -231,6 +231,26 @@ def _make_hybrid_kv_cache_config() -> KVCacheConfig:
231
)
232
233
234
+def _mamba_spec() -> MambaSpec:
235
+ return MambaSpec(
236
+ block_size=16,
237
+ shapes=((1, 1),),
238
+ dtypes=(torch.float32,),
239
+ mamba_cache_mode="align",
240
+ )
241
+
242
+
243
+def _uniform_spec(spec_kind: str) -> UniformTypeKVCacheSpecs:
244
+ # DSA models merge their indexer and MLA layers into a UniformType group,
245
+ # a container rather than an AttentionSpec, but still sharded across DCP.
246
+ specs: dict[str, Any] = (
247
+ {"mla_layer": _mla_spec(), "indexer_layer": _mla_spec(head_size=128)}
248
+ if spec_kind == "attention"
249
+ else {"mamba_layer0": _mamba_spec(), "mamba_layer1": _mamba_spec()}
250
+ )
251
+ return UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=specs)
252
+
253
+
254
def _make_mamba_hybrid_kv_cache_config() -> KVCacheConfig:
255
return KVCacheConfig(
256
num_blocks=4,
@@ -369,6 +389,61 @@ def test_dcp_scales_attention_but_not_mamba_group_blocks():
389
] == [1, 3]
390
391
392
+@pytest.mark.parametrize("dcp_size,expected", [(1, 16), (2, 32)])
393
+def test_dcp_scales_uniform_type_attention_group_blocks(dcp_size, expected):
394
+ config = _make_vllm_config(
395
+ tensor_parallel_size=2, decode_context_parallel_size=dcp_size
396
+ )
397
+ kv_cache_config = KVCacheConfig(
398
+ num_blocks=4,
399
+ kv_cache_tensors=[
400
+ KVCacheTensor(
401
+ size=64,
402
+ layers=["mla_layer", "indexer_layer"],
403
+ layer_stride=32,
404
+ block_stride=8,
405
+ )
406
+ ],
407
+ kv_cache_groups=[
408
+ KVCacheGroupSpec(["mla_layer", "indexer_layer"], _uniform_spec("attention"))
409
+ ],
410
+ )
411
+
412
+ offloading_config = build_offloading_config(config, kv_cache_config)
413
+
414
+ assert offloading_config.groups[0].tokens_per_block == expected
415
+ assert offloading_config.cache.tokens_per_hash == expected
416
+
417
+
418
+@pytest.mark.parametrize(
419
+ "spec_kind,expected",
420
+ [
421
+ ("attention", 32),
422
+ # Mamba state is replicated across DCP ranks, so a uniform group of
423
+ # Mamba layers keeps its span.
424
+ ("mamba", 16),
425
+ ],
426
+)
427
+def test_dcp_scales_uniform_type_group_alongside_mamba(spec_kind, expected):
428
+ config = _make_vllm_config(tensor_parallel_size=2, decode_context_parallel_size=2)
429
+ config.speculative_config = None
430
+ kv_cache_config = KVCacheConfig(
431
+ num_blocks=4,
432
+ kv_cache_tensors=[],
433
+ kv_cache_groups=[
434
+ KVCacheGroupSpec(["layer0", "layer1"], _uniform_spec(spec_kind)),
435
+ KVCacheGroupSpec(["mamba_layer"], _mamba_spec()),
436
+ ],
437
+ )
438
+
439
+ offloading_config = build_offloading_config(config, kv_cache_config)
440
+
441
+ assert tuple(group.tokens_per_block for group in offloading_config.groups) == (
442
+ expected,
443
+ 16,
444
+ )
445
+
446
+
447
def test_preserves_data_parallel_config():
448
config = _make_vllm_config()
449
config.parallel_config.data_parallel_index = 2
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+7
-8
@@ -4,7 +4,10 @@
4
5
from typing import TYPE_CHECKING
6
7
-from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes
7
+from vllm.v1.core.kv_cache_utils import (
8
+ resolve_dcp_kv_block_size,
9
+ resolve_kv_cache_block_sizes,
10
+)
11
from vllm.v1.kv_cache_interface import (
12
AttentionSpec,
13
FullAttentionSpec,
@@ -41,13 +44,9 @@ def build_offloading_config(
44
parallel_config = vllm_config.parallel_config
45
groups = tuple(
46
OffloadingGroupConfig(
44
- tokens_per_block=(
45
- group.kv_cache_spec.block_size
46
- * (
47
- parallel_config.decode_context_parallel_size
48
- if isinstance(group.kv_cache_spec, AttentionSpec)
49
- else 1
50
- )
47
+ tokens_per_block=resolve_dcp_kv_block_size(
48
+ group.kv_cache_spec,
49
+ parallel_config.decode_context_parallel_size,
50
),
51
layer_names=tuple(group.layer_names),
52
)