[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>
Itay Etelis committed
Sep 2, 2026 at 10:43 UTC
df09c767355fd1313001fbf14ad363bba5bf0ffd
2 files changed
+32
-5
tests/v1/kv_connector/unit/offloading_connector/test_config.py
+22
@@ -153,12 +153,14 @@ def _mla_spec(
153
block_size: int = 16,
154
head_size: int = 512,
155
dtype: torch.dtype = torch.float32,
156
+ tokens_per_state: int = 1,
157
) -> MLAAttentionSpec:
158
return MLAAttentionSpec(
159
block_size=block_size,
160
num_kv_heads=1,
161
head_size=head_size,
162
dtype=dtype,
163
+ tokens_per_state=tokens_per_state,
164
)
165
166
@@ -670,6 +672,15 @@ def _groups(*specs: KVCacheSpec) -> list[KVCacheGroupSpec]:
672
return [KVCacheGroupSpec([f"l{i}"], spec) for i, spec in enumerate(specs)]
673
674
675
+def _uniform_group(*specs: KVCacheSpec) -> KVCacheGroupSpec:
676
+ """GLM-5.2/DSv3.2-style wrapper: same-type layers whose specs differ."""
677
+ names = [f"l{i}" for i in range(len(specs))]
678
+ return KVCacheGroupSpec(
679
+ names,
680
+ UniformTypeKVCacheSpecs(block_size=16, kv_cache_specs=dict(zip(names, specs))),
681
+ )
682
+
683
+
684
@pytest.mark.parametrize(
685
"kv_cache_groups",
686
[
@@ -695,6 +706,17 @@ def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]):
706
False,
707
id="mamba-hybrid",
708
),
709
+ pytest.param(
710
+ [_uniform_group(_mla_spec(head_size=576), _mla_spec(head_size=128))],
711
+ True,
712
+ id="uniform-mla-wrapper",
713
+ ),
714
+ pytest.param(
715
+ [_uniform_group(_mla_spec(head_size=576), _mla_spec(tokens_per_state=2))],
716
+ False,
717
+ id="uniform-uncertifiable-inner",
718
+ ),
719
+ pytest.param([], False, id="attention-free"),
720
],
721
)
722
def test_canonical_layout_gate(kv_cache_groups, certified):
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+10
-5
@@ -12,6 +12,7 @@ from vllm.v1.kv_cache_interface import (
12
MLAAttentionSpec,
13
SlidingWindowMLASpec,
14
SlidingWindowSpec,
15
+ iter_layer_specs,
16
)
17
from vllm.v1.kv_offload.config import (
18
OffloadingCacheConfig,
@@ -168,12 +169,16 @@ def build_offloading_config(
169
total_kv_heads % tp_size == 0 or tp_size % total_kv_heads == 0
170
) and spec.num_kv_heads == max(1, total_kv_heads // tp_size)
171
172
+ # UniformTypeKVCacheSpecs groups (e.g. MLA plus its DSA indexer) hold
173
+ # one spec per layer; certify per layer, as the mapping derivation does.
174
+ layer_specs = [
175
+ spec
176
+ for group in kv_cache_config.kv_cache_groups
177
+ for spec in iter_layer_specs(group.kv_cache_spec)
178
+ ]
179
is_parallelism_agnostic = (
172
- len(kv_cache_config.kv_cache_groups) > 0
173
- and all(
174
- spec_certifiable(group.kv_cache_spec)
175
- for group in kv_cache_config.kv_cache_groups
176
- )
180
+ len(layer_specs) > 0
181
+ and all(spec_certifiable(spec) for spec in layer_specs)
182
and parallel_config.decode_context_parallel_size == 1
183
and parallel_config.prefill_context_parallel_size == 1
184
and parallel_config.world_size == tp_size