[KV Connector][Offloading] Certify attention-only hybrids in the canonical portability gate (#51689)
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>
Itay Etelis committed
Aug 31, 2026 at 15:53 UTC
c5d840ff6a50f544fca8524e8caffda9f63b7728
2 files changed
+91
-91
tests/v1/kv_connector/unit/offloading_connector/test_config.py
+56
-66
@@ -23,6 +23,7 @@ from vllm.v1.kv_cache_interface import (
23
HiddenStateCacheSpec,
24
KVCacheConfig,
25
KVCacheGroupSpec,
26
+ KVCacheSpec,
27
KVCacheTensor,
28
MambaSpec,
29
MLAAttentionSpec,
@@ -247,8 +248,16 @@ def _make_mamba_hybrid_kv_cache_config() -> KVCacheConfig:
248
)
249
250
250
-def _parallelism_agnostic(kv_cache_groups: list[KVCacheGroupSpec]) -> bool:
251
- config = _make_vllm_config()
251
+def _parallelism_agnostic(
252
+ kv_cache_groups: list[KVCacheGroupSpec],
253
+ *,
254
+ canonical: bool = False,
255
+ v2: bool = False,
256
+) -> bool:
257
+ config = _make_vllm_config(
258
+ extra_config={"canonical_layout": True} if canonical else None
259
+ )
260
+ config.use_v2_model_runner = v2
261
kv_cache_config = KVCacheConfig(
262
num_blocks=0,
263
kv_cache_tensors=[],
@@ -641,85 +650,66 @@ def test_parallelism_agnostic_for_single_full_attention_group():
650
assert _parallelism_agnostic([KVCacheGroupSpec(["l0"], _full_attention_spec())])
651
652
653
+_SWA_SPEC = SlidingWindowSpec(
654
+ block_size=16,
655
+ num_kv_heads=4,
656
+ head_size=128,
657
+ dtype=torch.float32,
658
+ sliding_window=128,
659
+)
660
+_SWA_MLA_SPEC = SlidingWindowMLASpec(
661
+ block_size=16,
662
+ num_kv_heads=1,
663
+ head_size=576,
664
+ dtype=torch.float32,
665
+ sliding_window=128,
666
+)
667
+
668
+
669
+def _groups(*specs: KVCacheSpec) -> list[KVCacheGroupSpec]:
670
+ return [KVCacheGroupSpec([f"l{i}"], spec) for i, spec in enumerate(specs)]
671
+
672
+
673
@pytest.mark.parametrize(
674
"kv_cache_groups",
675
[
647
- [KVCacheGroupSpec(["l0"], _mla_spec(head_size=576))],
648
- [
649
- KVCacheGroupSpec(
650
- ["l0"],
651
- SlidingWindowSpec(
652
- block_size=16,
653
- num_kv_heads=4,
654
- head_size=128,
655
- dtype=torch.float32,
656
- sliding_window=128,
657
- ),
658
- )
659
- ],
660
- [
661
- KVCacheGroupSpec(["l0"], _full_attention_spec()),
662
- KVCacheGroupSpec(["l1"], _full_attention_spec()),
663
- ],
676
+ _groups(_mla_spec(head_size=576)),
677
+ _groups(_SWA_SPEC),
678
+ _groups(_full_attention_spec(), _full_attention_spec()),
679
],
680
)
681
def test_parallelism_agnostic_excluded(kv_cache_groups: list[KVCacheGroupSpec]):
682
assert not _parallelism_agnostic(kv_cache_groups)
683
684
670
-def test_canonical_layout_widens_parallelism_agnostic_to_mla():
671
- """The canonical layout dedups the TP-replicated MLA latent into one
672
- portable copy, so the gate admits MLA — but only when canonical_layout
673
- is requested."""
674
- mla_groups = [KVCacheGroupSpec(["l0"], _mla_spec(head_size=576))]
675
- assert not _parallelism_agnostic(mla_groups)
676
-
677
- config = _make_vllm_config(extra_config={"canonical_layout": True})
678
- kv_cache_config = KVCacheConfig(
679
- num_blocks=0,
680
- kv_cache_tensors=[],
681
- kv_cache_groups=mla_groups,
682
- )
683
- offloading_config = build_offloading_config(config, kv_cache_config)
684
- assert offloading_config.parallel.is_parallelism_agnostic
685
- assert offloading_config.canonical_layout
686
-
687
- # hybrid groupings stay out: their non-full-attention layers can only
688
- # derive opaque (exact-topology) mappings
689
- hybrid_config = KVCacheConfig(
690
- num_blocks=0,
691
- kv_cache_tensors=[],
692
- kv_cache_groups=[
693
- KVCacheGroupSpec(["l0"], _full_attention_spec()),
694
- KVCacheGroupSpec(["l1"], _full_attention_spec()),
695
- ],
696
- )
697
- assert not build_offloading_config(
698
- config, hybrid_config
699
- ).parallel.is_parallelism_agnostic
685
+@pytest.mark.parametrize(
686
+ ("kv_cache_groups", "certified"),
687
+ [
688
+ pytest.param(_groups(_mla_spec(head_size=576)), True, id="mla-latent"),
689
+ pytest.param(
690
+ _groups(_full_attention_spec(), _SWA_SPEC), True, id="attention-hybrid"
691
+ ),
692
+ pytest.param(_groups(_SWA_MLA_SPEC), False, id="swa-mla"),
693
+ pytest.param(
694
+ list(_make_mamba_hybrid_kv_cache_config().kv_cache_groups),
695
+ False,
696
+ id="mamba-hybrid",
697
+ ),
698
+ ],
699
+)
700
+def test_canonical_layout_gate(kv_cache_groups, certified):
701
+ """The canonical layout certifies portability group by group; none of
702
+ these shapes are portable in the direct layout."""
703
+ assert not _parallelism_agnostic(kv_cache_groups)
704
+ assert _parallelism_agnostic(kv_cache_groups, canonical=True) is certified
705
706
707
def test_canonical_layout_certifies_v2_model_runner():
708
"""Canonical bytes are certified per layer against live tensor strides at
709
registration, so the static gate must not depend on the model-runner
710
version — the v2 runner is the case the canonical layout exists for."""
706
- kv_cache_config = KVCacheConfig(
707
- num_blocks=0,
708
- kv_cache_tensors=[],
709
- kv_cache_groups=[KVCacheGroupSpec(["l0"], _full_attention_spec())],
710
- )
711
-
712
- config = _make_vllm_config()
713
- config.use_v2_model_runner = True
714
- assert not build_offloading_config(
715
- config, kv_cache_config
716
- ).parallel.is_parallelism_agnostic
717
-
718
- config = _make_vllm_config(extra_config={"canonical_layout": True})
719
- config.use_v2_model_runner = True
720
- assert build_offloading_config(
721
- config, kv_cache_config
722
- ).parallel.is_parallelism_agnostic
711
+ groups = _groups(_full_attention_spec())
712
+ assert _parallelism_agnostic(groups, canonical=True, v2=True)
713
714
715
def test_parallelism_agnostic_disabled_on_v2_model_runner():
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+35
-25
@@ -8,7 +8,10 @@ from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes
8
from vllm.v1.kv_cache_interface import (
9
AttentionSpec,
10
FullAttentionSpec,
11
+ KVCacheSpec,
12
MLAAttentionSpec,
13
+ SlidingWindowMLASpec,
14
+ SlidingWindowSpec,
15
)
16
from vllm.v1.kv_offload.config import (
17
OffloadingCacheConfig,
@@ -139,35 +142,42 @@ def build_offloading_config(
142
and parallel_config.decode_context_parallel_size == 1
143
and parallel_config.prefill_context_parallel_size == 1
144
)
142
- # Canonical pages are topology-free by construction, so the canonical
143
- # layout widens the gate to every config whose mappings derive portable:
144
- # exactly-sharded or replicated GQA heads (writer rotation) and the
145
- # TP-replicated MLA latent (one canonical copy for all replicas). The
146
- # model-runner version is irrelevant here — certification happens per
147
- # layer against live tensor strides at registration, and create_worker
148
- # fails closed against this flag if any layer cannot be certified.
145
+ # Canonical pages are topology-free, so the gate widens to every config
146
+ # whose mappings derive portable, group by group; certification happens
147
+ # per layer at registration and create_worker fails closed on this flag.
148
if canonical_layout and not is_parallelism_agnostic:
149
tp_size = parallel_config.tensor_parallel_size
151
- if isinstance(single_group_spec, FullAttentionSpec):
152
- if type(single_group_spec) is MLAAttentionSpec:
153
- spec_certifiable = (
154
- single_group_spec.tokens_per_state == 1
155
- and single_group_spec.real_page_size_bytes
156
- % single_group_spec.block_size
157
- == 0
150
+ total_kv_heads = vllm_config.model_config.get_total_num_kv_heads()
151
+
152
+ def spec_certifiable(spec: KVCacheSpec) -> bool:
153
+ """Conservative static mirror of _layer_mapping's per-layer checks."""
154
+ if not isinstance(spec, AttentionSpec):
155
+ return False
156
+ if spec.kv_quant_mode.is_per_token_head:
157
+ return False
158
+ if type(spec) is MLAAttentionSpec:
159
+ return (
160
+ spec.tokens_per_state == 1
161
+ and spec.real_page_size_bytes % spec.block_size == 0
162
)
159
- else:
160
- total_kv_heads = vllm_config.model_config.get_total_num_kv_heads()
161
- spec_certifiable = (
162
- total_kv_heads % tp_size == 0 or tp_size % total_kv_heads == 0
163
- )
164
- is_parallelism_agnostic = (
165
- spec_certifiable
166
- and not single_group_spec.kv_quant_mode.is_per_token_head
167
- and parallel_config.decode_context_parallel_size == 1
168
- and parallel_config.prefill_context_parallel_size == 1
169
- and parallel_config.world_size == tp_size
163
+ if isinstance(spec, (SlidingWindowMLASpec, MLAAttentionSpec)):
164
+ return False
165
+ if not isinstance(spec, (FullAttentionSpec, SlidingWindowSpec)):
166
+ return False
167
+ return (
168
+ total_kv_heads % tp_size == 0 or tp_size % total_kv_heads == 0
169
+ ) and spec.num_kv_heads == max(1, total_kv_heads // tp_size)
170
+
171
+ 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
)
177
+ and parallel_config.decode_context_parallel_size == 1
178
+ and parallel_config.prefill_context_parallel_size == 1
179
+ and parallel_config.world_size == tp_size
180
+ )
181
182
kv_events_config = vllm_config.kv_events_config
183
cache_dtype = (