[Feature][Model Runner V2] Support extract_hidden_states speculation (#49811)
Signed-off-by: Zupeng Wang <71580390+zupengwang@users.noreply.github.com> Signed-off-by: Misha Goin <mgoin64@gmail.com> Co-authored-by: OpenAI <support@openai.com> Co-authored-by: Misha Goin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Zupeng Wang committed
Aug 20, 2026 at 23:13 UTC
01af92e175407231b1433b0aef01a1b9c983d955
6 files changed
+314
-2
tests/test_config.py
+15
@@ -5,6 +5,7 @@ import logging
5
import os
6
from dataclasses import MISSING, Field, asdict, dataclass, field
7
from types import SimpleNamespace
8
+from typing import cast
9
from unittest.mock import patch
10
11
import pydantic
@@ -263,6 +264,20 @@ def test_dsa_models_select_matching_mtp(model_type, expected_architecture):
264
assert hf_config.architectures == [expected_architecture]
265
266
267
+def test_v2_model_runner_supports_extract_hidden_states():
268
+ config = VllmConfig()
269
+ config.speculative_config = cast(
270
+ SpeculativeConfig,
271
+ SimpleNamespace(
272
+ method="extract_hidden_states",
273
+ parallel_drafting=False,
274
+ enable_adaptive_verification=False,
275
+ ),
276
+ )
277
+
278
+ assert config._get_v2_model_runner_unsupported_features() == []
279
+
280
+
281
@pytest.mark.parametrize(
282
("use_v2_model_runner", "expected_capture_sizes"),
283
[
tests/v1/worker/test_gpu_extract_hidden_states_speculator.py
new
+132
@@ -0,0 +1,132 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from contextlib import nullcontext
4
+from types import SimpleNamespace
5
+from typing import Any, cast
6
+
7
+import pytest
8
+import torch
9
+
10
+from vllm.v1.worker.gpu.spec_decode import extract_hidden_states as spec_module
11
+from vllm.v1.worker.gpu.spec_decode import init_speculator
12
+from vllm.v1.worker.gpu.spec_decode.extract_hidden_states import (
13
+ ExtractHiddenStatesSpeculator,
14
+)
15
+
16
+
17
+class _RecordingModel(torch.nn.Module):
18
+ def forward(self, *, hidden_states: torch.Tensor) -> None:
19
+ self.hidden_states = hidden_states.clone()
20
+
21
+
22
+def test_init_requires_greedy_draft_sampling():
23
+ vllm_config = cast(
24
+ Any,
25
+ SimpleNamespace(
26
+ speculative_config=SimpleNamespace(draft_sample_method="probabilistic")
27
+ ),
28
+ )
29
+
30
+ with pytest.raises(ValueError, match="only supports draft_sample_method='greedy'"):
31
+ ExtractHiddenStatesSpeculator(vllm_config, torch.device("cpu"))
32
+
33
+
34
+def test_init_speculator_dispatches_extract_hidden_states(monkeypatch):
35
+ vllm_config = cast(
36
+ Any,
37
+ SimpleNamespace(
38
+ speculative_config=SimpleNamespace(method="extract_hidden_states")
39
+ ),
40
+ )
41
+ device = torch.device("cpu")
42
+
43
+ def fake_speculator(config, target_device):
44
+ return config, target_device
45
+
46
+ monkeypatch.setattr(spec_module, "ExtractHiddenStatesSpeculator", fake_speculator)
47
+
48
+ assert init_speculator(vllm_config, device) == (vllm_config, device)
49
+
50
+
51
+def test_propose_caches_hidden_states_and_returns_sampled_tokens(monkeypatch):
52
+ contexts = []
53
+
54
+ def fake_set_forward_context(*args, **kwargs):
55
+ contexts.append((args, kwargs))
56
+ return nullcontext()
57
+
58
+ monkeypatch.setattr(spec_module, "set_forward_context", fake_set_forward_context)
59
+
60
+ layer_name = "cache_only_layers.2"
61
+ speculator = object.__new__(ExtractHiddenStatesSpeculator)
62
+ speculator.vllm_config = cast(Any, SimpleNamespace())
63
+ speculator.num_hidden_states = 2
64
+ speculator.hidden_states = torch.zeros(4, 2, 3)
65
+ speculator.draft_attn_layer_names = {layer_name}
66
+ speculator.model = _RecordingModel()
67
+
68
+ input_batch = cast(
69
+ Any,
70
+ SimpleNamespace(
71
+ idx_mapping=torch.tensor([2, 0], dtype=torch.int32),
72
+ is_padding=torch.zeros(4, dtype=torch.bool),
73
+ ),
74
+ )
75
+ aux_hidden_states = [
76
+ torch.full((4, 3), 1.0),
77
+ torch.full((4, 3), 2.0),
78
+ ]
79
+ attn_metadata = {layer_name: object(), "target_layer": object()}
80
+ slot_mappings = {
81
+ layer_name: torch.arange(4),
82
+ "target_layer": torch.arange(4),
83
+ }
84
+ last_sampled = torch.tensor([[10], [11], [12]], dtype=torch.int64)
85
+
86
+ draft_tokens = speculator.propose(
87
+ input_batch=input_batch,
88
+ attn_metadata=attn_metadata,
89
+ slot_mappings=slot_mappings,
90
+ last_hidden_states=torch.empty(0),
91
+ aux_hidden_states=aux_hidden_states,
92
+ num_sampled=torch.empty(0),
93
+ num_rejected=torch.empty(0),
94
+ last_sampled=last_sampled,
95
+ next_prefill_tokens=torch.empty(0),
96
+ temperature=torch.empty(0),
97
+ seeds=torch.empty(0),
98
+ )
99
+
100
+ expected_hidden_states = torch.stack(aux_hidden_states, dim=1)
101
+ assert torch.equal(speculator.model.hidden_states, expected_hidden_states)
102
+ assert torch.equal(draft_tokens, torch.tensor([[12], [10]]))
103
+
104
+ assert len(contexts) == 1
105
+ args, kwargs = contexts[0]
106
+ assert args[0] == {layer_name: attn_metadata[layer_name]}
107
+ assert kwargs["num_tokens"] == 4
108
+ assert set(kwargs["slot_mapping"]) == {layer_name}
109
+ assert torch.equal(kwargs["slot_mapping"][layer_name], slot_mappings[layer_name])
110
+
111
+
112
+def test_propose_requires_aux_hidden_states():
113
+ speculator = object.__new__(ExtractHiddenStatesSpeculator)
114
+ speculator.num_hidden_states = 2
115
+ input_batch = cast(
116
+ Any, SimpleNamespace(idx_mapping=torch.tensor([0], dtype=torch.int32))
117
+ )
118
+
119
+ with pytest.raises(ValueError, match="aux_hidden_states are required"):
120
+ speculator.propose(
121
+ input_batch=input_batch,
122
+ attn_metadata={},
123
+ slot_mappings={},
124
+ last_hidden_states=torch.empty(0),
125
+ aux_hidden_states=None,
126
+ num_sampled=torch.empty(0),
127
+ num_rejected=torch.empty(0),
128
+ last_sampled=torch.tensor([[10]]),
129
+ next_prefill_tokens=torch.empty(0),
130
+ temperature=torch.empty(0),
131
+ seeds=torch.empty(0),
132
+ )
vllm/config/vllm.py
+1
@@ -2429,6 +2429,7 @@ class VllmConfig:
2429
"mtp",
2430
"dflash",
2431
"dspark",
2432
+ "extract_hidden_states",
2433
):
2434
unsupported.append(f"speculative method '{speculative_config.method}'")
2435
vllm/v1/worker/gpu/model_runner.py
+6
-1
@@ -247,7 +247,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
247
if self.is_last_pp_rank:
248
self.speculator = init_speculator(self.vllm_config, self.device)
249
250
- if self.speculative_config.method in ("eagle3", "dflash", "dspark"):
250
+ if self.speculative_config.method in (
251
+ "eagle3",
252
+ "dflash",
253
+ "dspark",
254
+ "extract_hidden_states",
255
+ ):
256
# Drafting may require auxiliary hidden states from target model outputs
257
self.use_aux_hidden_state_outputs = True
258
if self.use_pp:
vllm/v1/worker/gpu/spec_decode/__init__.py
+7
-1
@@ -8,7 +8,13 @@ from vllm.config import VllmConfig
8
def init_speculator(vllm_config: VllmConfig, device: torch.device):
9
speculative_config = vllm_config.speculative_config
10
assert speculative_config is not None
11
- if speculative_config.method == "dflash":
11
+ if speculative_config.method == "extract_hidden_states":
12
+ from vllm.v1.worker.gpu.spec_decode.extract_hidden_states import (
13
+ ExtractHiddenStatesSpeculator,
14
+ )
15
+
16
+ return ExtractHiddenStatesSpeculator(vllm_config, device)
17
+ elif speculative_config.method == "dflash":
18
from vllm.v1.worker.gpu.spec_decode.dflash.speculator import (
19
DFlashSpeculator,
20
)
vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py
new
+153
@@ -0,0 +1,153 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import Any
4
+
5
+import torch
6
+import torch.nn as nn
7
+
8
+from vllm.compilation.backends import set_model_tag
9
+from vllm.config import VllmConfig
10
+from vllm.config.compilation import CUDAGraphMode
11
+from vllm.forward_context import set_forward_context
12
+from vllm.model_executor.model_loader import get_model
13
+from vllm.v1.worker.gpu.input_batch import InputBatch
14
+from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
15
+
16
+
17
+class ExtractHiddenStatesSpeculator(DraftModelSpeculator):
18
+ """Cache target hidden states while returning always-accepted draft tokens."""
19
+
20
+ def __init__(self, vllm_config: VllmConfig, device: torch.device):
21
+ assert vllm_config.speculative_config is not None
22
+ if vllm_config.speculative_config.draft_sample_method != "greedy":
23
+ raise ValueError(
24
+ "extract_hidden_states only supports draft_sample_method='greedy'"
25
+ )
26
+ super().__init__(vllm_config, device)
27
+
28
+ if self.num_speculative_steps != 1:
29
+ raise ValueError(
30
+ "extract_hidden_states requires num_speculative_tokens to be 1"
31
+ )
32
+ if self.speculative_config.disable_padded_drafter_batch:
33
+ raise ValueError(
34
+ "disable_padded_drafter_batch is not supported with "
35
+ "extract_hidden_states method"
36
+ )
37
+
38
+ self.supports_mm_inputs = False
39
+ layer_ids = getattr(
40
+ self.draft_model_config.hf_config,
41
+ "eagle_aux_hidden_state_layer_ids",
42
+ None,
43
+ )
44
+ if not layer_ids:
45
+ raise ValueError(
46
+ "eagle_aux_hidden_state_layer_ids must be set in the draft "
47
+ "model config for extract_hidden_states method"
48
+ )
49
+
50
+ self.num_hidden_states = len(layer_ids)
51
+ assert isinstance(self.dtype, torch.dtype)
52
+ self.hidden_states = torch.zeros(
53
+ self.max_num_tokens,
54
+ self.num_hidden_states,
55
+ self.vllm_config.model_config.get_hidden_size(),
56
+ dtype=self.dtype,
57
+ device=device,
58
+ )
59
+
60
+ def load_draft_model(
61
+ self,
62
+ target_model: nn.Module,
63
+ target_attn_layer_names: set[str],
64
+ ) -> nn.Module:
65
+ del target_model, target_attn_layer_names
66
+ with set_model_tag("extract_hidden_states"):
67
+ return get_model(
68
+ vllm_config=self.vllm_config,
69
+ model_config=self.draft_model_config,
70
+ )
71
+
72
+ def load_model(self, target_model: nn.Module) -> None:
73
+ super().load_model(target_model)
74
+ if len(self.draft_attn_layer_names) != 1:
75
+ raise ValueError(
76
+ "ExtractHiddenStatesModel should have exactly one attention "
77
+ f"layer, found {len(self.draft_attn_layer_names)}"
78
+ )
79
+
80
+ def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
81
+ del cudagraph_mode
82
+
83
+ def capture(self) -> None:
84
+ return None
85
+
86
+ @torch.inference_mode()
87
+ def propose(
88
+ self,
89
+ input_batch: InputBatch,
90
+ attn_metadata: dict[str, Any],
91
+ slot_mappings: dict[str, torch.Tensor],
92
+ last_hidden_states: torch.Tensor,
93
+ aux_hidden_states: list[torch.Tensor] | None,
94
+ num_sampled: torch.Tensor,
95
+ num_rejected: torch.Tensor,
96
+ last_sampled: torch.Tensor,
97
+ next_prefill_tokens: torch.Tensor,
98
+ temperature: torch.Tensor,
99
+ seeds: torch.Tensor,
100
+ num_tokens_across_dp: torch.Tensor | None = None,
101
+ dummy_run: bool = False,
102
+ skip_attn_for_dummy_run: bool = False,
103
+ mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
104
+ is_profile: bool = False,
105
+ ) -> torch.Tensor:
106
+ del (
107
+ last_hidden_states,
108
+ num_sampled,
109
+ num_rejected,
110
+ next_prefill_tokens,
111
+ temperature,
112
+ seeds,
113
+ dummy_run,
114
+ mm_inputs,
115
+ is_profile,
116
+ )
117
+
118
+ draft_tokens = last_sampled[input_batch.idx_mapping, :1]
119
+ if skip_attn_for_dummy_run:
120
+ return draft_tokens
121
+ if aux_hidden_states is None:
122
+ raise ValueError(
123
+ "aux_hidden_states are required when using extract_hidden_states"
124
+ )
125
+ if len(aux_hidden_states) != self.num_hidden_states:
126
+ raise ValueError(
127
+ f"Expected {self.num_hidden_states} auxiliary hidden states, "
128
+ f"got {len(aux_hidden_states)}"
129
+ )
130
+
131
+ stacked_hidden_states = torch.stack(aux_hidden_states, dim=1)
132
+ num_tokens = stacked_hidden_states.shape[0]
133
+ self.hidden_states[:num_tokens].copy_(stacked_hidden_states)
134
+
135
+ draft_attn_metadata = {
136
+ name: attn_metadata[name] for name in self.draft_attn_layer_names
137
+ }
138
+ draft_slot_mappings = {
139
+ name: slot_mappings[name][:num_tokens]
140
+ for name in self.draft_attn_layer_names
141
+ }
142
+ with set_forward_context(
143
+ draft_attn_metadata,
144
+ self.vllm_config,
145
+ num_tokens=num_tokens,
146
+ num_tokens_across_dp=num_tokens_across_dp,
147
+ cudagraph_runtime_mode=CUDAGraphMode.NONE,
148
+ slot_mapping=draft_slot_mappings,
149
+ is_padding=input_batch.is_padding[:num_tokens],
150
+ ):
151
+ self.model(hidden_states=self.hidden_states[:num_tokens])
152
+
153
+ return draft_tokens