[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>
Igor Margulis committed
Jun 30, 2026 at 18:19 UTC
d8f483dc30a9a74b8fdf1c6dcb11cff503fe40a3
3 files changed
+187
-15
tests/v1/kv_connector/unit/test_hidden_states_connector.py
new
+126
@@ -0,0 +1,126 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""CPU-only unit tests for ExampleHiddenStatesConnector KV-cache-group logic."""
4
+
5
+from types import SimpleNamespace
6
+
7
+import pytest
8
+import torch
9
+
10
+from vllm.distributed.kv_transfer.kv_connector.v1.example_hidden_states_connector import ( # noqa: E501
11
+ ExampleHiddenStatesConnector,
12
+)
13
+from vllm.v1.core.kv_cache_utils import get_kv_cache_groups
14
+from vllm.v1.kv_cache_interface import (
15
+ FullAttentionSpec,
16
+ HiddenStateCacheSpec,
17
+ KVCacheGroupSpec,
18
+ MLAAttentionSpec,
19
+ SlidingWindowMLASpec,
20
+)
21
+
22
+
23
+def _full(block_size: int) -> FullAttentionSpec:
24
+ return FullAttentionSpec(
25
+ block_size=block_size, num_kv_heads=8, head_size=128, dtype=torch.bfloat16
26
+ )
27
+
28
+
29
+def _hidden(block_size: int) -> HiddenStateCacheSpec:
30
+ return HiddenStateCacheSpec(
31
+ block_size=block_size, num_kv_heads=6, head_size=2048, dtype=torch.bfloat16
32
+ )
33
+
34
+
35
+def _config(*specs):
36
+ """Minimal stand-in exposing only ``kv_cache_groups`` (all the helpers read)."""
37
+ return SimpleNamespace(
38
+ kv_cache_groups=[
39
+ KVCacheGroupSpec(layer_names=[f"layer.{i}"], kv_cache_spec=spec)
40
+ for i, spec in enumerate(specs)
41
+ ]
42
+ )
43
+
44
+
45
+# ---- _find_cache_kv_group_id ------------------------------------------------
46
+
47
+
48
+def test_find_group_id_none_config_returns_zero():
49
+ assert ExampleHiddenStatesConnector._find_cache_kv_group_id(None) == 0
50
+
51
+
52
+def test_find_group_id_single_non_hidden_group_returns_zero():
53
+ # Uniform (dense) model: one group, no HiddenStateCacheSpec -> group 0.
54
+ cfg = _config(_full(16))
55
+ assert ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == 0
56
+
57
+
58
+def test_find_group_id_locates_hidden_group_when_not_first():
59
+ # Hybrid layout: the hidden-states group is not group 0.
60
+ cfg = _config(_full(528), _hidden(22), _full(528))
61
+ assert ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == 1
62
+
63
+
64
+def test_find_group_id_locates_hidden_group_last():
65
+ cfg = _config(_full(528), _full(528), _hidden(22))
66
+ assert ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg) == 2
67
+
68
+
69
+def test_find_group_id_raises_when_no_hidden_group_and_multiple_groups():
70
+ cfg = _config(_full(16), _full(16))
71
+ with pytest.raises(ValueError, match="Could not uniquely identify"):
72
+ ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg)
73
+
74
+
75
+def test_find_group_id_raises_when_multiple_hidden_groups():
76
+ cfg = _config(_hidden(22), _hidden(22))
77
+ with pytest.raises(ValueError, match="Could not uniquely identify"):
78
+ ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg)
79
+
80
+
81
+# ---- _get_cache_block_size --------------------------------------------------
82
+
83
+
84
+def test_get_block_size_reads_hidden_group_spec_not_global():
85
+ # Hidden group keeps block size 22; the global is bumped to 528 for hybrids.
86
+ vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=528))
87
+ cfg = _config(_full(528), _hidden(22))
88
+ block_size = ExampleHiddenStatesConnector._get_cache_block_size(
89
+ vllm_config, cfg, cache_kv_group_id=1
90
+ )
91
+ assert block_size == 22
92
+
93
+
94
+def test_get_block_size_falls_back_to_cache_config_when_no_kv_cache_config():
95
+ vllm_config = SimpleNamespace(cache_config=SimpleNamespace(block_size=16))
96
+ block_size = ExampleHiddenStatesConnector._get_cache_block_size(
97
+ vllm_config, None, cache_kv_group_id=0
98
+ )
99
+ assert block_size == 16
100
+
101
+
102
+# ---- MLA-verifier absorption ------------------------------------------------
103
+
104
+
105
+def test_find_group_id_errors_clearly_when_absorbed_by_mla_swa_verifier():
106
+ # HiddenStateCacheSpec subclasses MLAAttentionSpec, so an MLA + sliding-
107
+ # window MLA verifier absorbs it into the MLA group instead of isolating it.
108
+ dt = torch.bfloat16
109
+ spec = {
110
+ "layers.0.mla": MLAAttentionSpec(
111
+ block_size=64, num_kv_heads=1, head_size=576, dtype=dt
112
+ ),
113
+ "layers.1.swa": SlidingWindowMLASpec(
114
+ block_size=64, num_kv_heads=1, head_size=576, dtype=dt, sliding_window=512
115
+ ),
116
+ "cache_only_layers.61": _hidden(64),
117
+ }
118
+ vllm_config = SimpleNamespace(
119
+ scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False),
120
+ speculative_config=None,
121
+ )
122
+ groups = get_kv_cache_groups(vllm_config, spec)
123
+ assert not any(isinstance(g.kv_cache_spec, HiddenStateCacheSpec) for g in groups)
124
+ cfg = SimpleNamespace(kv_cache_groups=groups)
125
+ with pytest.raises(ValueError, match="MLA verifiers are unsupported"):
126
+ ExampleHiddenStatesConnector._find_cache_kv_group_id(cfg)
vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
+58
-15
@@ -109,6 +109,49 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
109
# Must be False so that drafter kv cache isn't merged with verifier's
110
return False
111
112
+ @classmethod
113
+ def _find_cache_kv_group_id(cls, kv_cache_config: "KVCacheConfig | None") -> int:
114
+ """Index of the KV cache group holding the extracted hidden states.
115
+
116
+ Located by spec type so it resolves on both scheduler and worker side.
117
+ """
118
+ if kv_cache_config is None:
119
+ return 0
120
+
121
+ from vllm.v1.kv_cache_interface import HiddenStateCacheSpec
122
+
123
+ groups = kv_cache_config.kv_cache_groups
124
+ group_ids = [
125
+ gid
126
+ for gid, group in enumerate(groups)
127
+ if isinstance(group.kv_cache_spec, HiddenStateCacheSpec)
128
+ ]
129
+ if len(group_ids) == 1:
130
+ return group_ids[0]
131
+ if not group_ids and len(groups) == 1:
132
+ return 0
133
+ raise ValueError(
134
+ "Could not uniquely identify the extract-hidden-states KV cache "
135
+ f"group among {len(groups)} groups; the hidden-states layer must be "
136
+ "isolated in its own group (MLA verifiers are unsupported)."
137
+ )
138
+
139
+ @staticmethod
140
+ def _get_cache_block_size(
141
+ vllm_config: "VllmConfig",
142
+ kv_cache_config: "KVCacheConfig | None",
143
+ cache_kv_group_id: int,
144
+ ) -> int:
145
+ """Block size of the hidden-states group, read from its own spec.
146
+
147
+ cache_config.block_size is bumped to a common multiple for hybrid
148
+ verifiers; the page-aligned hidden-states group keeps a smaller one.
149
+ """
150
+ if kv_cache_config is None:
151
+ return vllm_config.cache_config.block_size
152
+ cache_group = kv_cache_config.kv_cache_groups[cache_kv_group_id]
153
+ return cache_group.kv_cache_spec.block_size
154
+
155
def __init__(
156
self,
157
vllm_config: "VllmConfig",
@@ -120,7 +163,12 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
163
role=role,
164
kv_cache_config=kv_cache_config,
165
)
123
- self._block_size = vllm_config.cache_config.block_size
166
+ # Read the hidden-states group and its block size from the group spec;
167
+ # cache_config.block_size is bumped (wrong) for hybrid verifiers.
168
+ self._cache_kv_group_id = self._find_cache_kv_group_id(kv_cache_config)
169
+ self._block_size = self._get_cache_block_size(
170
+ vllm_config, kv_cache_config, self._cache_kv_group_id
171
+ )
172
self._storage_path = self._kv_transfer_config.get_from_extra_config(
173
"shared_storage_path", "/tmp"
174
)
@@ -151,13 +199,6 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
199
# Worker-side state (set by register_kv_caches).
200
self._kv_cache: torch.Tensor | None = None
201
154
- # Identify which KV cache group holds the hidden-states layer.
155
- self._hs_group_idx: int = 0
156
- if self._kv_cache_config is not None:
157
- for i, group in enumerate(self._kv_cache_config.kv_cache_groups):
158
- if any("cache_only_layers" in n for n in group.layer_names):
159
- self._hs_group_idx = i
160
- break
202
# Only TP rank 0 writes hidden states to disk; other TP ranks no-op.
203
# Set in register_kv_caches (after distributed init).
204
self._is_tp_rank_zero: bool = True
@@ -267,12 +308,14 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
308
)
309
self._kv_cache = kv_caches[self.cache_layers[0]]
310
270
- # Find the KV cache group index for hidden states
271
- if self._kv_cache_config is not None:
272
- for i, group in enumerate(self._kv_cache_config.kv_cache_groups):
273
- if self.cache_layers[0] in group.layer_names:
274
- self._hs_group_idx = i
275
- break
311
+ # Block size must match the indexed buffer, else reads hit the wrong
312
+ # slots. Raise (not assert) so the check survives `python -O`.
313
+ if self._block_size != self._kv_cache.shape[1]:
314
+ raise ValueError(
315
+ f"Hidden-states block-size mismatch: derived {self._block_size} "
316
+ f"but buffer block size is {self._kv_cache.shape[1]}; read slots "
317
+ "would be wrong (likely a hybrid block-size resolution bug)."
318
+ )
319
320
@staticmethod
321
def _write_tensors(
@@ -543,7 +586,7 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
586
request: "Request",
587
block_ids: tuple[list[int], ...],
588
) -> tuple[bool, dict[str, Any] | None]:
546
- return self.request_finished(request, block_ids[self._hs_group_idx])
589
+ return self.request_finished(request, block_ids[self._cache_kv_group_id])
590
591
@classmethod
592
def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:
vllm/model_executor/models/extract_hidden_states.py
+3
@@ -83,7 +83,10 @@ def basic_cache(
83
kv_cache: torch.Tensor, # shape: [num_blocks, block_size, num_heads, head_size]
84
slot_mapping: torch.Tensor, # shape: [seq_len]
85
):
86
+ # Padding slots are -1; redirect them to the null block (block 0, never
87
+ # allocated to a request) so the scatter stays branch-free and sync-free.
88
block_size = kv_cache.shape[1]
89
+ slot_mapping = slot_mapping.clamp_min(0)
90
kv_cache[slot_mapping // block_size, slot_mapping % block_size] = to_cache
91
92