From d8f483dc30a9a74b8fdf1c6dcb11cff503fe40a3 Mon Sep 17 00:00:00 2001 From: Igor Margulis Date: Tue, 30 Jun 2026 18:19:51 +0300 Subject: [PATCH] [Spec Decode] Fix hidden-state extraction block size for hybrid verifiers (#46301) Signed-off-by: Igor Margulis Signed-off-by: mgoin Co-authored-by: Cursor Co-authored-by: mgoin --- .../unit/test_hidden_states_connector.py | 126 ++++++++++++++++++ .../v1/example_hidden_states_connector.py | 73 +++++++--- .../models/extract_hidden_states.py | 3 + 3 files changed, 187 insertions(+), 15 deletions(-) create mode 100644 tests/v1/kv_connector/unit/test_hidden_states_connector.py diff --git a/tests/v1/kv_connector/unit/test_hidden_states_connector.py b/tests/v1/kv_connector/unit/test_hidden_states_connector.py new file mode 100644 index 0000000000..ffd6482891 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_hidden_states_connector.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""CPU-only unit tests for ExampleHiddenStatesConnector KV-cache-group logic.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.distributed.kv_transfer.kv_connector.v1.example_hidden_states_connector import ( # noqa: E501 + ExampleHiddenStatesConnector, +) +from vllm.v1.core.kv_cache_utils import get_kv_cache_groups +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + HiddenStateCacheSpec, + KVCacheGroupSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, +) + + +def _full(block_size: int) -> FullAttentionSpec: + return FullAttentionSpec( + block_size=block_size, num_kv_heads=8, head_size=128, dtype=torch.bfloat16 + ) + + +def _hidden(block_size: int) -> HiddenStateCacheSpec: + return HiddenStateCacheSpec( + block_size=block_size, num_kv_heads=6, head_size=2048, dtype=torch.bfloat16 + ) + + +def _config(*specs): + """Minimal stand-in exposing only ``kv_cache_groups`` (all the helpers read).""" + return SimpleNamespace( + kv_cache_groups=[ + KVCacheGroupSpec(layer_names=[f"layer.{i}"], kv_cache_spec=spec) + for i, spec in enumerate(specs) + ] + ) + + +# ---- _find_cache_kv_group_id ------------------------------------------------ + + +def test_find_group_id_none_config_returns_zero(): + assert ExampleHiddenStatesConnector._find_cache_kv_group_id(None) == 0 + + +def test_find_group_id_single_non_hidden_group_returns_zero(): + # Uniform (dense) model: one group, no HiddenStateCacheSpec -> group 0. + cfg = _config(_full(16)) + assert ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == 0 + + +def test_find_group_id_locates_hidden_group_when_not_first(): + # Hybrid layout: the hidden-states group is not group 0. + cfg = _config(_full(528), _hidden(22), _full(528)) + assert ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == 1 + + +def test_find_group_id_locates_hidden_group_last(): + cfg = _config(_full(528), _full(528), _hidden(22)) + assert ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == 2 + + +def test_find_group_id_raises_when_no_hidden_group_and_multiple_groups(): + cfg = _config(_full(16), _full(16)) + with pytest.raises(ValueError, match="Could not uniquely identify"): + ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) + + +def test_find_group_id_raises_when_multiple_hidden_groups(): + cfg = _config(_hidden(22), _hidden(22)) + with pytest.raises(ValueError, match="Could not uniquely identify"): + ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) + + +# ---- _get_cache_block_size -------------------------------------------------- + + +def test_get_block_size_reads_hidden_group_spec_not_global(): + # Hidden group keeps block size 22; the global is bumped to 528 for hybrids. + vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=528)) + cfg = _config(_full(528), _hidden(22)) + block_size = ExampleHiddenStatesConnector._get_cache_block_size( + vllm_config, cfg, cache_kv_group_id=1 + ) + assert block_size == 22 + + +def test_get_block_size_falls_back_to_cache_config_when_no_kv_cache_config(): + vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=16)) + block_size = ExampleHiddenStatesConnector._get_cache_block_size( + vllm_config, None, cache_kv_group_id=0 + ) + assert block_size == 16 + + +# ---- MLA-verifier absorption ------------------------------------------------ + + +def test_find_group_id_errors_clearly_when_absorbed_by_mla_swa_verifier(): + # HiddenStateCacheSpec subclasses MLAAttentionSpec, so an MLA + sliding- + # window MLA verifier absorbs it into the MLA group instead of isolating it. + dt = torch.bfloat16 + spec = { + "layers.0.mla": MLAAttentionSpec( + block_size=64, num_kv_heads=1, head_size=576, dtype=dt + ), + "layers.1.swa": SlidingWindowMLASpec( + block_size=64, num_kv_heads=1, head_size=576, dtype=dt, sliding_window=512 + ), + "cache_only_layers.61": _hidden(64), + } + vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), + speculative_config=None, + ) + groups = get_kv_cache_groups(vllm_config, spec) + assert not any(isinstance(g.kv_cache_spec, HiddenStateCacheSpec) for g in groups) + cfg = SimpleNamespace(kv_cache_groups=groups) + with pytest.raises(ValueError, match="MLA verifiers are unsupported"): + ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py index 7e6c95bf8f..299ff037ad 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py @@ -109,6 +109,49 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): # Must be False so that drafter kv cache isn't merged with verifier's return False + @classmethod + def _find_cache_kv_group_id(cls, kv_cache_config: "KVCacheConfig | None") -> int: + """Index of the KV cache group holding the extracted hidden states. + + Located by spec type so it resolves on both scheduler and worker side. + """ + if kv_cache_config is None: + return 0 + + from vllm.v1.kv_cache_interface import HiddenStateCacheSpec + + groups = kv_cache_config.kv_cache_groups + group_ids = [ + gid + for gid, group in enumerate(groups) + if isinstance(group.kv_cache_spec, HiddenStateCacheSpec) + ] + if len(group_ids) == 1: + return group_ids[0] + if not group_ids and len(groups) == 1: + return 0 + raise ValueError( + "Could not uniquely identify the extract-hidden-states KV cache " + f"group among {len(groups)} groups; the hidden-states layer must be " + "isolated in its own group (MLA verifiers are unsupported)." + ) + + @staticmethod + def _get_cache_block_size( + vllm_config: "VllmConfig", + kv_cache_config: "KVCacheConfig | None", + cache_kv_group_id: int, + ) -> int: + """Block size of the hidden-states group, read from its own spec. + + cache_config.block_size is bumped to a common multiple for hybrid + verifiers; the page-aligned hidden-states group keeps a smaller one. + """ + if kv_cache_config is None: + return vllm_config.cache_config.block_size + cache_group = kv_cache_config.kv_cache_groups[cache_kv_group_id] + return cache_group.kv_cache_spec.block_size + def __init__( self, vllm_config: "VllmConfig", @@ -120,7 +163,12 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): role=role, kv_cache_config=kv_cache_config, ) - self._block_size = vllm_config.cache_config.block_size + # Read the hidden-states group and its block size from the group spec; + # cache_config.block_size is bumped (wrong) for hybrid verifiers. + self._cache_kv_group_id = self._find_cache_kv_group_id(kv_cache_config) + self._block_size = self._get_cache_block_size( + vllm_config, kv_cache_config, self._cache_kv_group_id + ) self._storage_path = self._kv_transfer_config.get_from_extra_config( "shared_storage_path", "/tmp" ) @@ -151,13 +199,6 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): # Worker-side state (set by register_kv_caches). self._kv_cache: torch.Tensor | None = None - # Identify which KV cache group holds the hidden-states layer. - self._hs_group_idx: int = 0 - if self._kv_cache_config is not None: - for i, group in enumerate(self._kv_cache_config.kv_cache_groups): - if any("cache_only_layers" in n for n in group.layer_names): - self._hs_group_idx = i - break # Only TP rank 0 writes hidden states to disk; other TP ranks no-op. # Set in register_kv_caches (after distributed init). self._is_tp_rank_zero: bool = True @@ -267,12 +308,14 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): ) self._kv_cache = kv_caches[self.cache_layers[0]] - # Find the KV cache group index for hidden states - if self._kv_cache_config is not None: - for i, group in enumerate(self._kv_cache_config.kv_cache_groups): - if self.cache_layers[0] in group.layer_names: - self._hs_group_idx = i - break + # Block size must match the indexed buffer, else reads hit the wrong + # slots. Raise (not assert) so the check survives `python -O`. + if self._block_size != self._kv_cache.shape[1]: + raise ValueError( + f"Hidden-states block-size mismatch: derived {self._block_size} " + f"but buffer block size is {self._kv_cache.shape[1]}; read slots " + "would be wrong (likely a hybrid block-size resolution bug)." + ) @staticmethod def _write_tensors( @@ -543,7 +586,7 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): request: "Request", block_ids: tuple[list[int], ...], ) -> tuple[bool, dict[str, Any] | None]: - return self.request_finished(request, block_ids[self._hs_group_idx]) + return self.request_finished(request, block_ids[self._cache_kv_group_id]) @classmethod def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: diff --git a/vllm/model_executor/models/extract_hidden_states.py b/vllm/model_executor/models/extract_hidden_states.py index 8df4823b69..87dba7b5c9 100644 --- a/vllm/model_executor/models/extract_hidden_states.py +++ b/vllm/model_executor/models/extract_hidden_states.py @@ -83,7 +83,10 @@ def basic_cache( kv_cache: torch.Tensor, # shape: [num_blocks, block_size, num_heads, head_size] slot_mapping: torch.Tensor, # shape: [seq_len] ): + # Padding slots are -1; redirect them to the null block (block 0, never + # allocated to a request) so the scatter stays branch-free and sync-free. block_size = kv_cache.shape[1] + slot_mapping = slot_mapping.clamp_min(0) kv_cache[slot_mapping // block_size, slot_mapping % block_size] = to_cache