[Core][KV Connector] Keep the connector's head order when its KV layout is unsupported

tests/v1/core/test_kv_cache_utils.py

@@ -2174,3 +2174,34 @@ def test_null_head_len():
         assert null_head_len(table) == expected
     blocks = KVCacheBlocks(([null] * 4 + real, [null] + real[:1] + [null]))
     assert blocks.get_unhashed_block_ids_all_groups() == [[1, 2, 3], [1]]
+
+
+@pytest.mark.parametrize(
+    "supported,expected",
+    [
+        # A hybrid model's backends only support block-outer layouts.
+        (["BLNHC", "BLHNC"], KVCacheLayout.BLHNC),
+        # No head-major candidate: the backends' first choice.
+        (["BLNHC"], KVCacheLayout.BLNHC),
+    ],
+)
+def test_connector_layout_fallback_keeps_head_order(monkeypatch, supported, expected):
+    """A connector that slices by head keeps heads outside tokens when it can."""
+    from vllm.v1.attention.backends import utils as backend_utils
+
+    monkeypatch.delenv("VLLM_KV_CACHE_LAYOUT", raising=False)
+    monkeypatch.setattr(
+        backend_utils, "get_kv_connector_cache_layout", lambda _: "LBHNC"
+    )
+    config = VllmConfig(model_config=ModelConfig(max_model_len=1024))
+    config.cache_config.kv_cache_layout = None
+    assert backend_utils.resolve_kv_cache_layout(config, [supported]) == expected
+
+
+def test_kv_cache_layout_is_head_major():
+    assert {layout for layout in KVCacheLayout if layout.is_head_major} == {
+        KVCacheLayout.LBHNC,
+        KVCacheLayout.LHBNC,
+        KVCacheLayout.BLHNC,
+        KVCacheLayout.BHLNC,
+    }

vllm/v1/attention/backends/utils.py

@@ -270,11 +270,11 @@ def resolve_kv_cache_layout(
     shapes narrow the candidates to block-compact layouts. An explicit
     ``VLLM_KV_CACHE_LAYOUT`` must be one of the candidates or resolution fails,
     with the legacy ``NHD``/``HND`` names as aliases for ``LBNHC``/``LBHNC``; the
-    connector's preference is used when compatible and dropped with a warning
-    otherwise. A layout already present on ``cache_config`` wins outright, and
-    the result is recorded there (see ``CacheConfig.kv_cache_layout``); it
-    reaches workers through the ``set_kv_cache_layout`` RPC and
-    ``KVCacheConfig.kv_cache_layout``.
+    connector's preference is used when compatible and otherwise replaced, with a
+    warning, by the first candidate with the same head/token order. A layout
+    already present on ``cache_config`` wins outright, and the result is recorded
+    there (see ``CacheConfig.kv_cache_layout``); it reaches workers through the
+    ``set_kv_cache_layout`` RPC and ``KVCacheConfig.kv_cache_layout``.
     """
     cache_config = vllm_config.cache_config
     if cache_config.kv_cache_layout is not None:
@@ -336,12 +336,18 @@ def resolve_kv_cache_layout(
     elif (connector := get_kv_connector_cache_layout(vllm_config)) is not None:
         layout = _layout_from_name(connector)
         if layout not in candidates:
+            # Connectors that slice pages by head (heterogeneous TP) need heads
+            # outside tokens; the layer and block order can follow the backends.
+            fallback = next(
+                (c for c in candidates if c.is_head_major == layout.is_head_major),
+                candidates[0],
+            )
             logger.warning_once(
                 f"KV connector cache layout {connector} does not satisfy every "
                 f"supported set; valid layouts: {_layout_names(candidates)}. "
-                f"Using {candidates[0].name} instead."
+                f"Using {fallback.name} instead."
             )
-            layout = candidates[0]
+            layout = fallback
     else:
         layout = candidates[0]
 

vllm/v1/kv_cache_layout.py

@@ -46,6 +46,11 @@ def is_block_contiguous(self) -> bool:
         """True when [H, N, C] is contiguous within a block."""
         return self.value[-3:] == (_DIM_H, _DIM_N, _DIM_C)
 
+    @property
+    def is_head_major(self) -> bool:
+        """True when H is outside N, so each head's tokens are contiguous."""
+        return self.value.index(_DIM_H) < self.value.index(_DIM_N)
+
     @property
     def is_block_compact(self) -> bool:
         """True when each page's [H, N, C] bytes form one contiguous run; i.e.