[Spec Decode] Fix hidden-state extraction block size for hybrid verifiers (#46301)

Signed-off-by: Igor Margulis <igor.margulis@intel.com>
Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: mgoin <mgoin64@gmail.com>

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)

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:

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