diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index 99d2ccfb82..5e2ef5acf1 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -153,12 +153,14 @@ def _mla_spec( block_size: int = 16, head_size: int = 512, dtype: torch.dtype = torch.float32, + tokens_per_state: int = 1, ) -> MLAAttentionSpec: return MLAAttentionSpec( block_size=block_size, num_kv_heads=1, head_size=head_size, dtype=dtype, + tokens_per_state=tokens_per_state, ) @@ -670,6 +672,15 @@ def _groups(*specs: KVCacheSpec) -> list[KVCacheGroupSpec]: return [KVCacheGroupSpec([f"l{i}"], spec) for i, spec in enumerate(specs)] +def _uniform_group(*specs: KVCacheSpec) -> KVCacheGroupSpec: + """GLM-5.2/DSv3.2-style wrapper: same-type layers whose specs differ.""" + names = [f"l{i}" for i in range(len(specs))] + return KVCacheGroupSpec( + names, + UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=dict(zip(names, specs))), + ) + + @pytest.mark.parametrize( "kv_cache_groups", [ @@ -695,6 +706,17 @@ def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]): False, id="mamba-hybrid", ), + pytest.param( + [_uniform_group(_mla_spec(head_size=576), _mla_spec(head_size=128))], + True, + id="uniform-mla-wrapper", + ), + pytest.param( + [_uniform_group(_mla_spec(head_size=576), _mla_spec(tokens_per_state=2))], + False, + id="uniform-uncertifiable-inner", + ), + pytest.param([], False, id="attention-free"), ], ) def test_canonical_layout_gate(kv_cache_groups, certified): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index 1faa4f5495..4fd2b972c7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -12,6 +12,7 @@ from vllm.v1.kv_cache_interface import ( MLAAttentionSpec, SlidingWindowMLASpec, SlidingWindowSpec, + iter_layer_specs, ) from vllm.v1.kv_offload.config import ( OffloadingCacheConfig, @@ -168,12 +169,16 @@ def build_offloading_config( total_kv_heads % tp_size == 0 or tp_size % total_kv_heads == 0 ) and spec.num_kv_heads == max(1, total_kv_heads // tp_size) + # UniformTypeKVCacheSpecs groups (e.g. MLA plus its DSA indexer) hold + # one spec per layer; certify per layer, as the mapping derivation does. + layer_specs = [ + spec + for group in kv_cache_config.kv_cache_groups + for spec in iter_layer_specs(group.kv_cache_spec) + ] is_parallelism_agnostic = ( - len(kv_cache_config.kv_cache_groups) > 0 - and all( - spec_certifiable(group.kv_cache_spec) - for group in kv_cache_config.kv_cache_groups - ) + len(layer_specs) > 0 + and all(spec_certifiable(spec) for spec in layer_specs) and parallel_config.decode_context_parallel_size == 1 and parallel_config.prefill_context_parallel_size == 1 and parallel_config.world_size == tp_size