From 0f2ea973beebfe2959c74f06e2a412c55c1c7007 Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Tue, 11 Aug 2026 06:03:59 +0300 Subject: [PATCH] [KV Connector][Offloading] Keep per-layer KV registration when canonical_layout is requested (#51688) Signed-off-by: Itay Etelis Co-authored-by: Itay Etelis --- .../unit/offloading_connector/test_config.py | 28 +++++++++++++++++++ .../kv_connector/v1/offloading_connector.py | 4 ++- 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index cc58aaf7a7..0b0b6b874f 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/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 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index a32d5d9cc8..f003d0d6f8 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/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