A high-throughput and memory-efficient inference and serving engine for LLMs
Hidden states extraction improvements (#43805)
Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Fynn Schmitt-Ulms committed
Jun 11, 2026 at 09:44 UTC
ebc6ef971a71b1a43ec728fae52237524b3056ca
7 files changed
+457
-236
benchmarks/benchmark_hidden_state_extraction.py
-2
@@ -92,7 +92,6 @@ def run_baseline(
92
llm = LLM(
93
model=model,
94
enable_prefix_caching=False,
95
- enable_chunked_prefill=False,
95
**extra_args,
96
)
97
sampling_params = SamplingParams(max_tokens=1)
@@ -194,7 +193,6 @@ async def _run_extraction_async(
193
engine_args = AsyncEngineArgs(
194
model=model,
195
enable_prefix_caching=False,
197
- enable_chunked_prefill=False,
196
max_num_batched_tokens=40960,
197
max_model_len=40960,
198
speculative_config={
docs/features/speculative_decoding/extract_hidden_states.md
+45
-5
@@ -19,7 +19,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import (
19
with tempfile.TemporaryDirectory() as tmpdir:
20
llm = LLM(
21
model="Qwen/Qwen3-8B",
22
- enable_chunked_prefill=False,
22
speculative_config={
23
"method": "extract_hidden_states",
24
"num_speculative_tokens": 1,
@@ -59,17 +58,58 @@ For improved performance, it is recommended to use a RAM-mounted file system suc
58
```bash
59
vllm serve Qwen/Qwen3-8B \
60
--speculative_config '{"method": "extract_hidden_states", "num_speculative_tokens": 1, "draft_model_config": {"hf_config": {"eagle_aux_hidden_state_layer_ids": [1, 2, 3, 4]}}}' \
62
- --kv_transfer_config '{"kv_connector": "ExampleHiddenStatesConnector", "kv_role": "kv_producer", "kv_connector_extra_config": {"shared_storage_path": "/dev/shm/hidden_states"}}' \
63
- --no-enable-chunked-prefill
61
+ --kv_transfer_config '{"kv_connector": "ExampleHiddenStatesConnector", "kv_role": "kv_producer", "kv_connector_extra_config": {"shared_storage_path": "/dev/shm/hidden_states"}}'
62
+```
63
+
64
+## Per-Request Options
65
+
66
+Both offline and online modes support per-request options via `kv_transfer_params`:
67
+
68
+| Parameter | Default | Description |
69
+| --- | --- | --- |
70
+| `hidden_states_path` | Auto-generated | Custom file path for saving hidden states. If not set, files are saved to `<shared_storage_path>/<request_id>.safetensors`. Requires `allow_custom_save_path` to be enabled in the server config. |
71
+| `include_output_tokens` | `False` | When `True`, save hidden states for both prompt and generated output tokens. When `False`, only prompt token hidden states are saved. |
72
+
73
+### Offline usage
74
+
75
+Pass per-request options via `extra_args` on `SamplingParams`:
76
+
77
+```python
78
+SamplingParams(
79
+ max_tokens=32,
80
+ extra_args={
81
+ "kv_transfer_params": {
82
+ "hidden_states_path": "/tmp/my_output.safetensors",
83
+ "include_output_tokens": True,
84
+ }
85
+ },
86
+)
87
+```
88
+
89
+### Online usage
90
+
91
+Pass `kv_transfer_params` as a top-level field in the API request:
92
+
93
+```json
94
+{
95
+ "model": "Qwen/Qwen3-8B",
96
+ "messages": [{"role": "user", "content": "Hello"}],
97
+ "max_tokens": 32,
98
+ "kv_transfer_params": {
99
+ "hidden_states_path": "/tmp/my_output.safetensors",
100
+ "include_output_tokens": true
101
+ }
102
+}
103
```
104
105
## Configuration
106
68
-The `kv_connector_extra_config` dict accepts these options:
107
+The `kv_connector_extra_config` dict accepts these server-level options:
108
109
| Parameter | Default | Description |
110
| --- | --- | --- |
72
-| `shared_storage_path` | `/tmp` | Directory where hidden state files are saved |
111
+| `shared_storage_path` | `/tmp` | Directory where hidden state files are saved (used when `hidden_states_path` is not set per-request) |
112
+| `allow_custom_save_path` | `False` | Allow API clients to specify custom file paths via `hidden_states_path`. When disabled, client-provided paths are ignored with a warning. Enable only with trusted clients — custom paths can write to arbitrary locations on the server. |
113
| `num_writer_threads` | `8` | Thread pool size for async disk writes |
114
| `use_synchronization_lock` | `True` | Use file locks so concurrent readers block until writes complete. Can be disabled for batch generation where synchronization is not needed. |
115
examples/features/speculative_decoding/extract_hidden_states_offline.py
+23
-6
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import os
4
import tempfile
5
6
from vllm import LLM, SamplingParams
@@ -18,7 +19,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import (
19
with tempfile.TemporaryDirectory() as tmpdirname:
20
llm = LLM(
21
model="Qwen/Qwen3-8B", # Your target model
21
- enable_chunked_prefill=False, # required
22
speculative_config={
23
"method": "extract_hidden_states",
24
"num_speculative_tokens": 1,
@@ -38,13 +38,30 @@ with tempfile.TemporaryDirectory() as tmpdirname:
38
kv_role="kv_producer",
39
kv_connector_extra_config={
40
"shared_storage_path": tmpdirname,
41
+ "allow_custom_save_path": True,
42
},
43
),
44
)
45
46
prompts = ["Generate a sentence with hidden states", "Write a python function"]
46
- sampling_params = SamplingParams(max_tokens=1)
47
- outputs = llm.generate(prompts, sampling_params)
47
+
48
+ # One request uses defaults, the other uses a custom save path and
49
+ # includes output token hidden states via per-request kv_transfer_params.
50
+ sampling_params_list = [
51
+ SamplingParams(max_tokens=1),
52
+ SamplingParams(
53
+ max_tokens=10,
54
+ extra_args={
55
+ "kv_transfer_params": {
56
+ "hidden_states_path": os.path.join(
57
+ tmpdirname, "custom_output.safetensors"
58
+ ),
59
+ "include_output_tokens": True,
60
+ }
61
+ },
62
+ ),
63
+ ]
64
+ outputs = llm.generate(prompts, sampling_params_list)
65
66
for output in outputs:
67
print("\nPrompt:", output.prompt)
@@ -52,16 +69,16 @@ with tempfile.TemporaryDirectory() as tmpdirname:
69
70
hidden_states_path = output.kv_transfer_params.get("hidden_states_path")
71
assert hidden_states_path is not None
55
- print("Prompt hidden states path:", hidden_states_path)
72
+ print("Hidden states path:", hidden_states_path)
73
74
obj = example_hidden_states_connector.load_hidden_states(hidden_states_path)
75
token_ids = obj["token_ids"]
76
hidden_states = obj["hidden_states"]
77
61
- print("Extracted token ids:", token_ids) # Matches prompt token ids
78
+ print("Extracted token ids:", token_ids)
79
print(
80
"Extracted hidden states shape:", hidden_states.shape
64
- ) # [prompt_len, num_extracted_layers, hidden_size]
81
+ ) # [num_tokens, num_extracted_layers, hidden_size]
82
print("Extracted hidden states:", hidden_states)
83
84
example_hidden_states_connector.cleanup_hidden_states(hidden_states_path)
tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py
+185
-52
@@ -1,44 +1,40 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import gc
4
import os
5
+import tempfile
6
7
import pytest
8
import torch
9
-from safetensors import safe_open
9
10
+from tests.utils import create_new_process_for_each_test, multi_gpu_test
11
from vllm import LLM, ModelRegistry, SamplingParams
12
+from vllm.distributed.kv_transfer.kv_connector.v1 import (
13
+ example_hidden_states_connector,
14
+)
15
16
17
def get_and_check_output(output, expected_shape):
18
assert output.kv_transfer_params is not None
19
hidden_states_path = output.kv_transfer_params.get("hidden_states_path")
20
assert hidden_states_path is not None
18
- assert os.path.exists(hidden_states_path)
21
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
22
+ obj = example_hidden_states_connector.load_hidden_states(hidden_states_path)
23
+ token_ids = obj["token_ids"]
24
+ hidden_states = obj["hidden_states"]
25
27
- token_ids = f.get_tensor("token_ids")
28
- hidden_states = f.get_tensor("hidden_states")
26
+ prompt_token_ids = output.prompt_token_ids
27
+ assert torch.equal(token_ids, torch.tensor(prompt_token_ids))
28
30
- prompt_token_ids = output.prompt_token_ids
31
- assert torch.equal(token_ids, torch.tensor(prompt_token_ids))
29
+ assert hidden_states.shape == expected_shape
30
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))
31
+ # Verify hidden_states are not all zeros (i.e., they were actually computed)
32
+ assert not torch.allclose(hidden_states, torch.zeros_like(hidden_states))
33
34
return token_ids, hidden_states
35
36
41
-@pytest.fixture(scope="module")
37
+@pytest.fixture
38
def predictable_llama_config_path(tmp_path_factory):
39
"""Create a minimal LlamaConfig for PredictableLlamaForCausalLM."""
40
from transformers import LlamaConfig, LlamaTokenizerFast
@@ -53,7 +49,7 @@ def predictable_llama_config_path(tmp_path_factory):
49
num_hidden_layers=24, # Enough layers to test various layer_ids
50
num_attention_heads=4,
51
num_key_value_heads=4,
56
- max_position_embeddings=128,
52
+ max_position_embeddings=1024,
53
architectures=["PredictableLlamaForCausalLM"],
54
)
55
@@ -85,24 +81,25 @@ def register_predictable_model():
81
def test_extract_hidden_states_with_predictable_dummy_model(
82
predictable_llama_config_path, tmp_path, monkeypatch
83
):
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
84
+ """Test hidden-state extraction with a predictable dummy model.
85
+
86
+ Tests 3 scenarios:
87
+
88
+ 1. **Basic extraction**: non-sequential layer ordering, multiple prompts
89
+ of varying length — verifies correct layer association and
90
+ deterministic values.
91
+ 2. **Chunked prefill**: max_num_batched_tokens=128 with ~500-token
92
+ prompts so each is split across multiple scheduler iterations —
93
+ verifies hidden states are reassembled correctly.
94
+ 3. **Per-request options**: custom hidden_states_path and
95
+ include_output_tokens — verifies per-request kv_transfer_params
96
+ plumbing.
97
"""
97
- # Force fork so the engine worker inherits the autouse fixture's
98
- # ModelRegistry.register_model("PredictableLlamaForCausalLM", ...).
99
- # Spawn (the CI default) starts a fresh Python process that wouldn't
100
- # see the registration.
98
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "fork")
99
103
- # Test with non-sequential layer ordering to verify correct association
100
layer_ids = [5, 2, 10]
101
num_layers = len(layer_ids)
102
+ max_num_batched_tokens = 128
103
104
llm = LLM(
105
model=predictable_llama_config_path,
@@ -116,16 +113,21 @@ def test_extract_hidden_states_with_predictable_dummy_model(
113
kv_transfer_config={
114
"kv_connector": "ExampleHiddenStatesConnector",
115
"kv_role": "kv_producer",
119
- "kv_connector_extra_config": {"shared_storage_path": tmp_path},
116
+ "kv_connector_extra_config": {
117
+ "shared_storage_path": tmp_path,
118
+ "allow_custom_save_path": True,
119
+ },
120
},
121
- max_model_len=128,
121
+ max_model_len=1024,
122
+ max_num_batched_tokens=max_num_batched_tokens,
123
enforce_eager=True,
123
- enable_chunked_prefill=False,
124
trust_remote_code=True,
125
- load_format="dummy", # Don't try to load real weights
125
+ load_format="dummy",
126
)
127
128
- # Test with multiple prompts of different lengths
128
+ hidden_size = llm.llm_engine.model_config.get_hidden_size()
129
+
130
+ # --- Scenario 1: basic extraction with non-sequential layers ----------
131
prompts = [
132
"Short",
133
"Medium length",
@@ -133,15 +135,10 @@ def test_extract_hidden_states_with_predictable_dummy_model(
135
"Much longer prompt with many tokens", # repeated prompt
136
]
137
sampling_params = SamplingParams(max_tokens=1, temperature=0.0)
136
- hidden_size = llm.llm_engine.model_config.get_hidden_size()
138
outputs = llm.generate(prompts, sampling_params)
138
- del llm
139
- gc.collect()
139
140
assert len(outputs) == len(prompts)
142
-
141
for output in outputs:
144
- # hidden_states shape is [prompt_len, num_hidden_layers, hidden_size]
142
expected_shape = (
143
len(output.prompt_token_ids),
144
num_layers,
@@ -156,12 +153,100 @@ def test_extract_hidden_states_with_predictable_dummy_model(
153
torch.full_like(layer_hidden, layer_id),
154
atol=1e-5,
155
), (
159
- f"Layer {layer_id} at position {idx} should output {float(layer_id)}, "
160
- f"but got mean={layer_hidden.mean():.3f}, "
161
- f"min={layer_hidden.min():.3f}, max={layer_hidden.max():.3f}"
156
+ f"Layer {layer_id} at position {idx} should output "
157
+ f"{float(layer_id)}, but got mean="
158
+ f"{layer_hidden.mean():.3f}, min="
159
+ f"{layer_hidden.min():.3f}, max={layer_hidden.max():.3f}"
160
)
161
162
+ # --- Scenario 2: chunked prefill with long prompts --------------------
163
+ long_prompt = " ".join(["word"] * 500)
164
+ chunked_prompts = [
165
+ long_prompt,
166
+ long_prompt + " extra tokens here",
167
+ "Short",
168
+ ]
169
+ outputs = llm.generate(chunked_prompts, sampling_params)
170
171
+ assert len(outputs) == len(chunked_prompts)
172
+ for output in outputs:
173
+ prompt_len = len(output.prompt_token_ids)
174
+ expected_shape = (prompt_len, num_layers, hidden_size)
175
+ _token_ids, hidden_states = get_and_check_output(output, expected_shape)
176
+
177
+ for idx, layer_id in enumerate(layer_ids):
178
+ layer_hidden = hidden_states[:, idx, :]
179
+ assert torch.allclose(
180
+ layer_hidden,
181
+ torch.full_like(layer_hidden, layer_id),
182
+ atol=1e-5,
183
+ ), (
184
+ f"Layer {layer_id} at position {idx} should output "
185
+ f"{float(layer_id)}, but got mean="
186
+ f"{layer_hidden.mean():.3f}, min="
187
+ f"{layer_hidden.min():.3f}, max="
188
+ f"{layer_hidden.max():.3f}. "
189
+ f"prompt_len={prompt_len}, "
190
+ f"max_num_batched_tokens={max_num_batched_tokens}"
191
+ )
192
+
193
+ # --- Scenario 3: per-request options ----------------------------------
194
+ max_tokens = 5
195
+ custom_path = os.path.join(tmp_path, "subdir", "custom.safetensors")
196
+
197
+ sampling_params_list = [
198
+ SamplingParams(max_tokens=max_tokens, temperature=0.0),
199
+ SamplingParams(
200
+ max_tokens=max_tokens,
201
+ temperature=0.0,
202
+ extra_args={
203
+ "kv_transfer_params": {
204
+ "hidden_states_path": custom_path,
205
+ "include_output_tokens": True,
206
+ }
207
+ },
208
+ ),
209
+ ]
210
+ per_req_prompts = ["Short", "Medium length"]
211
+ outputs = llm.generate(per_req_prompts, sampling_params_list)
212
+
213
+ # First output: prompt-only hidden states, default path
214
+ out0 = outputs[0]
215
+ path0 = out0.kv_transfer_params["hidden_states_path"]
216
+ assert path0 != custom_path
217
+ obj0 = example_hidden_states_connector.load_hidden_states(path0)
218
+ assert torch.equal(obj0["token_ids"], torch.tensor(out0.prompt_token_ids))
219
+ assert obj0["hidden_states"].shape == (
220
+ len(out0.prompt_token_ids),
221
+ num_layers,
222
+ hidden_size,
223
+ )
224
+ example_hidden_states_connector.cleanup_hidden_states(path0)
225
+
226
+ # Second output: prompt + output tokens, custom path
227
+ out1 = outputs[1]
228
+ assert out1.kv_transfer_params["hidden_states_path"] == custom_path
229
+ obj1 = example_hidden_states_connector.load_hidden_states(custom_path)
230
+ token_ids = obj1["token_ids"]
231
+ hidden_states = obj1["hidden_states"]
232
+ # The final output token was never an input to the model, so its hidden
233
+ # state is not in the cache — hence the -1.
234
+ total_tokens = len(out1.prompt_token_ids) + len(out1.outputs[0].token_ids) - 1
235
+ assert token_ids.shape[0] == total_tokens
236
+ assert hidden_states.shape == (total_tokens, num_layers, hidden_size)
237
+
238
+ # Verify predictable layer values hold for all tokens (prompt + output)
239
+ for idx, layer_id in enumerate(layer_ids):
240
+ layer_hidden = hidden_states[:, idx, :]
241
+ assert torch.allclose(
242
+ layer_hidden,
243
+ torch.full_like(layer_hidden, layer_id),
244
+ atol=1e-5,
245
+ )
246
+ example_hidden_states_connector.cleanup_hidden_states(custom_path)
247
+
248
+
249
+@create_new_process_for_each_test()
250
def test_extract_hidden_states_qwen35_hybrid_smoke(tmp_path):
251
"""Smoke test for Qwen3.5 hybrid (mamba + full-attention) models.
252
Uses load_format="dummy" to just check shape/plumbing.
@@ -185,7 +270,6 @@ def test_extract_hidden_states_qwen35_hybrid_smoke(tmp_path):
270
},
271
max_model_len=256,
272
enforce_eager=True,
188
- enable_chunked_prefill=False,
273
gpu_memory_utilization=0.4,
274
load_format="dummy",
275
)
@@ -193,19 +277,68 @@ def test_extract_hidden_states_qwen35_hybrid_smoke(tmp_path):
277
prompts = ["Hello world", "Test prompt with several tokens"]
278
sampling_params = SamplingParams(max_tokens=1, temperature=0.0)
279
outputs = llm.generate(prompts, sampling_params)
196
- del llm
197
- gc.collect()
280
281
assert len(outputs) == len(prompts)
282
for output in outputs:
283
assert output.kv_transfer_params is not None
284
hidden_states_path = output.kv_transfer_params.get("hidden_states_path")
285
assert hidden_states_path is not None
204
- assert os.path.exists(hidden_states_path)
286
206
- with safe_open(hidden_states_path, "pt") as f:
207
- token_ids = f.get_tensor("token_ids")
208
- hidden_states = f.get_tensor("hidden_states")
287
+ obj = example_hidden_states_connector.load_hidden_states(hidden_states_path)
288
+ token_ids = obj["token_ids"]
289
+ hidden_states = obj["hidden_states"]
290
+
291
+ assert torch.equal(token_ids, torch.tensor(output.prompt_token_ids))
292
+ assert hidden_states.shape == (
293
+ len(output.prompt_token_ids),
294
+ len(layer_ids),
295
+ hidden_size,
296
+ )
297
+
298
+
299
+@pytest.mark.timeout(60)
300
+@multi_gpu_test(num_gpus=2)
301
+@create_new_process_for_each_test()
302
+def test_extract_hidden_states_tp2():
303
+ """Test that hidden states extraction works with tensor_parallel_size=2."""
304
+ tmp_dir = tempfile.mkdtemp()
305
+ layer_ids = [5, 11, 17]
306
+ hidden_size = 1024 # Qwen/Qwen3-0.6B hidden_size
307
+
308
+ llm = LLM(
309
+ model="Qwen/Qwen3-0.6B",
310
+ tensor_parallel_size=2,
311
+ speculative_config={
312
+ "method": "extract_hidden_states",
313
+ "num_speculative_tokens": 1,
314
+ "draft_model_config": {
315
+ "hf_config": {"eagle_aux_hidden_state_layer_ids": layer_ids}
316
+ },
317
+ },
318
+ kv_transfer_config={
319
+ "kv_connector": "ExampleHiddenStatesConnector",
320
+ "kv_role": "kv_producer",
321
+ "kv_connector_extra_config": {"shared_storage_path": tmp_dir},
322
+ },
323
+ max_model_len=256,
324
+ enforce_eager=True,
325
+ gpu_memory_utilization=0.4,
326
+ load_format="dummy",
327
+ )
328
+
329
+ prompts = ["Hello world", "Test prompt with several tokens"]
330
+ sampling_params = SamplingParams(max_tokens=1, temperature=0.0)
331
+ outputs = llm.generate(prompts, sampling_params)
332
+
333
+ assert len(outputs) == len(prompts)
334
+ for output in outputs:
335
+ assert output.kv_transfer_params is not None
336
+ hidden_states_path = output.kv_transfer_params.get("hidden_states_path")
337
+ assert hidden_states_path is not None
338
+
339
+ obj = example_hidden_states_connector.load_hidden_states(hidden_states_path)
340
+ token_ids = obj["token_ids"]
341
+ hidden_states = obj["hidden_states"]
342
343
assert torch.equal(token_ids, torch.tensor(output.prompt_token_ids))
344
assert hidden_states.shape == (
tests/v1/spec_decode/test_extract_hidden_states.py
-2
@@ -69,7 +69,6 @@ def _create_proposer(
69
scheduler_config=SchedulerConfig(
70
max_model_len=model_config.max_model_len,
71
is_encoder_decoder=model_config.is_encoder_decoder,
72
- enable_chunked_prefill=False,
72
),
73
attention_config=AttentionConfig(),
74
)
@@ -120,7 +119,6 @@ def test_proposer_initialization_missing_layer_ids():
119
scheduler_config=SchedulerConfig(
120
max_model_len=model_config.max_model_len,
121
is_encoder_decoder=model_config.is_encoder_decoder,
123
- enable_chunked_prefill=False,
122
),
123
attention_config=AttentionConfig(),
124
)
vllm/config/vllm.py
-11
@@ -756,17 +756,6 @@ class VllmConfig:
756
Right now, this function reads the offloading settings from
757
CacheConfig and configures the KVTransferConfig accordingly.
758
"""
759
- # Check if KV connector requires chunked prefill to be disabled.
760
- if (
761
- self.kv_transfer_config is not None
762
- and self.kv_transfer_config.kv_connector == "ExampleHiddenStatesConnector"
763
- and self.scheduler_config.enable_chunked_prefill
764
- ):
765
- raise ValueError(
766
- "ExampleHiddenStatesConnector does not support chunked prefill. "
767
- "Please disable chunked prefill (--no-enable-chunked-prefill)."
768
- )
769
-
759
# KV offloading is only activated when kv_offloading_size is set.
760
if (kv_offloading_size := self.cache_config.kv_offloading_size) is None:
761
return
vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
+204
-158
@@ -19,10 +19,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
19
KVConnectorRole,
20
SupportsHMA,
21
)
22
-from vllm.forward_context import get_forward_context
22
+from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
23
from vllm.logger import init_logger
24
from vllm.v1.attention.backend import AttentionMetadata
25
-from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput
25
+from vllm.v1.core.sched.output import SchedulerOutput
26
27
if TYPE_CHECKING:
28
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
@@ -76,43 +76,20 @@ def cleanup_hidden_states(path: str, keep_hidden_states: bool = False) -> None:
76
77
78
@dataclass
79
-class ReqMeta:
80
- # Request ID
79
+class PendingSave:
80
req_id: str
82
- # Request filename
81
filename: str
84
- # Request tokens
82
token_ids: torch.Tensor
86
- # Whether this request is a new request or partially computed already
87
- new_req: bool
88
-
89
- @staticmethod
90
- def make_meta(
91
- req_id: str,
92
- filename: str,
93
- token_ids: list[int],
94
- new_req: bool,
95
- ) -> "ReqMeta":
96
- return ReqMeta(
97
- req_id=req_id,
98
- filename=filename,
99
- token_ids=torch.tensor(token_ids),
100
- new_req=new_req,
101
- )
83
+ block_ids: list[int]
84
85
86
@dataclass
87
class ExampleHiddenStatesConnectorMetadata(KVConnectorMetadata):
106
- requests: list[ReqMeta] = field(default_factory=list)
107
-
108
- def add_request(
109
- self,
110
- req_id: str,
111
- filename: str,
112
- token_ids: list[int],
113
- new_req: bool = True,
114
- ) -> None:
115
- self.requests.append(ReqMeta.make_meta(req_id, filename, token_ids, new_req))
88
+ pending_saves: list[PendingSave] = field(default_factory=list)
89
+ # req_id → filename for newly scheduled requests — the worker pre-creates
90
+ # lock files for these so the lock exists before the client receives the
91
+ # output path.
92
+ new_req_filenames: dict[str, str] = field(default_factory=dict)
93
94
95
class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
@@ -167,9 +144,16 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
144
getattr(spec_config, "eagle_aux_hidden_state_layer_ids", [])
145
)
146
147
+ # Scheduler-side state
148
+ self._pending_saves: dict[str, PendingSave] = {}
149
self._request_filenames: dict[str, str] = {}
171
- self._active_requests: dict[str, NewRequestData] = {}
172
- self._req_blocks: dict[str, list[int]] = {}
150
+
151
+ # Worker-side state (set by register_kv_caches).
152
+ self._kv_cache: torch.Tensor | None = None
153
+ self._hs_group_idx: int = 0
154
+ # Only TP rank 0 writes hidden states to disk; other TP ranks no-op.
155
+ # Set in register_kv_caches (after distributed init).
156
+ self._is_tp_rank_zero: bool = True
157
158
# Async write infrastructure (worker-side).
159
# Dedicated CUDA stream for DtoH copies so they don't block
@@ -184,14 +168,23 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
168
# Whether to use a filesystem lock when writing files to shared storage.
169
# This is necessary for online transfer clients to avoid incomplete reads,
170
# but can be disabled for offline tasks that run tasks in batches to completion
171
+ self.allow_custom_save_path = self._kv_transfer_config.get_from_extra_config(
172
+ "allow_custom_save_path", False
173
+ )
174
+ if self.allow_custom_save_path:
175
+ logger.warning(
176
+ "allow_custom_save_path is enabled. API clients can write "
177
+ "hidden states to arbitrary paths on the server filesystem. "
178
+ "Only enable this with trusted clients."
179
+ )
180
self.use_lock = self._kv_transfer_config.get_from_extra_config(
181
"use_synchronization_lock", True
182
)
190
- # (tensors_dict, copy_done_event, filename, req_id) queued by
191
- # save_kv_layer, submitted to thread pool by wait_for_save.
192
- self._pending_copies: list[
193
- tuple[dict[str, torch.Tensor], torch.cuda.Event, str, str]
194
- ] = []
183
+ # req_id → open fd on the .lock file with LOCK_EX held.
184
+ # Pre-created in wait_for_save when a request first arrives,
185
+ # consumed by _submit_async_write which passes the fd to the
186
+ # thread pool worker for release after writing.
187
+ self._lock_fds: dict[str, int] = {}
188
# req_id → in-flight disk-write Future for that req_id.
189
self._req_futures: dict[str, Future] = {}
190
# req_id → CUDA event marking completion of the DtoH copy. Once
@@ -218,42 +211,28 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
211
def wait_for_layer_load(self, layer_name: str) -> None:
212
pass # Store-only connector — nothing to load
213
221
- def wait_for_save(self):
222
- """Submit pending async copies to the thread pool for disk write.
214
+ def wait_for_save(self) -> None:
215
+ """Pre-create lock files for newly arrived requests.
216
224
- For each pending write we acquire an exclusive flock on a
225
- companion ``.lock`` file **before** submitting to the thread pool.
226
- The thread worker releases the lock after the data file is fully
227
- written. Clients call :func:`load_hidden_states` which takes a
228
- shared flock — the kernel sleeps the client until the writer is
229
- done. Because ``wait_for_save`` runs before the worker returns
230
- output to the scheduler, the lock file is guaranteed to exist
231
- (and be held) by the time the client receives the path.
232
-
233
- The lock can be disabled via the "use_synchronization_lock" extra config.
217
+ This runs on the worker BEFORE the scheduler returns the output
218
+ path to the client, guaranteeing that the lock file exists (and
219
+ LOCK_EX is held) by the time the client tries to open it.
220
"""
235
- for tensors, event, filename, req_id in self._pending_copies:
236
- prior = self._req_futures.get(req_id)
237
- assert prior is None, "Found another KV transfer request with same req_id!"
238
-
239
- lock_fd = None
240
- if self.use_lock:
241
- # Create/open the lock file and acquire an exclusive lock.
242
- # The lock is held by this fd; the thread worker will close
243
- # the fd after writing, which releases the lock.
244
- lock_path = filename + ".lock"
245
- lock_fd = os.open(
246
- lock_path, os.O_CREAT | os.O_WRONLY | os.O_TRUNC, 0o644
247
- )
248
- fcntl.flock(lock_fd, fcntl.LOCK_EX)
249
-
250
- future = self._executor.submit(
251
- self._write_tensors, tensors, event, filename, lock_fd
252
- )
253
- self._req_copy_events[req_id] = event
254
- self._req_futures[req_id] = future
255
- future.add_done_callback(partial(self._on_write_done, req_id))
256
- self._pending_copies.clear()
221
+ if not self._is_tp_rank_zero:
222
+ return
223
+ if not self.use_lock or not self.has_connector_metadata():
224
+ return
225
+ metadata = self._get_connector_metadata()
226
+ if not isinstance(metadata, ExampleHiddenStatesConnectorMetadata):
227
+ return
228
+ for req_id, filename in metadata.new_req_filenames.items():
229
+ if req_id in self._lock_fds:
230
+ continue
231
+ lock_path = filename + ".lock"
232
+ os.makedirs(os.path.dirname(lock_path), exist_ok=True)
233
+ lock_fd = os.open(lock_path, os.O_CREAT | os.O_WRONLY | os.O_TRUNC, 0o644)
234
+ fcntl.flock(lock_fd, fcntl.LOCK_EX)
235
+ self._lock_fds[req_id] = lock_fd
236
237
def _on_write_done(self, req_id: str, future: Future) -> None:
238
"""Surface any exception from the disk-write thread and drop the
@@ -264,6 +243,9 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
243
logger.error("Hidden-states write failed for req_id=%s: %r", req_id, exc)
244
245
def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
246
+ # Delay tp rank0 initialization until after distributed init
247
+ self._is_tp_rank_zero = get_tensor_model_parallel_rank() == 0
248
+
249
from vllm.model_executor.models.extract_hidden_states import (
250
CacheOnlyAttentionLayer,
251
)
@@ -276,6 +258,14 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
258
assert len(self.cache_layers) == 1, (
259
f"Expected 1 CacheOnlyAttentionLayer, got {len(self.cache_layers)}"
260
)
261
+ self._kv_cache = kv_caches[self.cache_layers[0]]
262
+
263
+ # Find the KV cache group index for hidden states
264
+ if self._kv_cache_config is not None:
265
+ for i, group in enumerate(self._kv_cache_config.kv_cache_groups):
266
+ if self.cache_layers[0] in group.layer_names:
267
+ self._hs_group_idx = i
268
+ break
269
270
@staticmethod
271
def _write_tensors(
@@ -304,35 +294,33 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
294
attn_metadata: AttentionMetadata,
295
**kwargs: Any,
296
) -> None:
307
- """Start saving the KV cache of the layer from vLLM's paged buffer
308
- to the connector.
297
+ # Hidden states are already cached by CacheOnlyAttentionLayer during
298
+ # forward. Extraction happens in get_finished once all tokens are done.
299
+ pass
300
310
- Launches an async DtoH copy on a dedicated CUDA stream. The
311
- actual disk write is deferred to wait_for_save() which submits
312
- it to a thread pool.
301
+ def _submit_async_write(
302
+ self,
303
+ pending: PendingSave,
304
+ ) -> None:
305
+ """Extract hidden states from KV cache and submit async DtoH + disk write.
306
314
- Args:
315
- layer_name (str): the name of the layer.
316
- kv_layer (torch.Tensor): the paged KV buffer of the current
317
- layer in vLLM.
318
- attn_metadata (AttentionMetadata): the attention metadata.
319
- **kwargs: additional arguments for the save operation.
307
+ Called from get_finished for each request that has finished generating.
308
"""
321
- if layer_name not in self.cache_layers:
309
+ if not self._is_tp_rank_zero:
310
return
323
-
324
- from vllm.model_executor.models.extract_hidden_states import (
325
- CacheOnlyAttentionMetadata,
326
- )
327
-
328
- assert isinstance(attn_metadata, CacheOnlyAttentionMetadata), (
329
- "ExampleHiddenStatesConnector only supports CacheOnlyAttentionBackend"
311
+ assert self._kv_cache is not None
312
+
313
+ # Compute slot mapping from block_ids
314
+ block_ids_t = torch.tensor(pending.block_ids, dtype=torch.long)
315
+ num_blocks = block_ids_t.shape[0]
316
+ block_offsets = torch.arange(0, self._block_size, dtype=torch.long)
317
+ slot_mapping = (
318
+ block_offsets.reshape((1, self._block_size))
319
+ + block_ids_t.reshape((num_blocks, 1)) * self._block_size
320
)
321
+ slot_mapping = slot_mapping.flatten()
322
332
- connector_metadata = self._get_connector_metadata()
333
- assert isinstance(connector_metadata, ExampleHiddenStatesConnectorMetadata)
334
-
335
- os.makedirs(self._storage_path, exist_ok=True)
323
+ num_tokens = pending.token_ids.shape[0]
324
325
copy_stream = self._get_copy_stream()
326
@@ -341,39 +329,56 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
329
ready_event.record()
330
copy_stream.wait_event(ready_event)
331
344
- slot_mapping = get_forward_context().slot_mapping[layer_name] # type: ignore
345
- offset = 0
346
- for request in connector_metadata.requests:
347
- num_tokens = request.token_ids.shape[0]
348
- with torch.cuda.stream(copy_stream):
349
- req_slot_mapping_gpu = slot_mapping[offset : offset + num_tokens]
350
- assert req_slot_mapping_gpu.device == kv_layer.device
351
- offset += num_tokens
352
-
353
- hidden_states_gpu = extract_from_kv_cache(
354
- kv_layer, req_slot_mapping_gpu, num_tokens
355
- )
356
- # Async DtoH copy into pinned host memory.
357
- pinned_hs = torch.empty_like(
358
- hidden_states_gpu, device="cpu", pin_memory=True
359
- )
360
- pinned_hs.copy_(hidden_states_gpu, non_blocking=True)
361
-
362
- # Record completion of this copy on the copy stream.
363
- copy_done = torch.cuda.Event()
364
- copy_done.record(copy_stream)
365
-
366
- # token_ids is already on CPU (created in ReqMeta.make_meta).
367
- assert not request.token_ids.is_cuda, (
368
- "Expected token_ids on CPU, got CUDA tensor"
332
+ with torch.cuda.stream(copy_stream):
333
+ # Move the CPU slot_mapping to GPU on the copy stream so the
334
+ # implicit H2D inside fancy indexing doesn't sync the default
335
+ # stream.
336
+ slot_mapping_gpu = slot_mapping.to(
337
+ device=self._kv_cache.device, non_blocking=True
338
+ )
339
+ hidden_states_gpu = extract_from_kv_cache(
340
+ self._kv_cache, slot_mapping_gpu, num_tokens
341
)
370
- tensors = {
371
- "hidden_states": pinned_hs,
372
- "token_ids": request.token_ids.clone(),
373
- }
374
- self._pending_copies.append(
375
- (tensors, copy_done, request.filename, request.req_id)
342
+ # Async DtoH copy into pinned host memory.
343
+ pinned_hs = torch.empty_like(
344
+ hidden_states_gpu, device="cpu", pin_memory=True
345
)
346
+ pinned_hs.copy_(hidden_states_gpu, non_blocking=True)
347
+
348
+ # Record completion of this copy on the copy stream.
349
+ copy_done = torch.cuda.Event()
350
+ copy_done.record(copy_stream)
351
+
352
+ # token_ids is already on CPU (created in request_finished).
353
+ assert not pending.token_ids.is_cuda, (
354
+ "Expected token_ids on CPU, got CUDA tensor"
355
+ )
356
+ tensors = {
357
+ "hidden_states": pinned_hs,
358
+ "token_ids": pending.token_ids.clone(),
359
+ }
360
+
361
+ # Submit to thread pool for disk write.
362
+ prior = self._req_futures.get(pending.req_id)
363
+ assert prior is None, "Found another KV transfer request with same req_id!"
364
+
365
+ os.makedirs(os.path.dirname(pending.filename), exist_ok=True)
366
+
367
+ # Use the pre-created lock fd from wait_for_save (already holds
368
+ # LOCK_EX). Falls back to creating one here if use_lock is True
369
+ # but no pre-created fd exists (shouldn't happen in normal flow).
370
+ lock_fd = self._lock_fds.pop(pending.req_id, None)
371
+ if lock_fd is None and self.use_lock:
372
+ lock_path = pending.filename + ".lock"
373
+ lock_fd = os.open(lock_path, os.O_CREAT | os.O_WRONLY | os.O_TRUNC, 0o644)
374
+ fcntl.flock(lock_fd, fcntl.LOCK_EX)
375
+
376
+ future = self._executor.submit(
377
+ self._write_tensors, tensors, copy_done, pending.filename, lock_fd
378
+ )
379
+ self._req_copy_events[pending.req_id] = copy_done
380
+ self._req_futures[pending.req_id] = future
381
+ future.add_done_callback(partial(self._on_write_done, pending.req_id))
382
383
# ==============================
384
# Scheduler-side methods
@@ -421,17 +426,34 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
426
scheduler_output (SchedulerOutput): the scheduler output object.
427
"""
428
meta = ExampleHiddenStatesConnectorMetadata()
429
+
430
+ # Transfer pending saves into metadata (scheduler → worker bridge)
431
+ meta.pending_saves = list(self._pending_saves.values())
432
+ self._pending_saves.clear()
433
+
434
+ # Resolve save paths for new requests and tell the worker so it can
435
+ # pre-create lock files before the client receives the output path.
436
for new_req in scheduler_output.scheduled_new_reqs:
425
- token_ids = new_req.prompt_token_ids or []
426
- filename = os.path.join(self._storage_path, f"{new_req.req_id}.safetensors")
427
- meta.add_request(
428
- new_req.req_id,
429
- filename=filename,
430
- token_ids=token_ids,
437
+ default_path = os.path.join(
438
+ self._storage_path, f"{new_req.req_id}.safetensors"
439
)
440
+ kv_params = (
441
+ new_req.sampling_params.extra_args.get("kv_transfer_params")
442
+ if new_req.sampling_params and new_req.sampling_params.extra_args
443
+ else None
444
+ ) or {}
445
+ custom_path = kv_params.get("hidden_states_path")
446
+ if custom_path is not None and not self.allow_custom_save_path:
447
+ logger.warning(
448
+ "Request %s provided hidden_states_path but "
449
+ "allow_custom_save_path is disabled. Ignoring "
450
+ "custom path and using default.",
451
+ new_req.req_id,
452
+ )
453
+ custom_path = None
454
+ filename = custom_path or default_path
455
self._request_filenames[new_req.req_id] = filename
433
- self._active_requests[new_req.req_id] = new_req
434
- self._req_blocks[new_req.req_id] = list(new_req.block_ids[0])
456
+ meta.new_req_filenames[new_req.req_id] = filename
457
458
return meta
459
@@ -444,35 +466,54 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
466
Called exactly once when a request has finished, before its blocks are
467
freed.
468
447
- The connector may assumes responsibility for freeing the blocks
448
- asynchronously by returning True.
449
-
450
- Returns:
451
- True if the request is being saved/sent asynchronously and blocks
452
- should not be freed until the request_id is returned from
453
- get_finished().
454
- Optional KVTransferParams to be included in the request outputs
455
- returned by the engine.
469
+ Returns True to delay block freeing until get_finished extracts
470
+ the hidden states from the KV cache.
471
"""
472
req_id = request.request_id
458
- req_filename = self._request_filenames.pop(req_id, None)
459
- _ = self._active_requests.pop(req_id, None)
460
- _ = self._req_blocks.pop(req_id, None)
461
-
462
- return True, {"hidden_states_path": req_filename}
473
+ filename = self._request_filenames.pop(req_id)
474
+ kv_params = request.kv_transfer_params or {}
475
+ if kv_params.get("include_output_tokens", False):
476
+ # Exclude the final token — it was the model's output, never an
477
+ # input to a forward pass, so its hidden state is not in the cache.
478
+ token_ids = torch.tensor(list(request.all_token_ids)[:-1])
479
+ elif request.prompt_token_ids is not None:
480
+ token_ids = torch.tensor(request.prompt_token_ids)
481
+ else:
482
+ logger.warning(
483
+ "Request %s has no prompt_token_ids (prompt_embeds only). "
484
+ "Saved token_ids will be empty.",
485
+ req_id,
486
+ )
487
+ token_ids = torch.tensor([], dtype=torch.long)
488
+ self._pending_saves[req_id] = PendingSave(
489
+ req_id=req_id,
490
+ filename=filename,
491
+ token_ids=token_ids,
492
+ block_ids=list(block_ids),
493
+ )
494
+ return True, {"hidden_states_path": filename}
495
496
def get_finished(
497
self, finished_req_ids: set[str]
498
) -> tuple[set[str] | None, set[str] | None]:
467
- """Poll DtoH-copy completion for requests that finished generating.
468
-
469
- The scheduler passes finished_req_ids to tell the worker which
470
- requests are done generating. We accumulate these across calls
471
- and return a request as "finished sending" once its DtoH copy
472
- event is complete (or if it never had a pending copy). The
473
- subsequent disk write may still be in flight; clients block on
474
- the per-file flock to wait for it.
499
+ """Extract hidden states and poll DtoH-copy completion.
500
+
501
+ On the worker side, connector metadata carries pending saves from the
502
+ scheduler. For each one we extract from the KV cache and launch an
503
+ async DtoH copy + thread-pool disk write.
504
+
505
+ We then poll accumulated finished req_ids: a request is "done sending"
506
+ once its DtoH copy event is complete. The subsequent disk write may
507
+ still be in flight; clients block on the per-file flock to wait for it.
508
"""
509
+ # Extract and submit async writes for newly finished requests
510
+ if self.has_connector_metadata():
511
+ connector_metadata = self._get_connector_metadata()
512
+ if isinstance(connector_metadata, ExampleHiddenStatesConnectorMetadata):
513
+ for pending in connector_metadata.pending_saves:
514
+ self._submit_async_write(pending)
515
+
516
+ # Poll for completed DtoH copies
517
self._accumulated_finished_req_ids.update(finished_req_ids)
518
519
done_sending: set[str] = set()
@@ -482,6 +523,11 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
523
self._req_copy_events.pop(req_id, None)
524
done_sending.add(req_id)
525
self._accumulated_finished_req_ids.discard(req_id)
526
+ # Clean up any leftover lock fds (e.g. aborted requests
527
+ # that never went through _submit_async_write).
528
+ lock_fd = self._lock_fds.pop(req_id, None)
529
+ if lock_fd is not None:
530
+ os.close(lock_fd)
531
532
return done_sending or None, None
533
@@ -490,7 +536,7 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
536
request: "Request",
537
block_ids: tuple[list[int], ...],
538
) -> tuple[bool, dict[str, Any] | None]:
493
- return self.request_finished(request, block_ids[0])
539
+ return self.request_finished(request, block_ids[self._hs_group_idx])
540
541
@classmethod
542
def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: