@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2f821faeae

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 19:00 UTC

[Spec Decode] Support hybrid attention models in extract_hidden_states (#39949)

Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Michael Goin committed May 13, 2026 at 13:45 UTC 2f821faeaec6cd8add4e40ed6a7467e0f7f16de6

14 files changed +225 -115

.buildkite/test_areas/misc.yaml

+13

@@ -129,6 +129,19 @@ steps:

129 - pytest -v -s -m 'cpu_test' v1/kv_connector/unit

130 - pytest -v -s -m 'cpu_test' v1/metrics

131

132 +- label: Extract Hidden States Integration

133 + key: extract-hidden-states-integration

134 + timeout_in_minutes: 20

135 + device: h200_18gb

136 + source_file_dependencies:

137 + - vllm/v1/spec_decode/extract_hidden_states.py

138 + - vllm/model_executor/models/extract_hidden_states.py

139 + - vllm/transformers_utils/configs/extract_hidden_states.py

140 + - tests/v1/kv_connector/extract_hidden_states_integration

141 + commands:

142 + - export VLLM_WORKER_MULTIPROC_METHOD=spawn

143 + - pytest -v -s v1/kv_connector/extract_hidden_states_integration

144 +

145 - label: Regression

146 key: regression

147 timeout_in_minutes: 20

tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py

+59 -1

@@ -83,7 +83,7 @@ def register_predictable_model():

85 def test_extract_hidden_states_with_predictable_dummy_model(

86 - predictable_llama_config_path, tmp_path

86 + predictable_llama_config_path, tmp_path, monkeypatch

87 ):

88 """Comprehensive test using a predictable dummy model with synthetic weights.

@@ -94,6 +94,12 @@ def test_extract_hidden_states_with_predictable_dummy_model(

94 3. Layer ordering is preserved correctly (non-sequential layer IDs)

95 4. Multiple prompts of different lengths produce consistent layer values

96 """

97 + # Force fork so the engine worker inherits the autouse fixture's

98 + # ModelRegistry.register_model("PredictableLlamaForCausalLM", ...).

99 + # Spawn (the CI default) starts a fresh Python process that wouldn't

100 + # see the registration.

101 + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "fork")

102 +

103 # Test with non-sequential layer ordering to verify correct association

104 layer_ids = [5, 2, 10]

105 num_layers = len(layer_ids)

@@ -153,3 +159,55 @@ def test_extract_hidden_states_with_predictable_dummy_model(

159 f"but got mean={layer_hidden.mean():.3f}, "

160 f"min={layer_hidden.min():.3f}, max={layer_hidden.max():.3f}"

161 )

162 +

163 +

164 +def test_extract_hidden_states_qwen35_hybrid_smoke(tmp_path):

165 + """Smoke test for Qwen3.5 hybrid (mamba + full-attention) models.

166 + Uses load_format="dummy" to just check shape/plumbing.

167 + """

168 + layer_ids = [5, 11, 17]

169 + hidden_size = 1024 # Qwen/Qwen3.5-0.8B hidden_size

170 +

171 + llm = LLM(

172 + model="Qwen/Qwen3.5-0.8B",

173 + speculative_config={

174 + "method": "extract_hidden_states",

175 + "num_speculative_tokens": 1,

176 + "draft_model_config": {

177 + "hf_config": {"eagle_aux_hidden_state_layer_ids": layer_ids}

178 + },

179 + },

180 + kv_transfer_config={

181 + "kv_connector": "ExampleHiddenStatesConnector",

182 + "kv_role": "kv_producer",

183 + "kv_connector_extra_config": {"shared_storage_path": str(tmp_path)},

184 + },

185 + max_model_len=256,

186 + enforce_eager=True,

187 + gpu_memory_utilization=0.4,

188 + load_format="dummy",

189 + )

190 +

191 + prompts = ["Hello world", "Test prompt with several tokens"]

192 + sampling_params = SamplingParams(max_tokens=1, temperature=0.0)

193 + outputs = llm.generate(prompts, sampling_params)

194 + del llm

195 + gc.collect()

196 +

197 + assert len(outputs) == len(prompts)

198 + for output in outputs:

199 + assert output.kv_transfer_params is not None

200 + hidden_states_path = output.kv_transfer_params.get("hidden_states_path")

201 + assert hidden_states_path is not None

202 + assert os.path.exists(hidden_states_path)

203 +

204 + with safe_open(hidden_states_path, "pt") as f:

205 + token_ids = f.get_tensor("token_ids")

206 + hidden_states = f.get_tensor("hidden_states")

207 +

208 + assert torch.equal(token_ids, torch.tensor(output.prompt_token_ids))

209 + assert hidden_states.shape == (

210 + len(output.prompt_token_ids),

211 + len(layer_ids),

212 + hidden_size,

213 + )

tests/v1/kv_connector/unit/test_decode_bench_connector.py

+2 -5

@@ -11,7 +11,6 @@ import pytest

11 import torch

13 from vllm import SamplingParams

14 -from vllm.config import KVTransferConfig

14 from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole

16 # ruff: noqa: E501

@@ -44,11 +43,9 @@ class DecodeBenchTestRunner:

44 # Create vllm config with DecodeBenchConnector

45 vllm_config = create_vllm_config(

47 - block_size=block_size, max_num_batched_tokens=1000

48 - )

49 - vllm_config.kv_transfer_config = KVTransferConfig(

46 + block_size=block_size,

47 + max_num_batched_tokens=1000,

48 kv_connector="DecodeBenchConnector",

51 - kv_role="kv_both",

49 )

51 self.vllm_config = vllm_config

tests/v1/kv_connector/unit/test_kv_connector_lifecycle.py

+5 -4

@@ -35,10 +35,11 @@ def _make_empty_scheduler_output():

37 def test_kv_connector_mixin_clears_metadata():

38 - vllm_config = create_vllm_config()

39 - vllm_config.kv_transfer_config.kv_connector = "TestExampleConnector"

40 - vllm_config.kv_transfer_config.kv_role = "kv_both"

41 - vllm_config.kv_transfer_config.kv_connector_extra_config["name"] = "unit"

38 + vllm_config = create_vllm_config(

39 + kv_connector="TestExampleConnector",

40 + kv_role="kv_both",

41 + kv_connector_extra_config={"name": "unit"},

42 + )

44 kv_cache_config = KVCacheConfig(

45 num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]

tests/v1/kv_connector/unit/test_multi_connector.py

+24 -33

@@ -130,17 +130,14 @@ KVConnectorFactory.register_connector(

130 @pytest.fixture

131 def mc() -> MultiConnector:

132 """MultiConnector using two mocked connectors"""

133 - vllm_config = create_vllm_config()

134 -

133 mock_connector_config = {

134 "kv_connector": "MockConnector",

135 "kv_role": "kv_both",

136 "kv_connector_module_path": "tests.v1.kv_connector.unit.test_multi_connector",

137 }

138

141 - vllm_config.kv_transfer_config = KVTransferConfig(

139 + vllm_config = create_vllm_config(

140 kv_connector="MultiConnector",

143 - kv_role="kv_both",

141 kv_connector_extra_config={

142 "connectors": [mock_connector_config, mock_connector_config],

143 },

@@ -405,39 +402,35 @@ def test_multi_connector_handle_preemptions_integration():

402

403 try:

404 # Configure MultiConnector with two TestExampleConnectors

408 - kv_transfer_config = KVTransferConfig(

409 - kv_connector="MultiConnector",

410 - kv_role="kv_both",

411 - kv_connector_extra_config={

412 - "connectors": [

413 - {

414 - "kv_connector": "TestExampleConnector",

415 - "kv_role": "kv_both",

416 - "kv_connector_extra_config": {

417 - "shared_storage_path": str(storage_path / "s1"),

418 - "name": "preempt1",

419 - },

420 - "kv_connector_module_path": "tests.v1.kv_connector.unit.utils",

405 + connectors_extra_config = {

406 + "connectors": [

407 + {

408 + "kv_connector": "TestExampleConnector",

409 + "kv_role": "kv_both",

410 + "kv_connector_extra_config": {

411 + "shared_storage_path": str(storage_path / "s1"),

412 + "name": "preempt1",

413 },

422 - {

423 - "kv_connector": "TestExampleConnector",

424 - "kv_role": "kv_both",

425 - "kv_connector_extra_config": {

426 - "shared_storage_path": str(storage_path / "s2"),

427 - "name": "preempt2",

428 - },

429 - "kv_connector_module_path": "tests.v1.kv_connector.unit.utils",

414 + "kv_connector_module_path": "tests.v1.kv_connector.unit.utils",

415 + },

416 + {

417 + "kv_connector": "TestExampleConnector",

418 + "kv_role": "kv_both",

419 + "kv_connector_extra_config": {

420 + "shared_storage_path": str(storage_path / "s2"),

421 + "name": "preempt2",

422 },

431 - ]

432 - },

433 - )

423 + "kv_connector_module_path": "tests.v1.kv_connector.unit.utils",

424 + },

425 + ]

426 + }

427

428 vllm_config = create_vllm_config(

429 block_size=16,

430 max_num_batched_tokens=100,

438 - kv_connector_extra_config=kv_transfer_config.kv_connector_extra_config,

431 + kv_connector="MultiConnector",

432 + kv_connector_extra_config=connectors_extra_config,

433 )

440 - vllm_config.kv_transfer_config = kv_transfer_config

434

435 # Create scheduler - this initializes the MultiConnector with SCHEDULER role

436 scheduler = create_scheduler(vllm_config, num_blocks=10)

@@ -973,7 +966,6 @@ def test_multi_connector_worker_metadata(mc):

966

967 def _make_multi_connector(connector_names: list[str]) -> MultiConnector:

968 """Build a MultiConnector wrapping the given registered connectors."""

976 - vllm_config = create_vllm_config()

969 connectors = [

970 {

971 "kv_connector": name,

@@ -982,9 +974,8 @@ def _make_multi_connector(connector_names: list[str]) -> MultiConnector:

974 }

975 for name in connector_names

976 ]

985 - vllm_config.kv_transfer_config = KVTransferConfig(

977 + vllm_config = create_vllm_config(

978 kv_connector="MultiConnector",

987 - kv_role="kv_both",

979 kv_connector_extra_config={"connectors": connectors},

980 )

981 kv_cache_config = KVCacheConfig(

tests/v1/kv_connector/unit/utils.py

@@ -102,6 +102,7 @@ def create_vllm_config(

102 attention_backend: str | None = None,

103 kv_load_failure_policy: Literal["recompute", "fail"] = "fail",

104 kv_connector: str = "NixlConnector",

105 + kv_connector_module_path: str | None = None,

106 kv_role: str = "kv_both",

107 disable_hybrid_kv_cache_manager: bool | None = None,

108 ) -> VllmConfig:

@@ -130,6 +131,7 @@ def create_vllm_config(

131 )

132 kv_transfer_config = KVTransferConfig(

133 kv_connector=kv_connector,

134 + kv_connector_module_path=kv_connector_module_path,

135 kv_role=kv_role,

136 enable_permute_local_kv=enable_permute_local_kv,

137 kv_connector_extra_config=kv_connector_extra_config or {},

vllm/config/vllm.py

+39 -14

@@ -1316,6 +1316,10 @@ class VllmConfig:

1316 "the `reasoning_start_str` and `reasoning_end_str`."

1317 )

1318

1319 + # Resolve kv_offloading-derived connector name into kv_transfer_config

1320 + # before the HMA check below, which inspects the connector class.

1321 + self._post_init_kv_transfer_config()

1322 +

1323 # Hybrid KV cache manager (HMA) runtime rules:

1324 # - Explicit enable (--no-disable-kv-cache-manager): error if runtime

1325 # disables it

@@ -1353,18 +1357,42 @@ class VllmConfig:

1357 if self.scheduler_config.disable_hybrid_kv_cache_manager is None:

1358 # Default to disable HMA, but only if the user didn't express a preference.

1359 if self.kv_transfer_config is not None:

1356 - # NOTE(Kuntai): turn HMA off for connector unless specifically enabled.

1357 - need_disable_hybrid_kv_cache_manager = True

1358 - logger.warning(

1359 - "Turning off hybrid kv cache manager because "

1360 - "`--kv-transfer-config` is set. This will reduce the "

1361 - "performance of vLLM on LLMs with sliding window attention "

1362 - "or Mamba attention. If you are a developer of kv connector"

1363 - ", please consider supporting hybrid kv cache manager for "

1364 - "your connector by making sure your connector is a subclass"

1365 - " of `SupportsHMA` defined in kv_connector/v1/base.py and"

1366 - " use --no-disable-hybrid-kv-cache-manager to start vLLM."

1360 + from vllm.config.kv_transfer import KVTransferConfig

1361 + from vllm.distributed.kv_transfer.kv_connector.factory import (

1362 + KVConnectorFactory,

1363 + )

1364 + from vllm.distributed.kv_transfer.kv_connector.v1.base import (

1365 + supports_hma,

1366 + )

1367 +

1368 + connector_cls = KVConnectorFactory.get_connector_class(

1369 + self.kv_transfer_config

1370 )

1371 + all_support_hma = supports_hma(connector_cls)

1372 + # MultiConnector subclasses SupportsHMA; only effectively

1373 + # supports HMA when every sub-connector does.

1374 + if all_support_hma and connector_cls.__name__ == "MultiConnector":

1375 + sub_ktcs = self.kv_transfer_config.kv_connector_extra_config.get(

1376 + "connectors", []

1377 + )

1378 + all_support_hma = all(

1379 + supports_hma(

1380 + KVConnectorFactory.get_connector_class(

1381 + KVTransferConfig(**sub)

1382 + )

1383 + )

1384 + for sub in sub_ktcs

1385 + )

1386 + if not all_support_hma:

1387 + need_disable_hybrid_kv_cache_manager = True

1388 + logger.warning(

1389 + "Turning off hybrid kv cache manager because "

1390 + "connector %s does not subclass `SupportsHMA`. "

1391 + "This will reduce performance on models with "

1392 + "sliding window or Mamba attention. See "

1393 + "kv_connector/v1/base.py for details.",

1394 + connector_cls.__name__,

1395 + )

1396 self.scheduler_config.disable_hybrid_kv_cache_manager = (

1397 need_disable_hybrid_kv_cache_manager

1398 )

@@ -1406,10 +1434,7 @@ class VllmConfig:

1434 if "-quant_fp8" not in custom_ops:

1435 custom_ops.append("+quant_fp8")

1436

1409 - # Handle the KV connector configs

1410 - self._post_init_kv_transfer_config()

1437 self._verify_kv_transfer_compat()

1412 -

1438 # Log the custom passes that are enabled

1439 self.compilation_config.pass_config.log_enabled_passes()

1440

vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py

+22 -35

@@ -12,6 +12,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (

12 KVConnectorBase_V1,

13 KVConnectorMetadata,

14 KVConnectorRole,

15 + SupportsHMA,

16 )

17 from vllm.logger import init_logger

18 from vllm.v1.attention.backend import AttentionMetadata

@@ -30,13 +31,9 @@ def extract_from_kv_cache(

31 slot_mapping: torch.Tensor,

32 num_tokens: int,

33 ) -> torch.Tensor:

33 - """Extract data from KV cache

34 - Assume the shape of the kv_cache is (num_pages, page_size, num_heads, head_size)

35 - """

36 -

37 - padded_kv = kv_cache.flatten(0, 1)[slot_mapping]

38 - # shape: [len(slot_mapping), num_heads, head_size]

39 - return padded_kv[:num_tokens] # shape: [num_tokens, num_heads, head_size]

34 + """Extract data from KV cache."""

35 + block_size = kv_cache.shape[1]

36 + return kv_cache[slot_mapping // block_size, slot_mapping % block_size][:num_tokens]

39 @dataclass

@@ -47,8 +44,6 @@ class ReqMeta:

44 filename: str

45 # Request tokens

46 token_ids: torch.Tensor

50 - # Slot mappings, should have the same length as token_ids

51 - slot_mapping: torch.Tensor

47 # Whether this request is a new request or partially computed already

48 new_req: bool

@@ -57,24 +52,12 @@ class ReqMeta:

52 req_id: str,

53 filename: str,

54 token_ids: list[int],

60 - block_ids: list[int],

61 - block_size: int,

55 new_req: bool,

56 ) -> "ReqMeta":

64 - token_ids_tensor = torch.tensor(token_ids)

65 - block_ids_tensor = torch.tensor(block_ids)

66 - num_blocks = block_ids_tensor.shape[0]

67 - block_offsets = torch.arange(0, block_size)

68 - slot_mapping = (

69 - block_offsets.reshape((1, block_size))

70 - + block_ids_tensor.reshape((num_blocks, 1)) * block_size

71 - )

72 - slot_mapping = slot_mapping.flatten()

57 return ReqMeta(

58 req_id=req_id,

59 filename=filename,

76 - token_ids=token_ids_tensor,

77 - slot_mapping=slot_mapping,

60 + token_ids=torch.tensor(token_ids),

61 new_req=new_req,

62 )

@@ -88,18 +71,12 @@ class ExampleHiddenStatesConnectorMetadata(KVConnectorMetadata):

71 req_id: str,

72 filename: str,

73 token_ids: list[int],

91 - block_ids: list[int],

92 - block_size: int,

74 new_req: bool = True,

75 ) -> None:

95 - self.requests.append(

96 - ReqMeta.make_meta(

97 - req_id, filename, token_ids, block_ids, block_size, new_req

98 - )

99 - )

76 + self.requests.append(ReqMeta.make_meta(req_id, filename, token_ids, new_req))

102 -class ExampleHiddenStatesConnector(KVConnectorBase_V1):

79 +class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):

80 """

81 Simple debug implementation of a HiddenStatesConnector.

@@ -206,9 +183,16 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1):

183 assert isinstance(connector_metadata, ExampleHiddenStatesConnectorMetadata)

184

185 os.makedirs(self._storage_path, exist_ok=True)

186 +

187 + slot_mapping = attn_metadata.slot_mapping

188 + offset = 0

189 for request in connector_metadata.requests:

190 + num_tokens = request.token_ids.shape[0]

191 + req_slot_mapping = slot_mapping[offset : offset + num_tokens]

192 + offset += num_tokens

193 +

194 hidden_states = extract_from_kv_cache(

211 - kv_layer, request.slot_mapping, request.token_ids.shape[0]

195 + kv_layer, req_slot_mapping, num_tokens

196 )

197 tensors = {

198 "hidden_states": hidden_states.detach().cpu(),

@@ -269,8 +253,6 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1):

253 new_req.req_id,

254 filename=filename,

255 token_ids=token_ids,

272 - block_ids=new_req.block_ids[0],

273 - block_size=self._block_size,

256 )

257 self._request_filenames[new_req.req_id] = filename

258 self._active_requests[new_req.req_id] = new_req

@@ -298,8 +280,6 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1):

280 req_id=req_id,

281 filename=filename,

282 token_ids=cached_req.prompt_token_ids or [],

301 - block_ids=req_block_ids,

302 - block_size=self._block_size,

283 new_req=False,

284 )

285

@@ -331,6 +311,13 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1):

311

312 return False, {"hidden_states_path": req_filename}

313

314 + def request_finished_all_groups(

315 + self,

316 + request: "Request",

317 + block_ids: tuple[list[int], ...],

318 + ) -> tuple[bool, dict[str, Any] | None]:

319 + return self.request_finished(request, block_ids[0])

320 +

321 @classmethod

322 def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:

323 """

vllm/model_executor/models/extract_hidden_states.py

+8 -11

@@ -34,8 +34,8 @@ from vllm.v1.attention.backend import (

34 )

35 from vllm.v1.kv_cache_interface import (

36 AttentionSpec,

37 + HiddenStateCacheSpec,

38 KVCacheSpec,

38 - MLAAttentionSpec,

39 )

41 ########## Custom Ops ########

@@ -79,13 +79,12 @@ def dummy_attention(layer_name, _placeholder):

81 def basic_cache(

82 - to_cache: torch.Tensor, # shape: [num_blocks, block_size, num_heads, head_size]

83 - kv_cache: torch.Tensor, # shape: [seq_len, num_heads, head_size]

82 + to_cache: torch.Tensor, # shape: [seq_len, num_heads, head_size]

83 + kv_cache: torch.Tensor, # shape: [num_blocks, block_size, num_heads, head_size]

84 slot_mapping: torch.Tensor, # shape: [seq_len]

85 ):

86 - num_blocks, block_size, num_heads, head_size = kv_cache.shape

87 - token_kv_cache = kv_cache.view(num_blocks * block_size, num_heads, head_size)

88 - token_kv_cache[slot_mapping] = to_cache

86 + block_size = kv_cache.shape[1]

87 + kv_cache[slot_mapping // block_size, slot_mapping % block_size] = to_cache

90 ######### CacheOnlyAttentionBackend ########

@@ -322,11 +321,9 @@ class CacheOnlyAttentionLayer(nn.Module, AttentionLayerBase):

321 return self.attn_backend

322

323 def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:

325 - # Note: we use MLAAttentionSpec here to because it will

326 - # produce page sizes of (block_size * num_kv_heads * head_size * dtype_size)

327 - # whereas FullAttentionSpec will add an additional factor of 2

328 - return MLAAttentionSpec(

329 - block_size=self.block_size,

324 + # Re-read block_size: hybrid models may bump it after __init__.

325 + return HiddenStateCacheSpec(

326 + block_size=vllm_config.cache_config.block_size,

327 num_kv_heads=self.num_heads,

328 head_size=self.head_size,

329 dtype=self.kv_cache_torch_dtype,

vllm/v1/core/kv_cache_utils.py

+26 -6

@@ -18,9 +18,11 @@ from vllm.logger import init_logger

18 from vllm.utils.hashing import sha256_cbor, xxhash_cbor

19 from vllm.utils.math_utils import cdiv, round_up

20 from vllm.utils.mem_utils import format_gib

21 +from vllm.utils.torch_utils import get_dtype_size

22 from vllm.v1.kv_cache_interface import (

23 ChunkedLocalAttentionSpec,

24 FullAttentionSpec,

25 + HiddenStateCacheSpec,

26 KVCacheConfig,

27 KVCacheGroupSpec,

28 KVCacheSpec,

@@ -1650,15 +1652,33 @@ def get_kv_cache_groups(

1652 _annotate_eagle_groups_deepseek_v4(vllm_config, kv_cache_spec, kv_cache_groups)

1653 return kv_cache_groups

1654

1655 + # Pull HiddenStateCacheSpec layers out before the general multi-group

1656 + # path so they don't affect page-size unification or grouping.

1657 + hidden_specs = {

1658 + k: v for k, v in kv_cache_spec.items() if isinstance(v, HiddenStateCacheSpec)

1659 + }

1660 + filtered_spec = {

1661 + k: v

1662 + for k, v in kv_cache_spec.items()

1663 + if not isinstance(v, HiddenStateCacheSpec)

1664 + }

1665 +

1666 # As KVCacheManager can only allocate memory of one size, we need to unify

1667 # the page size of the layers. For cases cannot be unified, this function

1668 # will raise an error.

1656 - kv_cache_spec = unify_kv_cache_spec_page_size(kv_cache_spec)

1657 - # Model contains multiple attention types, but KV cache of all layers

1658 - # have the same physical memory per block per layer. Split the layers

1659 - # into groups with the same number of layers, and thus same total page

1660 - # size.

1661 - return _get_kv_cache_groups_uniform_page_size(kv_cache_spec)

1669 + filtered_spec = unify_kv_cache_spec_page_size(filtered_spec)

1670 + groups = _get_kv_cache_groups_uniform_page_size(filtered_spec)

1671 +

1672 + # Add hidden-state layers back with page aligned to the common page.

1673 + if hidden_specs:

1674 + common_page = get_uniform_page_size([g.kv_cache_spec for g in groups])

1675 + for name, spec in hidden_specs.items():

1676 + per_token = spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype)

1677 + new_bs = max(common_page // per_token, 1)

1678 + aligned = replace(spec, block_size=new_bs, page_size_padded=common_page)

1679 + groups.append(KVCacheGroupSpec([name], aligned))

1680 +

1681 + return groups

1682

1683

1684 def generate_scheduler_kv_cache_config(

vllm/v1/core/single_type_kv_cache_manager.py

@@ -16,6 +16,7 @@ from vllm.v1.kv_cache_interface import (

16 ChunkedLocalAttentionSpec,

17 CrossAttentionSpec,

18 FullAttentionSpec,

19 + HiddenStateCacheSpec,

20 KVCacheSpec,

21 MambaSpec,

22 MLAAttentionSpec,

@@ -1143,6 +1144,7 @@ spec_manager_map: dict[type[KVCacheSpec], type[SingleTypeKVCacheManager]] = {

1144 FullAttentionSpec: FullAttentionManager,

1145 TQFullAttentionSpec: FullAttentionManager,

1146 MLAAttentionSpec: FullAttentionManager,

1147 + HiddenStateCacheSpec: FullAttentionManager,

1148 SlidingWindowSpec: SlidingWindowManager,

1149 SlidingWindowMLASpec: SlidingWindowManager,

1150 ChunkedLocalAttentionSpec: ChunkedLocalAttentionManager,

vllm/v1/kv_cache_interface.py

@@ -396,6 +396,13 @@ class MLAAttentionSpec(FullAttentionSpec):

396 )

397

398

399 +@dataclass(frozen=True, kw_only=True)

400 +class HiddenStateCacheSpec(MLAAttentionSpec):

401 + """Marker for hidden-state cache layers used by extract_hidden_states."""

402 +

403 + pass

404 +

405 +

406 @dataclass(frozen=True, kw_only=True)

407 class ChunkedLocalAttentionSpec(AttentionSpec):

408 attention_chunk_size: int

vllm/v1/spec_decode/extract_hidden_states.py

+9 -5

@@ -42,6 +42,7 @@ class ExtractHiddenStatesProposer:

42 self.model: nn.Module | None = None

43 self.attn_layer_names: list[str] = []

44 self.attn_metadata_builder: AttentionMetadataBuilder | None = None

45 + self.kv_cache_gid: int = -1

47 # Maximum number of tokens for buffers

48 max_batch_size = vllm_config.scheduler_config.max_num_seqs

@@ -374,9 +375,12 @@ class ExtractHiddenStatesProposer:

375 )

376

377 def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:

377 - """Validate all drafting layers belong to the same KV cache group.

378 -

379 - With exactly one attention layer (asserted in load_model), this is

380 - trivially satisfied.

381 - """

378 + """Validate all drafting layers belong to the same KV cache group

379 + and record the group index for common_attn_metadata selection."""

380 assert len(self.attn_layer_names) == 1

381 + layer = self.attn_layer_names[0]

382 + for gid, group in enumerate(kv_cache_config.kv_cache_groups):

383 + if layer in group.layer_names:

384 + self.kv_cache_gid = gid

385 + return

386 + raise ValueError(f"Cache-only layer {layer!r} not in any KV cache group")

vllm/v1/worker/gpu_model_runner.py

+7 -1

@@ -2347,7 +2347,13 @@ class GPUModelRunner(

2347

2348 if self.speculative_config and spec_decode_common_attn_metadata is None:

2349 if isinstance(

2350 - self.drafter, (EagleProposer, DFlashProposer, Gemma4Proposer)

2350 + self.drafter,

2351 + (

2352 + EagleProposer,

2353 + DFlashProposer,

2354 + Gemma4Proposer,

2355 + ExtractHiddenStatesProposer,

2356 + ),

2357 ):

2358 if self.drafter.kv_cache_gid == kv_cache_gid:

2359 spec_decode_common_attn_metadata = cm