[KV Connector][Offloading] Look through UniformTypeKVCacheSpecs in the canonical portability gate (#51690)

Signed-off-by: Itay Etelis <itay.etelis@ibm.com>
Co-authored-by: Itay Etelis <itay.etelis@ibm.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

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):

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