diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 6a0b6393a225..27019173fc17 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/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, + } diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index 7e0a7fb6446d..a96850bd7e75 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/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] diff --git a/vllm/v1/kv_cache_layout.py b/vllm/v1/kv_cache_layout.py index 7c1decd393a9..f31f4656335e 100644 --- a/vllm/v1/kv_cache_layout.py +++ b/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.