@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9433acb8df
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 7, 11:30 UTC

[Spec Decode] Add hidden states extraction system (#33736)

Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>

Fynn Schmitt-Ulms committed Mar 2, 2026 at 14:29 UTC 9433acb8dfdafa560dbee4d67bc286ab3543db39
16 files changed +2102 -38
examples/offline_inference/extract_hidden_states.py new
+58
@@ -0,0 +1,58 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import tempfile
4 +
5 +from safetensors import safe_open
6 +
7 +from vllm import LLM, SamplingParams
8 +
9 +# Example: Using the custom "extract_hidden_states" speculator method and
10 +# ExampleHiddenStatesConnector to extract and save hidden states from vllm
11 +
12 +with tempfile.TemporaryDirectory() as tmpdirname:
13 + llm = LLM(
14 + model="Qwen/Qwen3-8B", # Your target model
15 + speculative_config={
16 + "method": "extract_hidden_states",
17 + "num_speculative_tokens": 1,
18 + "draft_model_config": {
19 + "hf_config": {
20 + "eagle_aux_hidden_state_layer_ids": [ # Target model layer indices
21 + 1,
22 + 2,
23 + 3,
24 + 4,
25 + ],
26 + }
27 + },
28 + },
29 + kv_transfer_config={
30 + "kv_connector": "ExampleHiddenStatesConnector",
31 + "kv_role": "kv_producer",
32 + "kv_connector_extra_config": {
33 + "shared_storage_path": tmpdirname,
34 + },
35 + },
36 + )
37 +
38 + prompts = ["Generate a sentence with hidden states", "Write a python function"]
39 + sampling_params = SamplingParams(max_tokens=1)
40 + outputs = llm.generate(prompts, sampling_params)
41 +
42 + for output in outputs:
43 + print("\nPrompt:", output.prompt)
44 + print("Prompt token ids:", output.prompt_token_ids)
45 +
46 + hidden_states_path = output.kv_transfer_params.get("hidden_states_path")
47 + assert hidden_states_path is not None
48 + print("Prompt hidden states path:", hidden_states_path)
49 +
50 + with safe_open(hidden_states_path, "pt") as f:
51 + token_ids = f.get_tensor("token_ids")
52 + hidden_states = f.get_tensor("hidden_states")
53 +
54 + print("Extracted token ids:", token_ids) # Matches prompt token ids
55 + print(
56 + "Extracted hidden states shape:", hidden_states.shape
57 + ) # [num_hidden_layers, prompt len, hidden size]
58 + print("Extracted hidden states:", hidden_states)
tests/models/registry.py
+5 -1
@@ -108,7 +108,7 @@ class _HfExamplesInfo:
108
109 use_original_num_layers: bool = False
110 """
111 - If True, use the original number of layers from the model config
111 + If True, use the original number of layers from the model config
112 instead of minimal layers for testing.
113 """
114
@@ -1156,6 +1156,10 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1156 speculative_model="LGAI-EXAONE/K-EXAONE-236B-A23B",
1157 min_transformers_version="5.1.0",
1158 ),
1159 + "ExtractHiddenStatesModel": _HfExamplesInfo(
1160 + "Qwen/Qwen3-8B",
1161 + speculative_method="extract_hidden_states",
1162 + ),
1163 "Glm4MoeMTPModel": _HfExamplesInfo(
1164 "zai-org/GLM-4.5",
1165 speculative_model="zai-org/GLM-4.5",
tests/v1/kv_connector/extract_hidden_states_integration/__init__.py
tests/v1/kv_connector/extract_hidden_states_integration/predictable_llama.py new
+120
@@ -0,0 +1,120 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""Predictable dummy model for testing extract_hidden_states.
5 +
6 +Subclasses LlamaForCausalLM but overrides the model to produce deterministic
7 +hidden states: layer i outputs values equal to (i).
8 +"""
9 +
10 +from collections.abc import Iterable
11 +
12 +import torch
13 +import torch.nn as nn
14 +
15 +from vllm.config import VllmConfig
16 +from vllm.model_executor.models.llama import LlamaForCausalLM
17 +from vllm.sequence import IntermediateTensors
18 +
19 +
20 +class PredictableLlamaModel(nn.Module):
21 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
22 + super().__init__()
23 + self.config = vllm_config.model_config.hf_config
24 + self.aux_hidden_state_layers = tuple[int, ...]()
25 +
26 + # Create minimal embed_tokens for embedding
27 + from vllm.model_executor.layers.vocab_parallel_embedding import (
28 + VocabParallelEmbedding,
29 + )
30 +
31 + self.embed_tokens = VocabParallelEmbedding(
32 + self.config.vocab_size,
33 + self.config.hidden_size,
34 + )
35 +
36 + # Required for pipeline parallelism
37 + from vllm.model_executor.models.utils import (
38 + make_empty_intermediate_tensors_factory,
39 + )
40 +
41 + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
42 + ["hidden_states", "residual"], self.config.hidden_size
43 + )
44 +
45 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
46 + """Embed input IDs."""
47 + return self.embed_tokens(input_ids)
48 +
49 + def forward(
50 + self,
51 + input_ids: torch.Tensor | None,
52 + positions: torch.Tensor,
53 + intermediate_tensors: IntermediateTensors | None,
54 + inputs_embeds: torch.Tensor | None = None,
55 + **extra_layer_kwargs,
56 + ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
57 + """Forward pass that produces predictable outputs.
58 +
59 + Returns:
60 + If aux_hidden_state_layers is set: (hidden_states, aux_hidden_states)
61 + Otherwise: hidden_states
62 + """
63 + # Determine sequence length
64 + if inputs_embeds is not None:
65 + seq_len = inputs_embeds.shape[0]
66 + device = inputs_embeds.device
67 + elif input_ids is not None:
68 + seq_len = input_ids.shape[0] if input_ids.ndim == 1 else input_ids.shape[-1]
69 + device = input_ids.device
70 + else:
71 + raise ValueError("Either input_ids or inputs_embeds must be provided")
72 +
73 + # Final hidden states (last layer value)
74 + hidden_states = torch.full(
75 + (seq_len, self.config.hidden_size),
76 + fill_value=float(self.config.num_hidden_layers),
77 + device=device,
78 + dtype=torch.bfloat16,
79 + )
80 +
81 + # Check if we need auxiliary hidden states
82 + if len(self.aux_hidden_state_layers) > 0:
83 + aux_hidden_states = []
84 + for layer_idx in self.aux_hidden_state_layers:
85 + # Fill with (layer_idx) for predictability
86 + layer_hidden = torch.full(
87 + (seq_len, self.config.hidden_size),
88 + fill_value=float(layer_idx),
89 + device=device,
90 + dtype=torch.bfloat16,
91 + )
92 + aux_hidden_states.append(layer_hidden)
93 +
94 + return hidden_states, aux_hidden_states
95 +
96 + return hidden_states
97 +
98 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
99 + """Skip weight loading."""
100 + return set()
101 +
102 +
103 +class PredictableLlamaForCausalLM(LlamaForCausalLM):
104 + """Predictable Llama model for testing.
105 +
106 + Overrides _init_model to use PredictableLlamaModel instead of LlamaModel.
107 + """
108 +
109 + def _init_model(
110 + self,
111 + vllm_config: VllmConfig,
112 + prefix: str = "",
113 + layer_type: type[nn.Module] | None = None,
114 + ):
115 + """Initialize with predictable model."""
116 + return PredictableLlamaModel(vllm_config=vllm_config, prefix=prefix)
117 +
118 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
119 + """Skip weight loading for dummy model."""
120 + return set()
tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py new
+155
@@ -0,0 +1,155 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import gc
5 +import os
6 +
7 +import pytest
8 +import torch
9 +from safetensors import safe_open
10 +
11 +from vllm import LLM, ModelRegistry, SamplingParams
12 +
13 +
14 +def get_and_check_output(output, expected_shape):
15 + assert output.kv_transfer_params is not None
16 + hidden_states_path = output.kv_transfer_params.get("hidden_states_path")
17 + assert hidden_states_path is not None
18 + assert os.path.exists(hidden_states_path)
19 +
20 + # Load and verify the saved tensors
21 + with safe_open(hidden_states_path, "pt") as f:
22 + # Check that token_ids and hidden_states are present
23 + tensor_names = f.keys()
24 + assert "token_ids" in tensor_names
25 + assert "hidden_states" in tensor_names
26 +
27 + token_ids = f.get_tensor("token_ids")
28 + hidden_states = f.get_tensor("hidden_states")
29 +
30 + prompt_token_ids = output.prompt_token_ids
31 + assert torch.equal(token_ids, torch.tensor(prompt_token_ids))
32 +
33 + assert hidden_states.shape == expected_shape
34 +
35 + # Verify hidden_states are not all zeros (i.e., they were actually computed)
36 + assert not torch.allclose(hidden_states, torch.zeros_like(hidden_states))
37 +
38 + return token_ids, hidden_states
39 +
40 +
41 +@pytest.fixture(scope="module")
42 +def predictable_llama_config_path(tmp_path_factory):
43 + """Create a minimal LlamaConfig for PredictableLlamaForCausalLM."""
44 + from transformers import LlamaConfig, LlamaTokenizerFast
45 +
46 + config_dir = tmp_path_factory.mktemp("predictable_llama")
47 +
48 + # Create a minimal Llama config with small dimensions
49 + config = LlamaConfig(
50 + vocab_size=1000,
51 + hidden_size=256,
52 + intermediate_size=512,
53 + num_hidden_layers=24, # Enough layers to test various layer_ids
54 + num_attention_heads=4,
55 + num_key_value_heads=4,
56 + max_position_embeddings=128,
57 + architectures=["PredictableLlamaForCausalLM"],
58 + )
59 +
60 + # Save config
61 + config.save_pretrained(config_dir)
62 +
63 + # Create a simple tokenizer
64 + tokenizer = LlamaTokenizerFast.from_pretrained(
65 + "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
66 + cache_dir=os.path.expanduser("~/.cache/huggingface"),
67 + )
68 + tokenizer.save_pretrained(config_dir)
69 +
70 + return str(config_dir)
71 +
72 +
73 +@pytest.fixture(scope="module", autouse=True)
74 +def register_predictable_model():
75 + """Register the PredictableLlamaForCausalLM model."""
76 + from .predictable_llama import PredictableLlamaForCausalLM
77 +
78 + if "PredictableLlamaForCausalLM" not in ModelRegistry.get_supported_archs():
79 + ModelRegistry.register_model(
80 + "PredictableLlamaForCausalLM", PredictableLlamaForCausalLM
81 + )
82 + yield
83 +
84 +
85 +def test_extract_hidden_states_with_predictable_dummy_model(
86 + predictable_llama_config_path, tmp_path
87 +):
88 + """Comprehensive test using a predictable dummy model with synthetic weights.
89 +
90 + The PredictableLlamaForCausalLM outputs deterministic hidden states where
91 + each layer produces values equal to (layer_index). This test verifies:
92 + 1. Hidden states are correctly extracted from requested layers
93 + 2. Values match the expected predictable pattern
94 + 3. Layer ordering is preserved correctly (non-sequential layer IDs)
95 + 4. Multiple prompts of different lengths produce consistent layer values
96 + """
97 + # Test with non-sequential layer ordering to verify correct association
98 + layer_ids = [5, 2, 10]
99 + num_layers = len(layer_ids)
100 +
101 + llm = LLM(
102 + model=predictable_llama_config_path,
103 + speculative_config={
104 + "method": "extract_hidden_states",
105 + "num_speculative_tokens": 1,
106 + "draft_model_config": {
107 + "hf_config": {"eagle_aux_hidden_state_layer_ids": layer_ids}
108 + },
109 + },
110 + kv_transfer_config={
111 + "kv_connector": "ExampleHiddenStatesConnector",
112 + "kv_role": "kv_producer",
113 + "kv_connector_extra_config": {"shared_storage_path": tmp_path},
114 + },
115 + max_model_len=128,
116 + enforce_eager=True,
117 + trust_remote_code=True,
118 + load_format="dummy", # Don't try to load real weights
119 + )
120 +
121 + # Test with multiple prompts of different lengths
122 + prompts = [
123 + "Short",
124 + "Medium length",
125 + "Much longer prompt with many tokens",
126 + "Much longer prompt with many tokens", # repeated prompt
127 + ]
128 + sampling_params = SamplingParams(max_tokens=1, temperature=0.0)
129 + hidden_size = llm.llm_engine.model_config.get_hidden_size()
130 + outputs = llm.generate(prompts, sampling_params)
131 + del llm
132 + gc.collect()
133 +
134 + assert len(outputs) == len(prompts)
135 +
136 + for output in outputs:
137 + # hidden_states shape is [prompt_len, num_hidden_layers, hidden_size]
138 + expected_shape = (
139 + len(output.prompt_token_ids),
140 + num_layers,
141 + hidden_size,
142 + )
143 + _token_ids, hidden_states = get_and_check_output(output, expected_shape)
144 +
145 + for idx, layer_id in enumerate(layer_ids):
146 + layer_hidden = hidden_states[:, idx, :]
147 + assert torch.allclose(
148 + layer_hidden,
149 + torch.full_like(layer_hidden, layer_id),
150 + atol=1e-5,
151 + ), (
152 + f"Layer {layer_id} at position {idx} should output {float(layer_id)}, "
153 + f"but got mean={layer_hidden.mean():.3f}, "
154 + f"min={layer_hidden.min():.3f}, max={layer_hidden.max():.3f}"
155 + )
tests/v1/spec_decode/test_extract_hidden_states.py new
+346
@@ -0,0 +1,346 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from unittest import mock
5 +
6 +import pytest
7 +import torch
8 +
9 +from tests.v1.attention.utils import (
10 + BatchSpec,
11 + create_common_attn_metadata,
12 +)
13 +from vllm.config import (
14 + AttentionConfig,
15 + CacheConfig,
16 + DeviceConfig,
17 + ModelConfig,
18 + ParallelConfig,
19 + SchedulerConfig,
20 + SpeculativeConfig,
21 + VllmConfig,
22 +)
23 +from vllm.config.load import LoadConfig
24 +from vllm.platforms import current_platform
25 +from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer
26 +from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
27 +
28 +model_dir = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
29 +
30 +
31 +def _create_proposer(
32 + num_speculative_tokens: int = 1,
33 + layer_ids: list[int] | None = None,
34 +) -> ExtractHiddenStatesProposer:
35 + """Create an ExtractHiddenStatesProposer for testing."""
36 + if layer_ids is None:
37 + layer_ids = [1, 2, 3, 4]
38 +
39 + model_config = ModelConfig(model=model_dir, runner="generate", max_model_len=100)
40 +
41 + speculative_config = SpeculativeConfig(
42 + target_model_config=model_config,
43 + target_parallel_config=ParallelConfig(),
44 + method="extract_hidden_states",
45 + num_speculative_tokens=num_speculative_tokens,
46 + draft_model_config={
47 + "hf_config": {
48 + "eagle_aux_hidden_state_layer_ids": layer_ids,
49 + }
50 + },
51 + )
52 +
53 + device = current_platform.device_type
54 + vllm_config = VllmConfig(
55 + model_config=model_config,
56 + cache_config=CacheConfig(),
57 + speculative_config=speculative_config,
58 + device_config=DeviceConfig(device=device),
59 + parallel_config=ParallelConfig(),
60 + load_config=LoadConfig(),
61 + scheduler_config=SchedulerConfig(
62 + max_model_len=model_config.max_model_len,
63 + is_encoder_decoder=model_config.is_encoder_decoder,
64 + ),
65 + attention_config=AttentionConfig(),
66 + )
67 +
68 + return ExtractHiddenStatesProposer(vllm_config=vllm_config, device=device)
69 +
70 +
71 +def test_proposer_initialization():
72 + """Test that the proposer initializes correctly with the right parameters."""
73 + layer_ids = [1, 2, 3, 4]
74 + proposer = _create_proposer(num_speculative_tokens=1, layer_ids=layer_ids)
75 +
76 + assert proposer.num_hidden_states == len(layer_ids)
77 + assert proposer.vllm_config.speculative_config is not None
78 + assert proposer.vllm_config.speculative_config.num_speculative_tokens == 1
79 +
80 + # Verify the hidden states buffer is correctly shaped
81 + expected_shape = (
82 + proposer.max_num_tokens,
83 + len(layer_ids),
84 + proposer.hidden_size,
85 + )
86 + assert proposer.hidden_states.shape == expected_shape
87 +
88 +
89 +def test_proposer_initialization_missing_layer_ids():
90 + """Test that initialization fails when layer_ids are not provided."""
91 + model_config = ModelConfig(model=model_dir, runner="generate", max_model_len=100)
92 +
93 + speculative_config = SpeculativeConfig(
94 + target_model_config=model_config,
95 + target_parallel_config=ParallelConfig(),
96 + method="extract_hidden_states",
97 + num_speculative_tokens=1,
98 + draft_model_config={
99 + "hf_config": {} # Missing eagle_aux_hidden_state_layer_ids
100 + },
101 + )
102 +
103 + device = current_platform.device_type
104 + vllm_config = VllmConfig(
105 + model_config=model_config,
106 + cache_config=CacheConfig(),
107 + speculative_config=speculative_config,
108 + device_config=DeviceConfig(device=device),
109 + parallel_config=ParallelConfig(),
110 + load_config=LoadConfig(),
111 + scheduler_config=SchedulerConfig(
112 + max_model_len=model_config.max_model_len,
113 + is_encoder_decoder=model_config.is_encoder_decoder,
114 + ),
115 + attention_config=AttentionConfig(),
116 + )
117 +
118 + with pytest.raises(
119 + ValueError, match="eagle_aux_hidden_state_layer_ids must be set"
120 + ):
121 + ExtractHiddenStatesProposer(vllm_config=vllm_config, device=device)
122 +
123 +
124 +def test_prepare_next_token_ids_padded():
125 + """
126 + Test for prepare_next_token_ids_padded with extract_hidden_states.
127 +
128 + Since num_speculative_tokens == 1, sampled_token_ids has shape (batch_size, 1).
129 + For each request we either use the sampled token (if valid and not discarded)
130 + or a backup token from the request state.
131 + """
132 + device = torch.device(current_platform.device_type)
133 +
134 + num_requests = 4
135 + batch_spec = BatchSpec(
136 + seq_lens=[5] * num_requests,
137 + query_lens=[5] * num_requests,
138 + )
139 +
140 + req_ids = [f"req_{i + 1}" for i in range(num_requests)]
141 + mock_input_batch = mock.MagicMock(spec=InputBatch)
142 + mock_input_batch.req_ids = req_ids
143 + mock_input_batch.num_reqs = num_requests
144 + mock_input_batch.vocab_size = 100
145 +
146 + mock_requests = {}
147 + for req_id in req_ids:
148 + mock_request = mock.MagicMock(spec=CachedRequestState)
149 + # Each request will have a backup next token id of 10, 20, 30, 40
150 + mock_request.get_token_id.return_value = int(req_id.split("_")[1]) * 10
151 + mock_requests[req_id] = mock_request
152 +
153 + # explicitly discard the last request
154 + discarded_req_mask = torch.tensor(
155 + [False, False, False, True], dtype=torch.bool, device=device
156 + )
157 +
158 + # With num_speculative_tokens=1, sampled_token_ids has shape [batch_size, 1]
159 + sampled_token_ids = torch.tensor(
160 + [
161 + [1], # valid, use 1
162 + [4], # valid, use 4
163 + [-1], # invalid, use backup token "30"
164 + [2], # explicitly discarded, use backup token "40"
165 + ],
166 + dtype=torch.int32,
167 + device=device,
168 + )
169 +
170 + expected_next_token_ids_cpu = [1, 4, 30, 40]
171 + expected_next_token_ids_tensor = torch.tensor(
172 + expected_next_token_ids_cpu, dtype=torch.int32, device=device
173 + )
174 +
175 + proposer = _create_proposer(num_speculative_tokens=1)
176 +
177 + common_attn_metadata = create_common_attn_metadata(
178 + batch_spec,
179 + block_size=16,
180 + device=device,
181 + )
182 +
183 + # valid_sampled_tokens_count tracks if token is valid (not -1 and in vocab range)
184 + # It doesn't depend on whether the request is discarded
185 + expected_valid_sampled_tokens_count = torch.tensor(
186 + [1, 1, 0, 1], dtype=torch.int32, device=device
187 + )
188 +
189 + next_token_ids, valid_sampled_tokens_count = proposer.prepare_next_token_ids_padded(
190 + common_attn_metadata,
191 + sampled_token_ids,
192 + mock_requests,
193 + mock_input_batch,
194 + discarded_req_mask,
195 + )
196 +
197 + assert torch.equal(next_token_ids, expected_next_token_ids_tensor)
198 + assert torch.equal(valid_sampled_tokens_count, expected_valid_sampled_tokens_count)
199 +
200 +
201 +def test_propose():
202 + """
203 + Test the propose() method of ExtractHiddenStatesProposer.
204 +
205 + This should:
206 + 1. Accept target hidden states and sampled token IDs
207 + 2. Return the sampled tokens as "draft" tokens (shape [batch_size, 1])
208 + 3. Cache the hidden states in the model's KV cache
209 + """
210 + device = torch.device(current_platform.device_type)
211 +
212 + # Setup test parameters
213 + batch_size = 2
214 + num_tokens = 5
215 + num_hidden_layers = 4
216 +
217 + proposer = _create_proposer(
218 + num_speculative_tokens=1, layer_ids=list(range(num_hidden_layers))
219 + )
220 + hidden_size = proposer.hidden_size
221 +
222 + # Create mock model
223 + model_mock = mock.MagicMock()
224 + proposer.model = model_mock
225 +
226 + # Mock attention layer names
227 + proposer.attn_layer_names = ["cache_only_layers.28"]
228 +
229 + # Mock attention metadata builder
230 + mock_attn_metadata = mock.MagicMock()
231 + mock_attn_metadata_builder = mock.MagicMock()
232 + mock_attn_metadata_builder.build_for_drafting.return_value = mock_attn_metadata
233 + proposer.attn_metadata_builder = mock_attn_metadata_builder
234 +
235 + # Create input tensors
236 + batch_spec = BatchSpec(
237 + seq_lens=[3, 2],
238 + query_lens=[3, 2],
239 + )
240 +
241 + common_attn_metadata = create_common_attn_metadata(
242 + batch_spec,
243 + block_size=16,
244 + device=device,
245 + )
246 +
247 + # Create target hidden states: list of tensors, one per layer
248 + # Each tensor has shape [num_tokens, hidden_size]
249 + target_hidden_states = [
250 + torch.randn(num_tokens, hidden_size, dtype=proposer.dtype, device=device)
251 + for _ in range(num_hidden_layers)
252 + ]
253 +
254 + # Sampled token IDs from target model
255 + sampled_token_ids = torch.tensor([42, 60], dtype=torch.int32, device=device)
256 +
257 + # Mock scheduler output
258 + mock_scheduler_output = mock.MagicMock()
259 +
260 + # Call propose
261 + with mock.patch(
262 + "vllm.v1.spec_decode.extract_hidden_states.has_kv_transfer_group"
263 + ) as mock_has_kv:
264 + mock_has_kv.return_value = False
265 +
266 + draft_tokens, kv_connector_output = proposer.propose(
267 + sampled_token_ids=sampled_token_ids,
268 + target_hidden_states=target_hidden_states,
269 + common_attn_metadata=common_attn_metadata,
270 + scheduler_output=mock_scheduler_output,
271 + slot_mappings=None,
272 + )
273 +
274 + # Verify draft tokens match sampled tokens
275 + # Shape should be [batch_size, 1] for num_speculative_tokens=1
276 + assert draft_tokens.shape == (batch_size, 1)
277 + assert torch.equal(draft_tokens[:, 0], sampled_token_ids)
278 +
279 + # Verify the model was called
280 + model_mock.assert_called_once()
281 +
282 + # Verify hidden states were copied to the buffer The stacked hidden states
283 + # should have shape [num_tokens, num_hidden_layers, hidden_size]
284 + expected_stacked = torch.stack(target_hidden_states, dim=1)
285 + assert torch.allclose(
286 + proposer.hidden_states[:num_tokens], expected_stacked, atol=1e-6
287 + )
288 +
289 +
290 +@pytest.mark.parametrize("num_hidden_layers", [1, 4, 8])
291 +def test_propose_different_layer_counts(num_hidden_layers):
292 + """Test that propose works correctly with different numbers of hidden layers."""
293 + device = torch.device(current_platform.device_type)
294 +
295 + batch_size = 2
296 + num_tokens = 5
297 +
298 + proposer = _create_proposer(
299 + num_speculative_tokens=1, layer_ids=list(range(num_hidden_layers))
300 + )
301 + hidden_size = proposer.hidden_size
302 +
303 + # Setup mocks
304 + model_mock = mock.MagicMock()
305 + proposer.model = model_mock
306 + proposer.attn_layer_names = ["cache_only_layers.28"]
307 +
308 + mock_attn_metadata_builder = mock.MagicMock()
309 + mock_attn_metadata_builder.build_for_drafting.return_value = mock.MagicMock()
310 + proposer.attn_metadata_builder = mock_attn_metadata_builder
311 +
312 + batch_spec = BatchSpec(
313 + seq_lens=[3, 2],
314 + query_lens=[3, 2],
315 + )
316 +
317 + common_attn_metadata = create_common_attn_metadata(
318 + batch_spec,
319 + block_size=16,
320 + device=device,
321 + )
322 +
323 + # Create target hidden states
324 + target_hidden_states = [
325 + torch.randn(num_tokens, hidden_size, dtype=proposer.dtype, device=device)
326 + for _ in range(num_hidden_layers)
327 + ]
328 +
329 + sampled_token_ids = torch.tensor([42, 60], dtype=torch.int32, device=device)
330 + mock_scheduler_output = mock.MagicMock()
331 +
332 + with mock.patch(
333 + "vllm.v1.spec_decode.extract_hidden_states.has_kv_transfer_group"
334 + ) as mock_has_kv:
335 + mock_has_kv.return_value = False
336 +
337 + draft_tokens, _ = proposer.propose(
338 + sampled_token_ids=sampled_token_ids,
339 + target_hidden_states=target_hidden_states,
340 + common_attn_metadata=common_attn_metadata,
341 + scheduler_output=mock_scheduler_output,
342 + slot_mappings=None,
343 + )
344 +
345 + assert draft_tokens.shape == (batch_size, 1)
346 + assert torch.equal(draft_tokens[:, 0], sampled_token_ids)
vllm/config/speculative.py
+80 -26
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import ast
5 +import copy
6 from typing import TYPE_CHECKING, Any, Literal, get_args
7
8 from pydantic import Field, SkipValidation, model_validator
@@ -45,7 +46,7 @@ MTPModelTypes = Literal[
46 "pangu_ultra_moe_mtp",
47 "step3p5_mtp",
48 ]
48 -EagleModelTypes = Literal["eagle", "eagle3", MTPModelTypes]
49 +EagleModelTypes = Literal["eagle", "eagle3", "extract_hidden_states", MTPModelTypes]
50 SpeculativeMethod = Literal[
51 "ngram",
52 "medusa",
@@ -181,9 +182,22 @@ class SpeculativeConfig:
182 the final hidden states.
183 """
184 factors: list[Any] = []
184 - # Eagle3 affects the computation graph because it returns intermediate
185 - # hidden states in addition to the final hidden state.
186 - factors.append(self.method == "eagle3")
185 + # Eagle3 and extract_hidden_states affect the computation graph because
186 + # they return intermediate hidden states in addition to the final hidden state.
187 + uses_aux_hidden_states = self.method in ("eagle3", "extract_hidden_states")
188 + factors.append(uses_aux_hidden_states)
189 +
190 + # The specific layers used also affect the computation graph
191 + if uses_aux_hidden_states and self.draft_model_config is not None:
192 + layer_ids = getattr(
193 + self.draft_model_config.hf_config,
194 + "eagle_aux_hidden_state_layer_ids",
195 + None,
196 + )
197 + if layer_ids is not None:
198 + # Convert to tuple to make it hashable
199 + factors.append(tuple(layer_ids))
200 +
201 hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
202 return hash_str
203
@@ -352,6 +366,8 @@ class SpeculativeConfig:
366 self.model = "ngram"
367 elif self.method == "suffix":
368 self.model = "suffix"
369 + elif self.method == "extract_hidden_states":
370 + self.model = "extract_hidden_states"
371 else:
372 raise ValueError(
373 "num_speculative_tokens was provided but without speculative model."
@@ -394,6 +410,34 @@ class SpeculativeConfig:
410 self.draft_parallel_config = self.target_parallel_config
411 elif self.method == "suffix":
412 self._validate_suffix_decoding()
413 + elif self.method == "extract_hidden_states":
414 + from vllm.transformers_utils.configs.extract_hidden_states import (
415 + ExtractHiddenStatesConfig,
416 + )
417 +
418 + # ExtractHiddenStatesModel is instantiated manually in load_model()
419 + # We just need to store the target model config for KV cache shape info
420 + self.model = "extract_hidden_states"
421 + self.prompt_lookup_max = 0
422 + self.prompt_lookup_min = 0
423 +
424 + if hasattr(self.draft_model_config, "hf_config"):
425 + hf_config = self.draft_model_config.hf_config.to_dict()
426 + elif (
427 + isinstance(self.draft_model_config, dict)
428 + and "hf_config" in self.draft_model_config
429 + ):
430 + hf_config = self.draft_model_config["hf_config"]
431 + else:
432 + hf_config = {}
433 +
434 + self.draft_model_config = copy.copy(self.target_model_config)
435 + self.draft_model_config.hf_config = ExtractHiddenStatesConfig(
436 + self.draft_model_config.hf_config, **hf_config
437 + )
438 + self.update_arch_()
439 + self.draft_parallel_config = self.target_parallel_config
440 +
441 else:
442 self.prompt_lookup_max = 0
443 self.prompt_lookup_min = 0
@@ -478,23 +522,8 @@ class SpeculativeConfig:
522 method=self.method,
523 model_type="eagle",
524 )
481 - # EAGLEConfig primarily updates architectures, so update
482 - # all architectures-related fields in draft_model_config
525 self.draft_model_config.hf_config = eagle_config
484 - self.draft_model_config.hf_text_config = get_hf_text_config(
485 - self.draft_model_config.hf_config
486 - )
487 - self.draft_model_config.model_arch_config = (
488 - self.draft_model_config.get_model_arch_config()
489 - )
490 - model_info, arch = (
491 - self.draft_model_config.registry.inspect_model_cls(
492 - self.draft_model_config.architectures,
493 - self.draft_model_config,
494 - )
495 - )
496 - self.draft_model_config._model_info = model_info
497 - self.draft_model_config._architecture = arch
526 + self.update_arch_()
527
528 if self.num_speculative_tokens is not None and hasattr(
529 self.draft_model_config.hf_config, "num_lookahead_tokens"
@@ -671,6 +700,24 @@ class SpeculativeConfig:
700 )
701 return speculative_draft_tensor_parallel_size
702
703 + def update_arch_(self):
704 + """
705 + EagleConfig and ExtractHiddenStatesConfig update architectures, so update all
706 + architectures-related fields in self.draft_model_config
707 + """
708 + self.draft_model_config.hf_text_config = get_hf_text_config(
709 + self.draft_model_config.hf_config
710 + )
711 + self.draft_model_config.model_arch_config = (
712 + self.draft_model_config.get_model_arch_config()
713 + )
714 + model_info, arch = self.draft_model_config.registry.inspect_model_cls(
715 + self.draft_model_config.architectures,
716 + self.draft_model_config,
717 + )
718 + self.draft_model_config._model_info = model_info
719 + self.draft_model_config._architecture = arch
720 +
721 @staticmethod
722 def create_draft_parallel_config(
723 target_parallel_config: ParallelConfig,
@@ -718,7 +765,7 @@ class SpeculativeConfig:
765 self.draft_parallel_config
766 )
767
721 - eagle3_target_supported = [
768 + aux_hidden_states_supported = [
769 "llama",
770 "qwen",
771 "minicpm",
@@ -729,16 +776,16 @@ class SpeculativeConfig:
776 "nemotron_h",
777 ]
778 if (
732 - self.method == "eagle3"
779 + self.method in ("eagle3", "extract_hidden_states")
780 and self.target_model_config
781 and not any(
782 supported_model in self.target_model_config.hf_text_config.model_type
736 - for supported_model in eagle3_target_supported
783 + for supported_model in aux_hidden_states_supported
784 )
785 ):
786 raise ValueError(
740 - f"Eagle3 is only supported for {eagle3_target_supported} models. " # noqa: E501
741 - f"Got {self.target_model_config.hf_text_config.model_type=}"
787 + f"{self.method} is only supported for {aux_hidden_states_supported}"
788 + f" models. Got {self.target_model_config.hf_text_config.model_type=}"
789 )
790 self.verify_equal_vocab_size_if_draft_model()
791 return self
@@ -782,8 +829,15 @@ class SpeculativeConfig:
829 def uses_draft_model(self) -> bool:
830 return self.method == "draft_model"
831
832 + def uses_extract_hidden_states(self) -> bool:
833 + return self.method == "extract_hidden_states"
834 +
835 def __repr__(self) -> str:
836 method = self.method
787 - model = None if method in ("ngram", "suffix") else self.draft_model_config.model
837 + model = (
838 + None
839 + if method in ("ngram", "suffix", "extract_hidden_states")
840 + else self.draft_model_config.model
841 + )
842 num_spec_tokens = self.num_speculative_tokens
843 return f"SpeculativeConfig({method=}, {model=}, {num_spec_tokens=})"
vllm/distributed/kv_events.py
+4
@@ -209,6 +209,10 @@ class KVConnectorKVEvents(ABC):
209 def clear_events(self) -> None:
210 raise NotImplementedError
211
212 + def merge(self, other: "KVConnectorKVEvents") -> "KVConnectorKVEvents":
213 + self.add_events(other.get_all_events())
214 + return self
215 +
216
217 class EventPublisher(ABC):
218 """Lightweight publisher for EventBatch batches with data parallelism
vllm/distributed/kv_transfer/kv_connector/factory.py
+6
@@ -149,6 +149,12 @@ KVConnectorFactory.register_connector(
149 "ExampleConnector",
150 )
151
152 +KVConnectorFactory.register_connector(
153 + "ExampleHiddenStatesConnector",
154 + "vllm.distributed.kv_transfer.kv_connector.v1.example_hidden_states_connector",
155 + "ExampleHiddenStatesConnector",
156 +)
157 +
158 KVConnectorFactory.register_connector(
159 "P2pNcclConnector",
160 "vllm.distributed.kv_transfer.kv_connector.v1.p2p.p2p_nccl_connector",
vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py new
+354
@@ -0,0 +1,354 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import os
4 +from dataclasses import dataclass, field
5 +from typing import TYPE_CHECKING, Any, Optional
6 +
7 +import safetensors
8 +import torch
9 +
10 +from vllm.config import VllmConfig, get_layers_from_vllm_config
11 +from vllm.distributed.kv_transfer.kv_connector.v1.base import (
12 + KVConnectorBase_V1,
13 + KVConnectorMetadata,
14 + KVConnectorRole,
15 +)
16 +from vllm.logger import init_logger
17 +from vllm.v1.attention.backend import AttentionMetadata
18 +from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput
19 +
20 +if TYPE_CHECKING:
21 + from vllm.v1.core.kv_cache_manager import KVCacheBlocks
22 + from vllm.v1.kv_cache_interface import KVCacheConfig
23 + from vllm.v1.request import Request
24 +
25 +logger = init_logger(__name__)
26 +
27 +
28 +def extract_from_kv_cache(
29 + kv_cache: torch.Tensor,
30 + slot_mapping: torch.Tensor,
31 + num_tokens: int,
32 +) -> 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]
40 +
41 +
42 +@dataclass
43 +class ReqMeta:
44 + # Request ID
45 + req_id: str
46 + # Request filename
47 + filename: str
48 + # Request tokens
49 + token_ids: torch.Tensor
50 + # Slot mappings, should have the same length as token_ids
51 + slot_mapping: torch.Tensor
52 + # Whether this request is a new request or partially computed already
53 + new_req: bool
54 +
55 + @staticmethod
56 + def make_meta(
57 + req_id: str,
58 + filename: str,
59 + token_ids: list[int],
60 + block_ids: list[int],
61 + block_size: int,
62 + new_req: bool,
63 + ) -> "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()
73 + return ReqMeta(
74 + req_id=req_id,
75 + filename=filename,
76 + token_ids=token_ids_tensor,
77 + slot_mapping=slot_mapping,
78 + new_req=new_req,
79 + )
80 +
81 +
82 +@dataclass
83 +class ExampleHiddenStatesConnectorMetadata(KVConnectorMetadata):
84 + requests: list[ReqMeta] = field(default_factory=list)
85 +
86 + def add_request(
87 + self,
88 + req_id: str,
89 + filename: str,
90 + token_ids: list[int],
91 + block_ids: list[int],
92 + block_size: int,
93 + new_req: bool = True,
94 + ) -> None:
95 + self.requests.append(
96 + ReqMeta.make_meta(
97 + req_id, filename, token_ids, block_ids, block_size, new_req
98 + )
99 + )
100 +
101 +
102 +class ExampleHiddenStatesConnector(KVConnectorBase_V1):
103 + """
104 + Simple debug implementation of a HiddenStatesConnector.
105 +
106 + Simply extracts the hidden states from the kv cache and stores them to disk.
107 + Must be used in conjunction with the `extract_hidden_states` spec decoding method.
108 + """
109 +
110 + @property
111 + def prefer_cross_layer_blocks(self) -> bool:
112 + """
113 + Indicates whether this connector prefers KV blocks that hold KV data for all
114 + layers, which can speed up KV data transfers. Defaults to False.
115 + """
116 + # Must be False so that drafter kv cache isn't merged with verifier's
117 + return False
118 +
119 + def __init__(
120 + self,
121 + vllm_config: "VllmConfig",
122 + role: KVConnectorRole,
123 + kv_cache_config: Optional["KVCacheConfig"] = None,
124 + ):
125 + super().__init__(
126 + vllm_config=vllm_config,
127 + role=role,
128 + kv_cache_config=kv_cache_config,
129 + )
130 + self._block_size = vllm_config.cache_config.block_size
131 + self._storage_path = self._kv_transfer_config.get_from_extra_config(
132 + "shared_storage_path", "/tmp"
133 + )
134 + self.cache_layers: list[str] = [] # set by self.register_kv_caches
135 + logger.info(self._kv_transfer_config)
136 + logger.info("Shared storage path is %s", self._storage_path)
137 +
138 + assert self._vllm_config.speculative_config is not None, (
139 + "ExampleHiddenStatesConnector only works when using "
140 + "'extract_hidden_states' speculative method"
141 + )
142 + spec_config = self._vllm_config.speculative_config.draft_model_config.hf_config
143 + self.num_hidden_states = len(
144 + getattr(spec_config, "eagle_aux_hidden_state_layer_ids", [])
145 + )
146 +
147 + self._request_filenames: dict[str, str] = {}
148 + self._active_requests: dict[str, NewRequestData] = {}
149 + self._req_blocks: dict[str, list[int]] = {}
150 +
151 + # ==============================
152 + # Worker-side methods
153 + # ==============================
154 + def start_load_kv(self, *args, **kwargs: Any) -> None:
155 + pass # Empty implementation of abstract method
156 +
157 + def wait_for_layer_load(self, layer_name: str) -> None:
158 + pass # Empty implementation of abstract method
159 +
160 + def wait_for_save(self):
161 + pass # Empty implementation of abstract method
162 +
163 + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
164 + from vllm.model_executor.models.extract_hidden_states import (
165 + CacheOnlyAttentionLayer,
166 + )
167 +
168 + # Filter layers to only include CacheOnlyAttentionLayers
169 + layers = get_layers_from_vllm_config(
170 + self._vllm_config, CacheOnlyAttentionLayer, list(kv_caches.keys())
171 + )
172 + self.cache_layers = list(layers.keys())
173 + assert len(self.cache_layers) == 1, (
174 + f"Expected 1 CacheOnlyAttentionLayer, got {len(self.cache_layers)}"
175 + )
176 +
177 + def save_kv_layer(
178 + self,
179 + layer_name: str,
180 + kv_layer: torch.Tensor,
181 + attn_metadata: AttentionMetadata,
182 + **kwargs: Any,
183 + ) -> None:
184 + """Start saving the KV cache of the layer from vLLM's paged buffer
185 + to the connector.
186 +
187 + Args:
188 + layer_name (str): the name of the layer.
189 + kv_layer (torch.Tensor): the paged KV buffer of the current
190 + layer in vLLM.
191 + attn_metadata (AttentionMetadata): the attention metadata.
192 + **kwargs: additional arguments for the save operation.
193 + """
194 + if layer_name not in self.cache_layers:
195 + return
196 +
197 + from vllm.model_executor.models.extract_hidden_states import (
198 + CacheOnlyAttentionMetadata,
199 + )
200 +
201 + assert isinstance(attn_metadata, CacheOnlyAttentionMetadata), (
202 + "ExampleHiddenStatesConnector only supports CacheOnlyAttentionBackend"
203 + )
204 +
205 + connector_metadata = self._get_connector_metadata()
206 + assert isinstance(connector_metadata, ExampleHiddenStatesConnectorMetadata)
207 +
208 + os.makedirs(self._storage_path, exist_ok=True)
209 + for request in connector_metadata.requests:
210 + hidden_states = extract_from_kv_cache(
211 + kv_layer, request.slot_mapping, request.token_ids.shape[0]
212 + )
213 + tensors = {
214 + "hidden_states": hidden_states.detach().cpu(),
215 + "token_ids": request.token_ids.detach().cpu(),
216 + }
217 + safetensors.torch.save_file(tensors, request.filename)
218 +
219 + # ==============================
220 + # Scheduler-side methods
221 + # ==============================
222 +
223 + def get_num_new_matched_tokens(
224 + self,
225 + request: "Request",
226 + num_computed_tokens: int,
227 + ) -> tuple[int | None, bool]:
228 + """
229 + Get number of new tokens that can be loaded from the
230 + external KV cache beyond the num_computed_tokens.
231 +
232 + Args:
233 + request (Request): the request object.
234 + num_computed_tokens (int): the number of locally
235 + computed tokens for this request
236 +
237 + Returns:
238 + the number of tokens that can be loaded from the
239 + external KV cache beyond what is already computed.
240 + """
241 + # This connector is store-only, so we don't need to load any tokens
242 + return 0, False
243 +
244 + def update_state_after_alloc(
245 + self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int
246 + ):
247 + # Usually used to handle allocation of new blocks for requests that are loading
248 + # tokens from connector's external kv cache. We never load from external cache
249 + # so this is a no-op.
250 + assert num_external_tokens == 0, "This connector is store-only"
251 +
252 + def build_connector_meta(
253 + self,
254 + scheduler_output: SchedulerOutput,
255 + ) -> KVConnectorMetadata:
256 + """Build the connector metadata for this step.
257 +
258 + This function should NOT modify any fields in the scheduler_output.
259 + Also, calling this function will reset the state of the connector.
260 +
261 + Args:
262 + scheduler_output (SchedulerOutput): the scheduler output object.
263 + """
264 + meta = ExampleHiddenStatesConnectorMetadata()
265 + for new_req in scheduler_output.scheduled_new_reqs:
266 + token_ids = new_req.prompt_token_ids or []
267 + filename = os.path.join(self._storage_path, f"{new_req.req_id}.safetensors")
268 + meta.add_request(
269 + new_req.req_id,
270 + filename=filename,
271 + token_ids=token_ids,
272 + block_ids=new_req.block_ids[0],
273 + block_size=self._block_size,
274 + )
275 + self._request_filenames[new_req.req_id] = filename
276 + self._active_requests[new_req.req_id] = new_req
277 + self._req_blocks[new_req.req_id] = list(new_req.block_ids[0])
278 +
279 + cached_reqs = scheduler_output.scheduled_cached_reqs
280 + for i, req_id in enumerate(cached_reqs.req_ids):
281 + if req_id not in self._active_requests:
282 + continue
283 +
284 + new_block_ids = cached_reqs.new_block_ids[i]
285 +
286 + cached_req = self._active_requests[req_id]
287 + req_block_ids = self._req_blocks[req_id]
288 +
289 + assert new_block_ids is not None
290 + block_ids = new_block_ids[0]
291 +
292 + req_block_ids.extend(block_ids)
293 + filename = os.path.join(self._storage_path, f"{req_id}.safetensors")
294 +
295 + meta.add_request(
296 + req_id=req_id,
297 + filename=filename,
298 + token_ids=cached_req.prompt_token_ids or [],
299 + block_ids=req_block_ids,
300 + block_size=self._block_size,
301 + new_req=False,
302 + )
303 +
304 + return meta
305 +
306 + def request_finished(
307 + self,
308 + request: "Request",
309 + block_ids: list[int],
310 + ) -> tuple[bool, dict[str, Any] | None]:
311 + """
312 + Called exactly once when a request has finished, before its blocks are
313 + freed.
314 +
315 + The connector may assumes responsibility for freeing the blocks
316 + asynchronously by returning True.
317 +
318 + Returns:
319 + True if the request is being saved/sent asynchronously and blocks
320 + should not be freed until the request_id is returned from
321 + get_finished().
322 + Optional KVTransferParams to be included in the request outputs
323 + returned by the engine.
324 + """
325 + req_id = request.request_id
326 + req_filename = self._request_filenames.pop(req_id, None)
327 + _ = self._active_requests.pop(req_id, None)
328 + _ = self._req_blocks.pop(req_id, None)
329 +
330 + return False, {"hidden_states_path": req_filename}
331 +
332 + @classmethod
333 + def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None:
334 + """
335 + Get the required KV cache layout for this connector.
336 + Args:
337 + vllm_config (VllmConfig): the vllm config.
338 +
339 + Returns:
340 + str: the required KV cache layout. e.g. HND, or NHD.
341 + None if the connector does not require a specific layout.
342 + """
343 +
344 + if cls is KVConnectorBase_V1:
345 + raise TypeError(
346 + "get_required_kvcache_layout should not be called "
347 + "on the abstract base class"
348 + )
349 + # NHD means we have (num_tokens, num_heads)
350 + # HND means we have (num_heads, num_tokens)
351 + # For now, we only support NHD layout since this keeps the
352 + # hidden states for each token together in memory.
353 + # HND is primarily used when sharding heads across devices.
354 + return "NHD"
vllm/model_executor/models/extract_hidden_states.py new
+394
@@ -0,0 +1,394 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""Hidden States Extractor Model.
5 +
6 +This model extracts and caches hidden states from the target model
7 +without performing actual token generation. It's used with the
8 +extract_hidden_states speculative decoding method.
9 +"""
10 +
11 +from collections.abc import Iterable
12 +from typing import ClassVar
13 +
14 +import torch
15 +import torch.nn as nn
16 +
17 +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
18 +from vllm.config.cache import CacheDType
19 +from vllm.forward_context import get_forward_context
20 +from vllm.model_executor.layers.attention.attention import set_default_quant_scales
21 +from vllm.model_executor.layers.attention.kv_transfer_utils import (
22 + maybe_transfer_kv_layer,
23 +)
24 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
25 +from vllm.model_executor.models.utils import maybe_prefix
26 +from vllm.utils.torch_utils import kv_cache_dtype_str_to_dtype
27 +from vllm.v1.attention.backend import (
28 + AttentionBackend,
29 + AttentionImpl,
30 + AttentionMetadataBuilder,
31 + AttentionType,
32 + CommonAttentionMetadata,
33 + is_quantized_kv_cache,
34 +)
35 +from vllm.v1.kv_cache_interface import (
36 + AttentionSpec,
37 + KVCacheSpec,
38 + MLAAttentionSpec,
39 +)
40 +
41 +########## Custom Ops ########
42 +
43 +
44 +def unified_kv_cache_update(
45 + to_cache: torch.Tensor,
46 + layer_name: str,
47 +) -> torch.Tensor:
48 + """
49 + Returns a dummy that is passed to unified_attention to signal a side effect and
50 + the data dependency between them to ensure torch.compile preserves ordering.
51 + """
52 + forward_context = get_forward_context()
53 + attn_layer = forward_context.no_compile_layers[layer_name]
54 + kv_cache = attn_layer.kv_cache[forward_context.virtual_engine]
55 +
56 + slot_mapping = forward_context.slot_mapping
57 + assert isinstance(slot_mapping, dict), (
58 + f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
59 + )
60 + layer_slot_mapping = slot_mapping.get(layer_name)
61 + if layer_slot_mapping is not None:
62 + assert hasattr(attn_layer.impl, "do_kv_cache_update"), (
63 + f"{attn_layer.impl.__class__.__name__} does not support kv cache update"
64 + )
65 + attn_layer.impl.do_kv_cache_update(
66 + attn_layer,
67 + to_cache,
68 + kv_cache,
69 + layer_slot_mapping,
70 + )
71 +
72 + return torch.empty(0, device=kv_cache.device, dtype=kv_cache.dtype)
73 +
74 +
75 +@maybe_transfer_kv_layer
76 +def dummy_attention(layer_name, _placeholder):
77 + # Note: layer_name arg required by @maybe_transfer_kv_layer
78 + return _placeholder
79 +
80 +
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]
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
89 +
90 +
91 +######### CacheOnlyAttentionBackend ########
92 +
93 +
94 +class CacheOnlyAttentionBackend(AttentionBackend):
95 + """Attention backend that only caches KV without computing attention."""
96 +
97 + accept_output_buffer: bool = False
98 + supported_dtypes: ClassVar[list[torch.dtype]] = [
99 + torch.float16,
100 + torch.bfloat16,
101 + torch.float32,
102 + ]
103 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
104 + "auto",
105 + "bfloat16",
106 + ]
107 + forward_includes_kv_cache_update: bool = False
108 +
109 + @staticmethod
110 + def get_name() -> str:
111 + return "CACHE_ONLY_ATTN"
112 +
113 + @classmethod
114 + def supports_attn_type(cls, attn_type: str) -> bool:
115 + return attn_type == AttentionType.DECODER
116 +
117 + @classmethod
118 + def supports_mm_prefix(cls) -> bool:
119 + return True
120 +
121 + @staticmethod
122 + def get_impl_cls() -> type["CacheOnlyAttentionImpl"]:
123 + return CacheOnlyAttentionImpl
124 +
125 + @staticmethod
126 + def get_kv_cache_shape(
127 + num_blocks: int,
128 + block_size: int,
129 + num_kv_heads: int,
130 + head_size: int,
131 + cache_dtype_str: str = "auto",
132 + ) -> tuple[int, ...]:
133 + # We set `num_kv_heads = num_hidden_layers` and `head_size = hidden_size`
134 + # We also don't use a k/v (2) dim
135 + return (num_blocks, block_size, num_kv_heads, head_size)
136 +
137 + @staticmethod
138 + def get_builder_cls() -> type["CacheOnlyAttentionMetadataBuilder"]:
139 + return CacheOnlyAttentionMetadataBuilder
140 +
141 + @staticmethod
142 + def use_cascade_attention(*args, **kwargs) -> bool:
143 + return False
144 +
145 + @classmethod
146 + def get_supported_head_sizes(cls) -> list[int]:
147 + return []
148 +
149 +
150 +class CacheOnlyAttentionMetadata:
151 + def __init__(self, slot_mapping: torch.Tensor):
152 + self.slot_mapping = slot_mapping
153 +
154 +
155 +class CacheOnlyAttentionMetadataBuilder(
156 + AttentionMetadataBuilder[CacheOnlyAttentionMetadata]
157 +):
158 + def __init__(
159 + self,
160 + kv_cache_spec: AttentionSpec,
161 + layer_names: list[str],
162 + vllm_config: VllmConfig,
163 + device: torch.device,
164 + ):
165 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
166 +
167 + def build(
168 + self,
169 + common_prefix_len: int,
170 + common_attn_metadata: CommonAttentionMetadata,
171 + fast_build: bool = False,
172 + ) -> CacheOnlyAttentionMetadata:
173 + use_cascade = common_prefix_len > 0
174 + if use_cascade:
175 + raise NotImplementedError(
176 + "Cascade attention not supported by CacheOnlyAttention"
177 + )
178 + causal = common_attn_metadata.causal
179 + if not causal:
180 + raise NotImplementedError(
181 + "Non-causal attention not supported by CacheOnlyAttention"
182 + )
183 +
184 + return CacheOnlyAttentionMetadata(
185 + slot_mapping=common_attn_metadata.slot_mapping,
186 + )
187 +
188 +
189 +class CacheOnlyAttentionImpl(AttentionImpl):
190 + """Attention implementation that only caches KV states."""
191 +
192 + def __init__(
193 + self,
194 + num_heads: int,
195 + head_size: int,
196 + kv_cache_dtype: str,
197 + kv_cache_torch_dtype: torch.dtype,
198 + attn_type: AttentionType = AttentionType.DECODER,
199 + ) -> None:
200 + self.num_heads = num_heads
201 + self.head_size = head_size
202 + self.kv_cache_dtype = kv_cache_dtype
203 + self.kv_cache_torch_dtype = kv_cache_torch_dtype
204 +
205 + if attn_type != AttentionType.DECODER:
206 + raise NotImplementedError(f"Unsupported attention type: {attn_type}")
207 + if is_quantized_kv_cache(kv_cache_dtype):
208 + raise NotImplementedError("Quantized KV cache not supported")
209 +
210 + self.num_queries_per_kv = 1
211 +
212 + def do_kv_cache_update(
213 + self,
214 + layer,
215 + to_cache,
216 + kv_cache,
217 + slot_mapping,
218 + ):
219 + assert to_cache.dtype == self.kv_cache_torch_dtype, (
220 + f"Data to cache must be {self.kv_cache_torch_dtype}, got {to_cache.dtype}"
221 + )
222 + assert kv_cache.dtype == self.kv_cache_torch_dtype, (
223 + f"KV cache must be {self.kv_cache_torch_dtype}, got {kv_cache.dtype}"
224 + )
225 +
226 + basic_cache(to_cache, kv_cache, slot_mapping)
227 +
228 + def forward(self, *args, **kwargs):
229 + # Empty implementation of abstract method
230 + pass
231 +
232 +
233 +############## CacheOnlyAttentionLayer (replaces Attention) ############
234 +
235 +
236 +class CacheOnlyAttentionLayer(nn.Module, AttentionLayerBase):
237 + """Attention layer that only caches key/value states without computing attention."""
238 +
239 + def __init__(
240 + self,
241 + num_heads: int,
242 + head_size: int,
243 + cache_config: CacheConfig | None = None,
244 + prefix: str = "",
245 + attn_type: str = AttentionType.DECODER,
246 + ):
247 + super().__init__()
248 +
249 + self.num_heads = num_heads
250 + self.head_size = head_size
251 + self.layer_name = prefix
252 +
253 + vllm_config = get_current_vllm_config()
254 +
255 + # KV cache configuration
256 + cache_config = cache_config or vllm_config.cache_config
257 + if cache_config is not None:
258 + kv_cache_dtype = cache_config.cache_dtype
259 + self.block_size = cache_config.block_size
260 + else:
261 + kv_cache_dtype = "auto"
262 + self.block_size = 16
263 +
264 + assert kv_cache_dtype in ["auto", "bfloat16", "float16"], (
265 + "CacheOnlyAttentionLayer doesn't currently support quantized kv cache but"
266 + f"kv cache dtype was set to {kv_cache_dtype}"
267 + )
268 + self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
269 + kv_cache_dtype, vllm_config.model_config
270 + )
271 +
272 + # Initialize KV cache quantization attributes
273 + set_default_quant_scales(self, register_buffer=True)
274 +
275 + # Attention backend
276 + self.attn_backend = CacheOnlyAttentionBackend
277 + impl_cls = self.attn_backend.get_impl_cls()
278 + self.impl = impl_cls(
279 + num_heads,
280 + head_size,
281 + kv_cache_dtype,
282 + self.kv_cache_torch_dtype,
283 + attn_type,
284 + )
285 +
286 + assert not self.attn_backend.forward_includes_kv_cache_update, (
287 + "KV cache update should be independent of forward"
288 + )
289 +
290 + # Placeholder KV cache (replaced by bind_kv_cache)
291 + self.kv_cache = [
292 + torch.tensor([])
293 + for _ in range(vllm_config.parallel_config.pipeline_parallel_size)
294 + ]
295 +
296 + # Register in compilation context
297 + compilation_config = vllm_config.compilation_config
298 + if prefix in compilation_config.static_forward_context:
299 + raise ValueError(f"Duplicate layer name: {prefix}")
300 + compilation_config.static_forward_context[prefix] = self
301 +
302 + def forward(self, to_cache: torch.Tensor) -> torch.Tensor:
303 + """Cache hidden states as KV pairs without computing attention.
304 +
305 + Args:
306 + to_cache: The tensor to insert into the kv cache.
307 + shape [num_tokens, num_heads, head_size]
308 +
309 + Returns:
310 + Dummy output tensor (not used)
311 + """
312 + # Note: we set num_heads to num_hidden_layers and
313 + # head_size to hidden_size for hidden states storage
314 + output = torch.empty(0, device=to_cache.device, dtype=to_cache.dtype)
315 +
316 + # Note: dummy_out is used to force torch.compile to preserve ordering between
317 + # cache update and attention op (which triggers kv_connector transfer)
318 + dummy_out = unified_kv_cache_update(to_cache, self.layer_name)
319 +
320 + # Triggers kv_connector transfer via decorator
321 + _ = dummy_attention(self.layer_name, dummy_out)
322 +
323 + return output
324 +
325 + def get_attn_backend(self) -> type[AttentionBackend]:
326 + return self.attn_backend
327 +
328 + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
329 + # Note: we use MLAAttentionSpec here to because it will
330 + # produce page sizes of (block_size * num_kv_heads * head_size * dtype_size)
331 + # whereas FullAttentionSpec will add an additional factor of 2
332 + return MLAAttentionSpec(
333 + block_size=self.block_size,
334 + num_kv_heads=self.num_heads,
335 + head_size=self.head_size,
336 + dtype=self.kv_cache_torch_dtype,
337 + )
338 +
339 +
340 +############ ExtractHiddenStatesModel definition ##########
341 +
342 +
343 +class ExtractHiddenStatesModel(nn.Module):
344 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
345 + super().__init__()
346 +
347 + self.vllm_config = vllm_config
348 + self.hf_config = vllm_config.speculative_config.draft_model_config.hf_config
349 + self.hidden_size = vllm_config.model_config.get_hidden_size()
350 + self.target_num_hidden_layers = (
351 + vllm_config.model_config.get_total_num_hidden_layers()
352 + )
353 + self.num_hidden_states = len(
354 + getattr(self.hf_config, "eagle_aux_hidden_state_layer_ids", [])
355 + )
356 +
357 + cache_config = vllm_config.cache_config
358 +
359 + # Create a single cache-only attention layer
360 + # Note: We set num_heads <- self.num_hidden_states
361 + # and head_size <- hidden_size so that we can insert
362 + # the hidden states directly into the cache without
363 + # reshaping
364 + self.cache_only_layers = nn.ModuleDict(
365 + {
366 + str(self.target_num_hidden_layers): CacheOnlyAttentionLayer(
367 + num_heads=self.num_hidden_states,
368 + head_size=self.hidden_size,
369 + cache_config=cache_config,
370 + prefix=maybe_prefix(
371 + prefix, f"cache_only_layers.{self.target_num_hidden_layers}"
372 + ),
373 + )
374 + }
375 + )
376 +
377 + def forward(self, hidden_states: torch.Tensor) -> None:
378 + """Process and cache hidden states.
379 +
380 + Args:
381 + hidden_states: Hidden states from target model
382 + shape: [num_tokens, num_hidden_states, hidden_size]
383 +
384 + Returns:
385 + Tuple of (dummy_output, dummy_output) - both unused
386 + """
387 +
388 + # Call dummy attention layer to cache hidden states
389 + # Output is ignored - we only care about the KV cache side effects
390 + _ = self.cache_only_layers[str(self.target_num_hidden_layers)](hidden_states)
391 +
392 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
393 + """No weights to load for this dummy model."""
394 + return set()
vllm/model_executor/models/registry.py
+1
@@ -512,6 +512,7 @@ _MULTIMODAL_MODELS = {
512 }
513
514 _SPECULATIVE_DECODING_MODELS = {
515 + "ExtractHiddenStatesModel": ("extract_hidden_states", "ExtractHiddenStatesModel"),
516 "MiMoMTPModel": ("mimo_mtp", "MiMoMTP"),
517 "EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"),
518 "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"),
vllm/transformers_utils/configs/extract_hidden_states.py new
+53
@@ -0,0 +1,53 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Config definitions for ExtractHiddenStatesModel, to be used with
4 +the extract_hidden_states spec decoding method."""
5 +
6 +import os
7 +
8 +from transformers import PretrainedConfig
9 +
10 +
11 +class ExtractHiddenStatesConfig(PretrainedConfig):
12 + model_type = "extract_hidden_states"
13 +
14 + def __init__(
15 + self,
16 + model: PretrainedConfig | dict | None = None,
17 + method: str | None = "extract_hidden_states",
18 + **kwargs,
19 + ):
20 + assert method == "extract_hidden_states"
21 +
22 + if isinstance(model, dict):
23 + model_dict = model
24 + elif isinstance(model, PretrainedConfig):
25 + model_dict = model.to_dict()
26 + else:
27 + model_dict = {}
28 +
29 + # Combine: model_dict first, then kwargs override
30 + combined = {**model_dict, **kwargs}
31 + # Remove architectures from the base, we'll set it explicitly
32 + combined = {k: v for k, v in combined.items() if k != "architectures"}
33 +
34 + combined["architectures"] = ["ExtractHiddenStatesModel"]
35 +
36 + super().__init__(**combined)
37 +
38 + @classmethod
39 + def from_pretrained(
40 + cls,
41 + pretrained_model_name_or_path: str | os.PathLike,
42 + **kwargs,
43 + ) -> "ExtractHiddenStatesConfig":
44 + config_dict, kwargs = cls.get_config_dict(
45 + pretrained_model_name_or_path, **kwargs
46 + )
47 + return cls.from_dict(config_dict, **kwargs)
48 +
49 + def to_json_string(self, use_diff: bool = True) -> str:
50 + # we override use_diff to False as initializing
51 + # ExtractHiddenStatesConfig with default arguments is not supported
52 + del use_diff
53 + return super().to_json_string(use_diff=False)
vllm/v1/outputs.py
+53 -1
@@ -2,8 +2,9 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from abc import ABC, abstractmethod
5 +from collections.abc import Callable
6 from dataclasses import dataclass, field
6 -from typing import TYPE_CHECKING, NamedTuple, TypeAlias
7 +from typing import TYPE_CHECKING, NamedTuple, TypeAlias, TypeVar
8
9 import numpy as np
10 import torch
@@ -120,6 +121,20 @@ class SamplerOutput:
121 logprobs_tensors: LogprobsTensors | None
122
123
124 +T = TypeVar("T")
125 +
126 +
127 +def _combine_non_none(f: Callable[[T, T], T], items: list[T | None]) -> T | None:
128 + non_none = [item for item in items if item is not None]
129 + if len(non_none) == 0:
130 + return None
131 +
132 + combined = non_none[0]
133 + for item in non_none[1:]:
134 + combined = f(combined, item)
135 + return combined
136 +
137 +
138 @dataclass
139 class KVConnectorOutput:
140 # [req_ids]
@@ -146,6 +161,43 @@ class KVConnectorOutput:
161 and not self.invalid_block_ids
162 )
163
164 + @classmethod
165 + def merge(cls, *outputs: "KVConnectorOutput"):
166 + assert len(outputs) > 0, "Cannot merge empty outputs"
167 + finished_sending = _combine_non_none(
168 + set.union, [output.finished_sending for output in outputs]
169 + )
170 + finished_recving = _combine_non_none(
171 + set.union, [output.finished_recving for output in outputs]
172 + )
173 + kv_connector_stats = _combine_non_none(
174 + lambda x, y: x.aggregate(y),
175 + [output.kv_connector_stats for output in outputs],
176 + )
177 + kv_cache_events = _combine_non_none(
178 + lambda x, y: x.merge(y),
179 + [output.kv_cache_events for output in outputs],
180 + )
181 + invalid_block_ids = _combine_non_none(
182 + set.union, [output.invalid_block_ids for output in outputs]
183 + )
184 + assert invalid_block_ids is not None
185 +
186 + assert all(
187 + output.expected_finished_count == outputs[0].expected_finished_count
188 + for output in outputs
189 + )
190 + expected_finished_count = outputs[0].expected_finished_count
191 +
192 + return cls(
193 + finished_sending=finished_sending,
194 + finished_recving=finished_recving,
195 + kv_connector_stats=kv_connector_stats,
196 + kv_cache_events=kv_cache_events,
197 + invalid_block_ids=invalid_block_ids,
198 + expected_finished_count=expected_finished_count,
199 + )
200 +
201
202 @dataclass
203 class ECConnectorOutput:
vllm/v1/spec_decode/extract_hidden_states.py new
+395
@@ -0,0 +1,395 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from __future__ import annotations
5 +
6 +from contextlib import nullcontext
7 +from typing import TYPE_CHECKING
8 +
9 +import torch
10 +import torch.nn as nn
11 +
12 +from vllm.config import CUDAGraphMode, VllmConfig, get_layers_from_vllm_config
13 +from vllm.distributed.kv_transfer import has_kv_transfer_group
14 +from vllm.forward_context import set_forward_context
15 +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
16 +from vllm.model_executor.model_loader import get_model
17 +from vllm.v1.attention.backend import AttentionMetadataBuilder, CommonAttentionMetadata
18 +from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
19 +from vllm.v1.outputs import KVConnectorOutput
20 +from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
21 +from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
22 +from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin
23 +
24 +if TYPE_CHECKING:
25 + from vllm.v1.core.sched.output import SchedulerOutput
26 + from vllm.v1.kv_cache_interface import KVCacheConfig
27 +
28 +PADDING_SLOT_ID = -1
29 +
30 +
31 +class ExtractHiddenStatesProposer:
32 + def __init__(self, vllm_config: VllmConfig, device):
33 + assert vllm_config.speculative_config is not None
34 +
35 + assert vllm_config.speculative_config.num_speculative_tokens == 1
36 + if vllm_config.speculative_config.disable_padded_drafter_batch:
37 + raise ValueError(
38 + "disable_padded_drafter_batch is not supported with "
39 + "extract_hidden_states method"
40 + )
41 + self.vllm_config = vllm_config
42 + self.device = device
43 + self.dtype = vllm_config.model_config.dtype
44 + self.dp_rank = vllm_config.parallel_config.data_parallel_rank
45 +
46 + # Model and attention layer tracking (initialized in load_model)
47 + self.model: nn.Module | None = None
48 + self.attn_layer_names: list[str] = []
49 + self.attn_metadata_builder: AttentionMetadataBuilder | None = None
50 +
51 + # Maximum number of tokens for buffers
52 + max_batch_size = vllm_config.scheduler_config.max_num_seqs
53 + self.max_num_tokens = (
54 + vllm_config.scheduler_config.max_num_batched_tokens + max_batch_size
55 + )
56 +
57 + self.hf_config = vllm_config.speculative_config.draft_model_config.hf_config
58 + layer_ids = getattr(self.hf_config, "eagle_aux_hidden_state_layer_ids", None)
59 + if not layer_ids:
60 + raise ValueError(
61 + "eagle_aux_hidden_state_layer_ids must be set in the draft "
62 + "model config for extract_hidden_states method"
63 + )
64 + self.num_hidden_states = len(layer_ids)
65 + self.hidden_size = vllm_config.model_config.get_hidden_size()
66 + self.hidden_states = torch.zeros(
67 + (self.max_num_tokens, self.num_hidden_states, self.hidden_size),
68 + dtype=self.dtype,
69 + device=device,
70 + )
71 + self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config)
72 +
73 + self._slot_mapping_buffer = torch.zeros(
74 + self.max_num_tokens, dtype=torch.int64, device=device
75 + )
76 +
77 + def propose(
78 + self,
79 + sampled_token_ids: torch.Tensor,
80 + target_hidden_states: list[torch.Tensor],
81 + common_attn_metadata: CommonAttentionMetadata,
82 + scheduler_output: SchedulerOutput,
83 + slot_mappings: dict[str, torch.Tensor]
84 + | list[dict[str, torch.Tensor]]
85 + | None = None,
86 + ) -> tuple[torch.Tensor, KVConnectorOutput | None]:
87 + """Propose draft tokens by calling the ExtractHiddenStatesModel model.
88 +
89 + The ExtractHiddenStatesModel caches the hidden states in the KV cache
90 + without performing actual attention computation. This allows us to
91 + extract and store hidden states for later use (e.g., KV transfer).
92 +
93 + This proposer doesn't actually perform speculation - it returns the
94 + sampled tokens as "draft" tokens, ensuring they always verify (match).
95 + The main purpose is to cache hidden states, not to speculate.
96 +
97 + Args:
98 + sampled_token_ids: Sampled token IDs from the target model
99 + target_hidden_states: List of hidden state tensors from target model
100 + (one per aux hidden state layer)
101 + common_attn_metadata: Attention metadata
102 + scheduler_output: Scheduler output for KV connector
103 + slot_mappings: Slot mappings for KV cache (unused, provided for
104 + interface compatibility)
105 +
106 + Returns:
107 + Tuple of:
108 + - Draft tokens matching sampled tokens, shape [batch_size, 1]
109 + - KV connector output (if KV transfer is active), else None
110 + """
111 + assert self.model is not None and isinstance(target_hidden_states, list)
112 +
113 + # target_hidden_states is a list of tensors (one per layer)
114 + # Each tensor has shape [num_tokens, hidden_size]
115 + # Stack to shape: [num_tokens, num_hidden_states, hidden_size]
116 + stacked_hidden_states = torch.stack(target_hidden_states, dim=1)
117 + num_tokens = stacked_hidden_states.shape[0]
118 +
119 + # Copy hidden states to buffer
120 + self.hidden_states[:num_tokens] = stacked_hidden_states
121 +
122 + assert self.attn_metadata_builder is not None
123 + attn_metadata = self.attn_metadata_builder.build_for_drafting(
124 + common_attn_metadata=common_attn_metadata, draft_index=0
125 + )
126 +
127 + # We assume all cache-only layers belong to the same KV cache group,
128 + # thus using the same attention metadata.
129 + per_layer_attn_metadata = {}
130 + for layer_name in self.attn_layer_names:
131 + per_layer_attn_metadata[layer_name] = attn_metadata
132 +
133 + cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (
134 + self._determine_batch_execution_and_padding(num_tokens)
135 + )
136 + if num_tokens_across_dp is not None:
137 + num_tokens_across_dp[self.dp_rank] = num_input_tokens
138 +
139 + with (
140 + set_forward_context(
141 + per_layer_attn_metadata,
142 + self.vllm_config,
143 + num_tokens=num_input_tokens,
144 + num_tokens_across_dp=num_tokens_across_dp,
145 + cudagraph_runtime_mode=cudagraph_runtime_mode,
146 + slot_mapping=self._get_slot_mapping(
147 + num_input_tokens, common_attn_metadata.slot_mapping
148 + ),
149 + ),
150 + (
151 + KVConnectorModelRunnerMixin._get_kv_connector_output(scheduler_output)
152 + if has_kv_transfer_group()
153 + else nullcontext()
154 + ) as kv_connector_output,
155 + ):
156 + self.model(
157 + hidden_states=self.hidden_states[:num_input_tokens],
158 + )
159 +
160 + # Return the sampled tokens as "draft" tokens
161 + # Shape: [batch_size, 1] to match num_speculative_tokens=1
162 + return sampled_token_ids.unsqueeze(-1), kv_connector_output
163 +
164 + def _get_slot_mapping(
165 + self,
166 + num_tokens: int,
167 + slot_mapping: torch.Tensor | None = None,
168 + ) -> dict[str, torch.Tensor]:
169 + """Return slot_mapping dict for cache-only attention layers.
170 +
171 + If slot_mapping is provided, copies it into the buffer first.
172 + """
173 + if slot_mapping is not None:
174 + num_actual = slot_mapping.shape[0]
175 + self._slot_mapping_buffer[:num_actual].copy_(slot_mapping)
176 + if num_tokens > num_actual:
177 + self._slot_mapping_buffer[num_actual:num_tokens].fill_(PADDING_SLOT_ID)
178 +
179 + view = self._slot_mapping_buffer[:num_tokens]
180 + return {name: view for name in self.attn_layer_names}
181 +
182 + def _determine_batch_execution_and_padding(
183 + self,
184 + num_tokens: int,
185 + use_cudagraphs: bool = True,
186 + ) -> tuple[CUDAGraphMode, int, torch.Tensor | None]:
187 + cudagraph_mode, batch_desc = self.cudagraph_dispatcher.dispatch(
188 + num_tokens,
189 + valid_modes=({CUDAGraphMode.NONE} if not use_cudagraphs else None),
190 + )
191 + num_tokens_padded = batch_desc.num_tokens
192 +
193 + # Extra coordination when running data-parallel since we need to
194 + # coordinate across ranks
195 + # TODO(Flechman): support DBO ubatching
196 + should_ubatch, num_tokens_across_dp = False, None
197 + if self.vllm_config.parallel_config.data_parallel_size > 1:
198 + should_ubatch, num_tokens_across_dp, synced_cudagraph_mode = (
199 + coordinate_batch_across_dp(
200 + num_tokens_unpadded=num_tokens,
201 + parallel_config=self.vllm_config.parallel_config,
202 + allow_microbatching=False,
203 + num_tokens_padded=num_tokens_padded,
204 + cudagraph_mode=cudagraph_mode.value,
205 + )
206 + )
207 + assert not should_ubatch, (
208 + "DBO ubatching not implemented for extract_hidden_states"
209 + )
210 +
211 + # Extract DP-synced values
212 + if num_tokens_across_dp is not None:
213 + dp_rank = self.dp_rank
214 + num_tokens_padded = int(num_tokens_across_dp[dp_rank].item())
215 + # Re-dispatch with DP padding so we have the correct
216 + # batch_descriptor
217 + cudagraph_mode, batch_desc = self.cudagraph_dispatcher.dispatch(
218 + num_tokens_padded,
219 + valid_modes={CUDAGraphMode(synced_cudagraph_mode)},
220 + )
221 + # Assert to make sure the agreed upon token count is correct
222 + # otherwise num_tokens_across_dp will no-longer be valid
223 + assert batch_desc.num_tokens == num_tokens_padded
224 + num_tokens_across_dp[dp_rank] = num_tokens_padded
225 +
226 + return cudagraph_mode, num_tokens_padded, num_tokens_across_dp
227 +
228 + def initialize_cudagraph_keys(self, cudagraph_mode: CUDAGraphMode) -> None:
229 + """Initialize cudagraph dispatcher keys.
230 +
231 + Only supports PIECEWISE cudagraphs (via mixed_mode).
232 + Should be called after adjust_cudagraph_sizes_for_spec_decode.
233 + """
234 + assert self.vllm_config.speculative_config is not None
235 + if (
236 + not self.vllm_config.speculative_config.enforce_eager
237 + and cudagraph_mode.mixed_mode()
238 + in [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL]
239 + ):
240 + proposer_cudagraph_mode = CUDAGraphMode.PIECEWISE
241 + else:
242 + proposer_cudagraph_mode = CUDAGraphMode.NONE
243 +
244 + self.cudagraph_dispatcher.initialize_cudagraph_keys(proposer_cudagraph_mode)
245 +
246 + @torch.inference_mode()
247 + def dummy_run(
248 + self,
249 + num_tokens: int,
250 + use_cudagraphs: bool = True,
251 + is_graph_capturing: bool = False,
252 + slot_mappings: dict[str, torch.Tensor] | None = None,
253 + ) -> None:
254 + assert self.model is not None, "Model must be initialized before dummy_run"
255 + cudagraph_runtime_mode, num_input_tokens, num_tokens_across_dp = (
256 + self._determine_batch_execution_and_padding(
257 + num_tokens, use_cudagraphs=use_cudagraphs
258 + )
259 + )
260 +
261 + if num_tokens_across_dp is not None:
262 + num_tokens_across_dp[self.dp_rank] = num_input_tokens
263 +
264 + # Use our own slot mapping buffer during cudagraph capture.
265 + if (
266 + self.attn_layer_names
267 + and slot_mappings is not None
268 + and self.attn_layer_names[0] in slot_mappings
269 + ):
270 + slot_mapping_dict = self._get_slot_mapping(num_input_tokens)
271 + else:
272 + slot_mapping_dict = slot_mappings or {}
273 +
274 + with set_forward_context(
275 + None,
276 + self.vllm_config,
277 + num_tokens=num_input_tokens,
278 + num_tokens_across_dp=num_tokens_across_dp,
279 + cudagraph_runtime_mode=cudagraph_runtime_mode,
280 + slot_mapping=slot_mapping_dict,
281 + ):
282 + self.model(
283 + hidden_states=self.hidden_states[:num_input_tokens],
284 + )
285 +
286 + def _build_attn_metadata_builder(
287 + self, draft_attn_layers: dict[str, AttentionLayerBase]
288 + ) -> AttentionMetadataBuilder:
289 + """Build the attention metadata builder from draft attention layers."""
290 + if not draft_attn_layers:
291 + raise ValueError("No attention layers found for ExtractHiddenStatesModel")
292 + layer = next(iter(draft_attn_layers.values()))
293 + attn_backend = layer.get_attn_backend()
294 + return attn_backend.get_builder_cls()(
295 + layer.get_kv_cache_spec(self.vllm_config),
296 + self.attn_layer_names,
297 + self.vllm_config,
298 + self.device,
299 + )
300 +
301 + def prepare_next_token_ids_padded(
302 + self,
303 + common_attn_metadata: CommonAttentionMetadata,
304 + sampled_token_ids: torch.Tensor,
305 + requests: dict[str, CachedRequestState],
306 + gpu_input_batch: InputBatch,
307 + discard_request_mask: torch.Tensor,
308 + ) -> tuple[torch.Tensor, torch.Tensor]:
309 + """
310 + Prepare next token IDs for speculative decoding.
311 +
312 + Since num_speculative_tokens == 1, sampled_token_ids has shape
313 + (batch_size, 1). For each request we either use the sampled token
314 + (if valid and not discarded) or a backup token from the request state.
315 + """
316 + num_reqs = gpu_input_batch.num_reqs
317 + device = sampled_token_ids.device
318 +
319 + # Compute backup tokens for discarded / invalid requests
320 + backup_tokens_gpu = torch.tensor(
321 + [
322 + requests[gpu_input_batch.req_ids[i]].get_token_id(
323 + common_attn_metadata.seq_lens_cpu[i].item()
324 + )
325 + for i in range(num_reqs)
326 + ],
327 + dtype=torch.int32,
328 + device=device,
329 + )
330 +
331 + assert discard_request_mask.dtype == torch.bool
332 +
333 + # With num_speculative_tokens == 1, there is exactly one token
334 + sampled = sampled_token_ids[:, 0]
335 + is_valid = (sampled >= 0) & (sampled < gpu_input_batch.vocab_size)
336 + valid_sampled_tokens_count = is_valid.to(torch.int32)
337 +
338 + use_sampled = is_valid & ~discard_request_mask[:num_reqs]
339 + next_token_ids = torch.where(
340 + use_sampled, sampled.to(torch.int32), backup_tokens_gpu
341 + )
342 +
343 + return next_token_ids, valid_sampled_tokens_count
344 +
345 + def load_model(self, target_model: nn.Module) -> None:
346 + """Load the ExtractHiddenStatesModel model.
347 +
348 + This method instantiates the ExtractHiddenStatesModel model which is used
349 + to cache hidden states during speculative decoding. The model uses
350 + cache-only attention (no computation, just caching KV states).
351 +
352 + Args:
353 + target_model: The target model (passed for compatibility with
354 + EagleProposer interface, but not used here)
355 + """
356 + # Get the target model's attention layers before loading draft model
357 + target_attn_layer_names = set(
358 + get_layers_from_vllm_config(self.vllm_config, AttentionLayerBase).keys() # type: ignore[type-abstract]
359 + )
360 +
361 + assert self.vllm_config.speculative_config is not None
362 + draft_model_config = self.vllm_config.speculative_config.draft_model_config
363 + from vllm.compilation.backends import set_model_tag
364 +
365 + with set_model_tag("extract_hidden_states"):
366 + self.model = get_model(
367 + vllm_config=self.vllm_config, model_config=draft_model_config
368 + )
369 +
370 + # Identify draft model's attention layers (difference from target)
371 + all_attn_layers = get_layers_from_vllm_config(
372 + self.vllm_config,
373 + AttentionLayerBase, # type: ignore[type-abstract]
374 + )
375 + draft_attn_layers = {
376 + name: layer
377 + for name, layer in all_attn_layers.items()
378 + if name not in target_attn_layer_names
379 + }
380 + self.attn_layer_names = list(draft_attn_layers.keys())
381 + assert len(draft_attn_layers) == 1, (
382 + "ExtractHiddenStatesModel should have exactly one "
383 + f"attention layer, found {len(draft_attn_layers)}"
384 + )
385 + self.attn_metadata_builder = self._build_attn_metadata_builder(
386 + draft_attn_layers
387 + )
388 +
389 + def validate_same_kv_cache_group(self, kv_cache_config: KVCacheConfig) -> None:
390 + """Validate all drafting layers belong to the same KV cache group.
391 +
392 + With exactly one attention layer (asserted in load_model), this is
393 + trivially satisfied.
394 + """
395 + assert len(self.attn_layer_names) == 1
vllm/v1/worker/gpu_model_runner.py
+78 -10
@@ -159,6 +159,7 @@ from vllm.v1.sample.rejection_sampler import RejectionSampler
159 from vllm.v1.sample.sampler import Sampler
160 from vllm.v1.spec_decode.draft_model import DraftModelProposer
161 from vllm.v1.spec_decode.eagle import EagleProposer
162 +from vllm.v1.spec_decode.extract_hidden_states import ExtractHiddenStatesProposer
163 from vllm.v1.spec_decode.medusa import MedusaProposer
164 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
165 from vllm.v1.spec_decode.suffix_decoding import SuffixDecodingProposer
@@ -495,6 +496,7 @@ class GPUModelRunner(
496 | EagleProposer
497 | DraftModelProposer
498 | MedusaProposer
499 + | ExtractHiddenStatesProposer
500 )
501 if self.speculative_config.method == "ngram":
502 from vllm.v1.spec_decode.ngram_proposer import NgramProposer
@@ -518,6 +520,11 @@ class GPUModelRunner(
520 self.drafter = MedusaProposer(
521 vllm_config=self.vllm_config, device=self.device
522 )
523 + elif self.speculative_config.method == "extract_hidden_states":
524 + self.drafter = ExtractHiddenStatesProposer(
525 + vllm_config=self.vllm_config, device=self.device
526 + )
527 + self.use_aux_hidden_state_outputs = True
528 else:
529 raise ValueError(
530 "Unknown speculative decoding method: "
@@ -3693,10 +3700,9 @@ class GPUModelRunner(
3700 def sample_tokens(
3701 self, grammar_output: "GrammarOutput | None"
3702 ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
3696 - kv_connector_output = self.kv_connector_output
3697 - self.kv_connector_output = None
3698 -
3703 if self.execute_model_state is None:
3704 + kv_connector_output = self.kv_connector_output
3705 + self.kv_connector_output = None
3706 # receive sampled token ids from the last PP rank.
3707 if self.use_async_scheduling and get_pp_group().world_size > 1:
3708 self._pp_receive_prev_sampled_token_ids_to_input_batch()
@@ -3778,12 +3784,17 @@ class GPUModelRunner(
3784 <= self.effective_drafter_max_model_len
3785 )
3786 use_gpu_toks = (
3781 - spec_config.use_eagle() or spec_config.uses_draft_model()
3787 + spec_config.use_eagle()
3788 + or spec_config.uses_draft_model()
3789 + or spec_config.uses_extract_hidden_states()
3790 ) and not spec_config.disable_padded_drafter_batch
3791 if use_gpu_toks:
3792 # EAGLE/DraftModel speculative decoding can use the GPU sampled tokens
3793 # as inputs, and does not need to wait for bookkeeping to finish.
3786 - assert isinstance(self.drafter, EagleProposer | DraftModelProposer)
3794 + assert isinstance(
3795 + self.drafter,
3796 + EagleProposer | DraftModelProposer | ExtractHiddenStatesProposer,
3797 + )
3798 sampled_token_ids = sampler_output.sampled_token_ids
3799 if input_fits_in_drafter:
3800 propose_draft_token_ids(sampled_token_ids)
@@ -3842,6 +3853,10 @@ class GPUModelRunner(
3853 with record_function_or_nullcontext("gpu_model_runner: eplb"):
3854 self.eplb_step()
3855
3856 + # self.kv_connector_output may be modified during drafting
3857 + kv_connector_output = self.kv_connector_output
3858 + self.kv_connector_output = None
3859 +
3860 with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
3861 if self.model_config.enable_return_routed_experts:
3862 capturer = RoutedExpertsCapturer.get_instance()
@@ -4068,6 +4083,48 @@ class GPUModelRunner(
4083 sampling_metadata=sampling_metadata,
4084 slot_mappings=slot_mappings,
4085 )
4086 + elif spec_config.uses_extract_hidden_states():
4087 + assert isinstance(self.drafter, ExtractHiddenStatesProposer)
4088 + assert isinstance(sampled_token_ids, torch.Tensor), (
4089 + "sampled_token_ids should be a torch.Tensor for "
4090 + "extract_hidden_states method."
4091 + )
4092 + if not self.use_aux_hidden_state_outputs or aux_hidden_states is None:
4093 + raise ValueError(
4094 + "aux_hidden_states are required when using `extract_hidden_states`"
4095 + )
4096 + target_hidden_states = [h[:num_scheduled_tokens] for h in aux_hidden_states]
4097 +
4098 + draft_token_ids, drafter_kv_connector_output = self.drafter.propose(
4099 + sampled_token_ids=sampled_token_ids,
4100 + target_hidden_states=target_hidden_states,
4101 + common_attn_metadata=common_attn_metadata,
4102 + scheduler_output=scheduler_output,
4103 + slot_mappings=slot_mappings,
4104 + )
4105 + # Combine KVConnectorOutputs or select the non-empty one
4106 + if self.kv_connector_output and drafter_kv_connector_output:
4107 + self.kv_connector_output = KVConnectorOutput.merge(
4108 + self.kv_connector_output, drafter_kv_connector_output
4109 + )
4110 + else:
4111 + self.kv_connector_output = (
4112 + self.kv_connector_output or drafter_kv_connector_output
4113 + )
4114 +
4115 + next_token_ids, valid_sampled_tokens_count = (
4116 + self.drafter.prepare_next_token_ids_padded(
4117 + common_attn_metadata,
4118 + sampled_token_ids,
4119 + self.requests,
4120 + self.input_batch,
4121 + self.discard_request_mask.gpu,
4122 + )
4123 + )
4124 + self._copy_valid_sampled_token_count(
4125 + next_token_ids, valid_sampled_tokens_count
4126 + )
4127 +
4128 elif spec_config.use_eagle() or spec_config.uses_draft_model():
4129 assert isinstance(self.drafter, EagleProposer | DraftModelProposer)
4130
@@ -4946,8 +5003,12 @@ class GPUModelRunner(
5003 if self.speculative_config and (
5004 self.speculative_config.use_eagle()
5005 or self.speculative_config.uses_draft_model()
5006 + or self.speculative_config.uses_extract_hidden_states()
5007 ):
4950 - assert isinstance(self.drafter, EagleProposer | DraftModelProposer)
5008 + assert isinstance(
5009 + self.drafter,
5010 + EagleProposer | DraftModelProposer | ExtractHiddenStatesProposer,
5011 + )
5012 assert self.speculative_config is not None
5013 # Eagle currently only supports PIECEWISE cudagraphs.
5014 # Therefore only use cudagraphs if the main model uses PIECEWISE
@@ -5656,9 +5717,12 @@ class GPUModelRunner(
5717 cudagraph_mode, self.uniform_decode_query_len
5718 )
5719
5659 - # Initialize eagle's cudagraph dispatcher if using eagle spec decode.
5660 - if self.speculative_config and self.speculative_config.use_eagle():
5661 - assert isinstance(self.drafter, EagleProposer)
5720 + # Initialize drafter's cudagraph dispatcher if using spec decode.
5721 + if self.speculative_config and (
5722 + self.speculative_config.use_eagle()
5723 + or self.speculative_config.uses_extract_hidden_states()
5724 + ):
5725 + assert isinstance(self.drafter, EagleProposer | ExtractHiddenStatesProposer)
5726 self.drafter.initialize_cudagraph_keys(cudagraph_mode)
5727
5728 def calculate_reorder_batch_threshold(self) -> None:
@@ -6025,8 +6089,12 @@ class GPUModelRunner(
6089 if self.speculative_config and (
6090 self.speculative_config.use_eagle()
6091 or self.speculative_config.uses_draft_model()
6092 + or self.speculative_config.uses_extract_hidden_states()
6093 ):
6029 - assert isinstance(self.drafter, EagleProposer | DraftModelProposer)
6094 + assert isinstance(
6095 + self.drafter,
6096 + EagleProposer | DraftModelProposer | ExtractHiddenStatesProposer,
6097 + )
6098 # validate all draft model layers belong to the same kv cache
6099 # group
6100 self.drafter.validate_same_kv_cache_group(kv_cache_config)