[KV Connector][Offloading] Keep per-layer KV registration when canonical_layout is requested (#51688)
Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com>
Itay Etelis committed
Aug 11, 2026 at 06:03 UTC
0f2ea973beebfe2959c74f06e2a412c55c1c7007
2 files changed
+31
-1
tests/v1/kv_connector/unit/offloading_connector/test_config.py
+28
@@ -661,6 +661,34 @@ def test_canonical_layout_certifies_v2_model_runner():
661
).parallel.is_parallelism_agnostic
662
663
664
+def test_prefer_cross_layer_blocks_yields_to_canonical_layout():
665
+ """The connector must not request cross-layer blocks under
666
+ canonical_layout: cross-layer slabs have no per-layer refs to certify."""
667
+ from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
668
+ from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
669
+ OffloadingConnector,
670
+ )
671
+
672
+ kv_cache_config = _make_kv_cache_config()
673
+ connector_module = "vllm.distributed.kv_transfer.kv_connector.v1"
674
+
675
+ def make_connector(extra_config: dict[str, Any] | None) -> OffloadingConnector:
676
+ with (
677
+ patch(f"{connector_module}.offloading_connector.OffloadingSpecFactory"),
678
+ patch(
679
+ f"{connector_module}.offloading_connector.OffloadingConnectorScheduler"
680
+ ),
681
+ ):
682
+ return OffloadingConnector(
683
+ _make_vllm_config(extra_config=extra_config),
684
+ KVConnectorRole.SCHEDULER,
685
+ kv_cache_config,
686
+ )
687
+
688
+ assert make_connector(None).prefer_cross_layer_blocks
689
+ assert not make_connector({"canonical_layout": True}).prefer_cross_layer_blocks
690
+
691
+
692
def test_parallelism_agnostic_disabled_on_v2_model_runner():
693
config = _make_vllm_config()
694
config.use_v2_model_runner = True
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+3
-1
@@ -49,7 +49,8 @@ from vllm.v1.request import Request
49
class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
50
@property
51
def prefer_cross_layer_blocks(self) -> bool:
52
- return True
52
+ # Cross-layer slabs have no per-layer refs to certify canonical mappings
53
+ return not self._canonical_layout
54
55
@property
56
def requires_kv_delivery(self) -> bool:
@@ -66,6 +67,7 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
67
super().__init__(vllm_config, role, kv_cache_config)
68
69
offloading_config = build_offloading_config(vllm_config, kv_cache_config)
70
+ self._canonical_layout = offloading_config.canonical_layout
71
spec = OffloadingSpecFactory.create_spec(offloading_config)
72
73
self.connector_scheduler: OffloadingConnectorScheduler | None = None