[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>

tests/v1/kv_connector/unit/offloading_connector/test_config.py

@@ -661,6 +661,34 @@ def test_canonical_layout_certifies_v2_model_runner():
     ).parallel.is_parallelism_agnostic
 
 
+def test_prefer_cross_layer_blocks_yields_to_canonical_layout():
+    """The connector must not request cross-layer blocks under
+    canonical_layout: cross-layer slabs have no per-layer refs to certify."""
+    from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
+    from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
+        OffloadingConnector,
+    )
+
+    kv_cache_config = _make_kv_cache_config()
+    connector_module = "vllm.distributed.kv_transfer.kv_connector.v1"
+
+    def make_connector(extra_config: dict[str, Any] | None) -> OffloadingConnector:
+        with (
+            patch(f"{connector_module}.offloading_connector.OffloadingSpecFactory"),
+            patch(
+                f"{connector_module}.offloading_connector.OffloadingConnectorScheduler"
+            ),
+        ):
+            return OffloadingConnector(
+                _make_vllm_config(extra_config=extra_config),
+                KVConnectorRole.SCHEDULER,
+                kv_cache_config,
+            )
+
+    assert make_connector(None).prefer_cross_layer_blocks
+    assert not make_connector({"canonical_layout": True}).prefer_cross_layer_blocks
+
+
 def test_parallelism_agnostic_disabled_on_v2_model_runner():
     config = _make_vllm_config()
     config.use_v2_model_runner = True

vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py

@@ -49,7 +49,8 @@ from vllm.v1.request import Request
 class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
     @property
     def prefer_cross_layer_blocks(self) -> bool:
-        return True
+        # Cross-layer slabs have no per-layer refs to certify canonical mappings
+        return not self._canonical_layout
 
     @property
     def requires_kv_delivery(self) -> bool:
@@ -66,6 +67,7 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA):
         super().__init__(vllm_config, role, kv_cache_config)
 
         offloading_config = build_offloading_config(vllm_config, kv_cache_config)
+        self._canonical_layout = offloading_config.canonical_layout
         spec = OffloadingSpecFactory.create_spec(offloading_config)
 
         self.connector_scheduler: OffloadingConnectorScheduler | None = None