@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4e2eba28be
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 2, 23:15 UTC

[Perf] Optimize hidden state extraction logic (#37374)

Signed-off-by: Benjamin Chislett <bchislett@nvidia.com> Signed-off-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Benjamin Chislett committed May 22, 2026 at 18:23 UTC 4e2eba28beec9972445c338e8ad2080b3cab3246
9 files changed +774 -67
benchmarks/benchmark_hidden_state_extraction.py new
+415
@@ -0,0 +1,415 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Benchmark hidden state extraction throughput.
5 +
6 +Measures two modes:
7 + 1. Baseline: bulk inference with max_tokens=1, no extraction.
8 + 2. Extract: async hidden state extraction via ExampleHiddenStatesConnector
9 + with N concurrent clients, each consuming hidden states as
10 + soon as their request finishes (overlapping I/O with generation).
11 +
12 +Reports tokens/s and prompts/s for each mode.
13 +
14 +Usage:
15 + python benchmarks/benchmark_hidden_state_extraction.py \
16 + --model Qwen/Qwen3-0.6B \
17 + --num-prompts 64 \
18 + --num-clients 8 \
19 + --prompt-len 8192 \
20 + --layers 1 2 3 4
21 +"""
22 +
23 +import argparse
24 +import asyncio
25 +import time
26 +from concurrent.futures import ThreadPoolExecutor
27 +
28 +import torch
29 +from transformers import AutoConfig
30 +
31 +from vllm import LLM, SamplingParams
32 +from vllm.config.kv_transfer import KVTransferConfig
33 +from vllm.distributed.kv_transfer.kv_connector.v1 import (
34 + example_hidden_states_connector,
35 +)
36 +from vllm.engine.arg_utils import AsyncEngineArgs
37 +from vllm.sampling_params import RequestOutputKind
38 +from vllm.v1.engine.async_llm import AsyncLLM
39 +
40 +
41 +def _make_profiler_config(profile_dir: str) -> dict:
42 + """Build a profiler_config dict for torch profiling."""
43 + return {
44 + "profiler": "torch",
45 + "torch_profiler_dir": profile_dir,
46 + "torch_profiler_with_stack": True,
47 + }
48 +
49 +
50 +def make_random_prompts(
51 + num_prompts: int, prompt_len: int, vocab_size: int, seed: int = 42
52 +) -> list[list[int]]:
53 + """Generate lists of random token IDs."""
54 + # Set seed for reproducibility
55 + torch.manual_seed(seed)
56 + return [
57 + torch.randint(0, vocab_size, (prompt_len,)).tolist() for _ in range(num_prompts)
58 + ]
59 +
60 +
61 +def consume_hidden_states(path: str) -> float:
62 + """Load hidden states from disk and compute per-position mean.
63 +
64 + Returns a single float: the grand mean of all hidden state values.
65 + This forces the benchmark to actually read and reduce the data.
66 +
67 + Uses :func:`load_hidden_states` which acquires a shared flock,
68 + blocking (without polling) until the async writer releases its
69 + exclusive lock.
70 + """
71 + obj = example_hidden_states_connector.load_hidden_states(path)
72 + hs = obj["hidden_states"]
73 + total = hs.mean().item()
74 +
75 + example_hidden_states_connector.cleanup_hidden_states(path)
76 +
77 + return total
78 +
79 +
80 +def run_baseline(
81 + model: str,
82 + prompts: list[list[int]],
83 + extra_args: dict,
84 + profile_dir: str | None = None,
85 +) -> dict:
86 + """Baseline: bulk inference, no hidden state extraction."""
87 + if profile_dir:
88 + extra_args = {
89 + **extra_args,
90 + "profiler_config": _make_profiler_config(profile_dir),
91 + }
92 + llm = LLM(
93 + model=model,
94 + enable_prefix_caching=False,
95 + enable_chunked_prefill=False,
96 + **extra_args,
97 + )
98 + sampling_params = SamplingParams(max_tokens=1)
99 + prompt_inputs = [{"prompt_token_ids": p} for p in prompts]
100 +
101 + # Warmup
102 + llm.generate(prompt_inputs[:4], sampling_params, use_tqdm=False)
103 +
104 + if profile_dir:
105 + llm.start_profile()
106 +
107 + t0 = time.perf_counter()
108 + outputs = llm.generate(prompt_inputs, sampling_params, use_tqdm=True)
109 + elapsed = time.perf_counter() - t0
110 +
111 + if profile_dir:
112 + llm.stop_profile()
113 +
114 + total_prompt_tokens = sum(len(o.prompt_token_ids) for o in outputs)
115 + num_prompts = len(outputs)
116 +
117 + del llm
118 + torch.accelerator.empty_cache()
119 +
120 + return {
121 + "mode": "baseline",
122 + "elapsed_s": elapsed,
123 + "num_prompts": num_prompts,
124 + "total_prompt_tokens": total_prompt_tokens,
125 + "tokens_per_s": total_prompt_tokens / elapsed,
126 + "prompts_per_s": num_prompts / elapsed,
127 + }
128 +
129 +
130 +# ---- Async extraction benchmark ----
131 +
132 +
133 +async def _client_loop(
134 + engine: AsyncLLM,
135 + prompt_queue: asyncio.Queue,
136 + consume_pool: ThreadPoolExecutor,
137 + results: list[dict],
138 + client_id: int,
139 +):
140 + """A single async client: pulls prompts, submits to engine, consumes
141 + hidden states as soon as each request finishes."""
142 + loop = asyncio.get_event_loop()
143 + while True:
144 + item = await prompt_queue.get()
145 + if item is None:
146 + prompt_queue.task_done()
147 + break
148 + idx, token_ids = item
149 +
150 + request_id = f"req-{idx}"
151 + sampling_params = SamplingParams(
152 + max_tokens=1,
153 + output_kind=RequestOutputKind.FINAL_ONLY,
154 + )
155 +
156 + final_output = None
157 + async for output in engine.generate(
158 + request_id=request_id,
159 + prompt={"prompt_token_ids": token_ids},
160 + sampling_params=sampling_params,
161 + ):
162 + if output.finished:
163 + final_output = output
164 +
165 + # Consume hidden states on a thread (disk I/O)
166 + path = final_output.kv_transfer_params["hidden_states_path"]
167 + mean_val = await loop.run_in_executor(consume_pool, consume_hidden_states, path)
168 + num_tokens = len(final_output.prompt_token_ids)
169 +
170 + results.append(
171 + {
172 + "request_id": request_id,
173 + "num_prompt_tokens": num_tokens,
174 + "mean_hidden_value": mean_val,
175 + }
176 + )
177 + prompt_queue.task_done()
178 +
179 +
180 +async def _run_extraction_async(
181 + model: str,
182 + prompts: list[list[int]],
183 + num_clients: int,
184 + layers: list[int],
185 + tmpdir: str,
186 + extra_args: dict,
187 + profile_dir: str | None = None,
188 +) -> dict:
189 + if profile_dir:
190 + extra_args = {
191 + **extra_args,
192 + "profiler_config": _make_profiler_config(profile_dir),
193 + }
194 + engine_args = AsyncEngineArgs(
195 + model=model,
196 + enable_prefix_caching=False,
197 + enable_chunked_prefill=False,
198 + max_num_batched_tokens=40960,
199 + max_model_len=40960,
200 + speculative_config={
201 + "method": "extract_hidden_states",
202 + "num_speculative_tokens": 1,
203 + "draft_model_config": {
204 + "hf_config": {
205 + "eagle_aux_hidden_state_layer_ids": layers,
206 + },
207 + },
208 + },
209 + kv_transfer_config=KVTransferConfig(
210 + kv_connector="ExampleHiddenStatesConnector",
211 + kv_role="kv_producer",
212 + kv_connector_extra_config={
213 + "shared_storage_path": tmpdir,
214 + },
215 + ),
216 + **extra_args,
217 + )
218 + engine = AsyncLLM.from_engine_args(engine_args)
219 +
220 + try:
221 + # Warmup: run a few prompts sequentially, cleaning up generated files
222 + for i in range(min(4, len(prompts))):
223 + sp = SamplingParams(max_tokens=1, output_kind=RequestOutputKind.FINAL_ONLY)
224 + final_output = None
225 + async for output in engine.generate(
226 + request_id=f"warmup-{i}",
227 + prompt={"prompt_token_ids": prompts[i]},
228 + sampling_params=sp,
229 + ):
230 + if output.finished:
231 + final_output = output
232 + if final_output and final_output.kv_transfer_params:
233 + path = final_output.kv_transfer_params.get("hidden_states_path")
234 + if path:
235 + example_hidden_states_connector.cleanup_hidden_states(path)
236 +
237 + if profile_dir:
238 + await engine.start_profile()
239 +
240 + # Fill prompt queue
241 + prompt_queue: asyncio.Queue = asyncio.Queue()
242 + for idx, token_ids in enumerate(prompts):
243 + prompt_queue.put_nowait((idx, token_ids))
244 + # Sentinel per client
245 + for _ in range(num_clients):
246 + prompt_queue.put_nowait(None)
247 +
248 + results: list[dict] = []
249 + consume_pool = ThreadPoolExecutor(max_workers=num_clients)
250 +
251 + t0 = time.perf_counter()
252 + tasks = [
253 + asyncio.create_task(
254 + _client_loop(engine, prompt_queue, consume_pool, results, i)
255 + )
256 + for i in range(num_clients)
257 + ]
258 + await asyncio.gather(*tasks)
259 + elapsed = time.perf_counter() - t0
260 +
261 + consume_pool.shutdown(wait=True)
262 +
263 + if profile_dir:
264 + await engine.stop_profile()
265 +
266 + total_prompt_tokens = sum(r["num_prompt_tokens"] for r in results)
267 + num_prompts = len(results)
268 + mean_hidden = sum(r["mean_hidden_value"] for r in results) / max(
269 + len(results), 1
270 + )
271 +
272 + return {
273 + "mode": "extract",
274 + "elapsed_s": elapsed,
275 + "num_prompts": num_prompts,
276 + "total_prompt_tokens": total_prompt_tokens,
277 + "tokens_per_s": total_prompt_tokens / elapsed,
278 + "prompts_per_s": num_prompts / elapsed,
279 + "mean_hidden_value": mean_hidden,
280 + }
281 + finally:
282 + engine.shutdown()
283 +
284 +
285 +def run_extraction(
286 + model: str,
287 + prompts: list[list[int]],
288 + num_clients: int,
289 + layers: list[int],
290 + extra_args: dict,
291 + profile_dir: str | None = None,
292 +) -> dict:
293 + return asyncio.run(
294 + _run_extraction_async(
295 + model,
296 + prompts,
297 + num_clients,
298 + layers,
299 + "/dev/shm",
300 + extra_args,
301 + profile_dir=profile_dir,
302 + )
303 + )
304 +
305 +
306 +def print_results(results: dict):
307 + mode = results["mode"]
308 + print(f"\n{'=' * 60}")
309 + print(f" {mode.upper()} RESULTS")
310 + print(f"{'=' * 60}")
311 + print(f" Prompts: {results['num_prompts']}")
312 + print(f" Total prompt tokens: {results['total_prompt_tokens']:,}")
313 + print(f" Wall time: {results['elapsed_s']:.2f}s")
314 + print(f" Tokens/s: {results['tokens_per_s']:,.0f}")
315 + print(f" Prompts/s: {results['prompts_per_s']:.2f}")
316 + if mode == "extract":
317 + print(f" Mean hidden value: {results['mean_hidden_value']:.6f}")
318 + print(f"{'=' * 60}\n")
319 +
320 +
321 +def main():
322 + parser = argparse.ArgumentParser(
323 + description="Benchmark hidden state extraction throughput"
324 + )
325 + parser.add_argument("--model", type=str, required=True)
326 + parser.add_argument("--num-prompts", type=int, default=64)
327 + parser.add_argument("--num-clients", type=int, default=8)
328 + parser.add_argument("--prompt-len", type=int, default=8192)
329 + parser.add_argument("--layers", type=int, nargs="+", default=[1, 2, 3, 4])
330 + parser.add_argument("--skip-baseline", action="store_true")
331 + parser.add_argument("--skip-extract", action="store_true")
332 + parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
333 + parser.add_argument("--max-num-batched-tokens", type=int, default=None)
334 + parser.add_argument("--max-cudagraph-capture-size", type=int, default=None)
335 + parser.add_argument("--max-model-len", type=int, default=None)
336 + parser.add_argument("--enforce-eager", action="store_true")
337 + parser.add_argument("--load-format", type=str, default=None)
338 + parser.add_argument(
339 + "--profile",
340 + action="store_true",
341 + help="Enable torch profiler for both baseline and extraction runs.",
342 + )
343 + parser.add_argument(
344 + "--torch-profiler-dir",
345 + type=str,
346 + default="./vllm_profile",
347 + help="Directory to save torch profiler traces (default: ./vllm_profile).",
348 + )
349 + parser.add_argument(
350 + "--enable-flashinfer-autotune",
351 + action="store_true",
352 + default=False,
353 + help="Enable FlashInfer autotuning (can be slow).",
354 + )
355 + args = parser.parse_args()
356 +
357 + extra_args = {
358 + "gpu_memory_utilization": args.gpu_memory_utilization,
359 + }
360 + if args.max_model_len is not None:
361 + extra_args["max_model_len"] = args.max_model_len
362 + if args.max_num_batched_tokens is not None:
363 + extra_args["max_num_batched_tokens"] = args.max_num_batched_tokens
364 + if args.max_model_len and args.max_num_batched_tokens < args.max_model_len:
365 + raise ValueError(
366 + "max_num_batched_tokens must be >= max_model_len since chunked prefill"
367 + " is not supported by hidden state extraction."
368 + )
369 + if args.enforce_eager:
370 + extra_args["enforce_eager"] = True
371 + if args.load_format is not None:
372 + extra_args["load_format"] = args.load_format
373 + if args.max_cudagraph_capture_size is not None:
374 + extra_args["max_cudagraph_capture_size"] = args.max_cudagraph_capture_size
375 + extra_args["enable_flashinfer_autotune"] = args.enable_flashinfer_autotune
376 +
377 + # Get vocab size from HF config without loading the full model
378 + hf_config = AutoConfig.from_pretrained(args.model, trust_remote_code=True)
379 + vocab_size = hf_config.vocab_size
380 + prompts = make_random_prompts(args.num_prompts, args.prompt_len, vocab_size)
381 + print(
382 + f"Generated {args.num_prompts} prompts, "
383 + f"{args.prompt_len} tokens each (vocab {vocab_size})"
384 + )
385 +
386 + profile_dir = args.torch_profiler_dir if args.profile else None
387 + if profile_dir:
388 + print(f"Torch profiler enabled, traces will be saved to {profile_dir}/")
389 +
390 + if not args.skip_baseline:
391 + baseline_profile_dir = f"{profile_dir}/baseline" if profile_dir else None
392 + baseline = run_baseline(
393 + args.model, prompts, extra_args, profile_dir=baseline_profile_dir
394 + )
395 + print_results(baseline)
396 +
397 + if not args.skip_extract:
398 + extract_profile_dir = f"{profile_dir}/extract" if profile_dir else None
399 + extract = run_extraction(
400 + args.model,
401 + prompts,
402 + args.num_clients,
403 + args.layers,
404 + extra_args,
405 + profile_dir=extract_profile_dir,
406 + )
407 + print_results(extract)
408 +
409 + if not args.skip_baseline and not args.skip_extract:
410 + slowdown = baseline["tokens_per_s"] / extract["tokens_per_s"]
411 + print("Extraction slowdown factor: {:.2f}x".format(slowdown))
412 +
413 +
414 +if __name__ == "__main__":
415 + main()
docs/features/speculative_decoding/README.md
+1
@@ -15,6 +15,7 @@ vLLM supports a variety of methods of speculative decoding. Model-based methods
15 - [Multi-Layer Perceptron](mlp.md)
16 - [N-Gram](n_gram.md)
17 - [Suffix Decoding](suffix.md)
18 +- [Hidden State Extraction](extract_hidden_states.md)
19 - [Custom Proposer Backend (Experimental)](#custom-proposer-backend-experimental)
20
21 ## Method Selection at a Glance
docs/features/speculative_decoding/extract_hidden_states.md new
+86
@@ -0,0 +1,86 @@
1 +# Hidden State Extraction
2 +
3 +The Hidden State Extraction feature allows vLLM to save intermediate layer activations from a target model during inference. This is useful for training [EAGLE](eagle.md)-style draft models, knowledge distillation, or offline analysis of model internals.
4 +
5 +!!! note
6 + It is possible to save the last-layer's output hidden states by passing `num_hidden_layers` as a layer id. Note that these are _not_ normalized using the output norm.
7 +
8 +## Offline Example
9 +
10 +```python
11 +import tempfile
12 +
13 +from vllm import LLM, SamplingParams
14 +from vllm.config.kv_transfer import KVTransferConfig
15 +from vllm.distributed.kv_transfer.kv_connector.v1 import (
16 + example_hidden_states_connector,
17 +)
18 +
19 +with tempfile.TemporaryDirectory() as tmpdir:
20 + llm = LLM(
21 + model="Qwen/Qwen3-8B",
22 + enable_chunked_prefill=False,
23 + speculative_config={
24 + "method": "extract_hidden_states",
25 + "num_speculative_tokens": 1,
26 + "draft_model_config": {
27 + "hf_config": {
28 + "eagle_aux_hidden_state_layer_ids": [1, 2, 3, 4],
29 + },
30 + },
31 + },
32 + kv_transfer_config=KVTransferConfig(
33 + kv_connector="ExampleHiddenStatesConnector",
34 + kv_role="kv_producer",
35 + kv_connector_extra_config={
36 + "shared_storage_path": tmpdir,
37 + },
38 + ),
39 + )
40 +
41 + outputs = llm.generate(
42 + ["The future of AI is"],
43 + SamplingParams(max_tokens=1),
44 + )
45 +
46 + for output in outputs:
47 + path = output.kv_transfer_params["hidden_states_path"]
48 + obj = example_hidden_states_connector.load_hidden_states(path)
49 + print(f"token_ids: {obj['token_ids'].shape}")
50 + print(f"hidden_states: {obj['hidden_states'].shape}")
51 +```
52 +
53 +A complete example is available at [`examples/features/speculative_decoding/extract_hidden_states_offline.py`](../../../examples/features/speculative_decoding/extract_hidden_states_offline.py).
54 +
55 +## Online Example
56 +
57 +For improved performance, it is recommended to use a RAM-mounted file system such as `/dev/shm/` for online usage in which the client cleans up the files soon after they are generated.
58 +
59 +```bash
60 +vllm serve Qwen/Qwen3-8B \
61 + --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
64 +```
65 +
66 +## Configuration
67 +
68 +The `kv_connector_extra_config` dict accepts these options:
69 +
70 +| Parameter | Default | Description |
71 +| --- | --- | --- |
72 +| `shared_storage_path` | `/tmp` | Directory where hidden state files are saved |
73 +| `num_writer_threads` | `8` | Thread pool size for async disk writes |
74 +| `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. |
75 +
76 +## Output Format
77 +
78 +Each request produces a `.safetensors` file containing:
79 +
80 +- **`hidden_states`** — shape `[num_tokens, num_extracted_layers, hidden_size]`
81 +- **`token_ids`** — shape `[num_tokens]`
82 +
83 +The file path is returned in `output.kv_transfer_params["hidden_states_path"]`. Use `load_hidden_states()` from the connector module to read the file with proper synchronization.
84 +
85 +!!! note
86 + Chunked prefill is not compatible with this feature and must be disabled.
examples/features/speculative_decoding/extract_hidden_states_offline.py
+25 -16
@@ -2,9 +2,15 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import tempfile
4
5 -from safetensors import safe_open
6 -
5 from vllm import LLM, SamplingParams
6 +from vllm.config.kv_transfer import KVTransferConfig
7 +from vllm.distributed.kv_transfer.kv_connector.v1 import (
8 + example_hidden_states_connector,
9 +)
10 +
11 +# NOTE: If changing the interface of the ExampleHiddenStatesConnector, please also
12 +# update the benchmark in benchmarks/benchmark_hidden_state_extraction.py
13 +# and the docs in docs/features/speculative_decoding/extract_hidden_states.md
14
15 # Example: Using the custom "extract_hidden_states" speculator method and
16 # ExampleHiddenStatesConnector to extract and save hidden states from vllm
@@ -12,6 +18,7 @@ from vllm import LLM, SamplingParams
18 with tempfile.TemporaryDirectory() as tmpdirname:
19 llm = LLM(
20 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,
@@ -23,16 +30,16 @@ with tempfile.TemporaryDirectory() as tmpdirname:
30 3,
31 4,
32 ],
26 - }
33 + },
34 },
35 },
29 - kv_transfer_config={
30 - "kv_connector": "ExampleHiddenStatesConnector",
31 - "kv_role": "kv_producer",
32 - "kv_connector_extra_config": {
36 + kv_transfer_config=KVTransferConfig(
37 + kv_connector="ExampleHiddenStatesConnector",
38 + kv_role="kv_producer",
39 + kv_connector_extra_config={
40 "shared_storage_path": tmpdirname,
41 },
35 - },
42 + ),
43 )
44
45 prompts = ["Generate a sentence with hidden states", "Write a python function"]
@@ -47,12 +54,14 @@ with tempfile.TemporaryDirectory() as tmpdirname:
54 assert hidden_states_path is not None
55 print("Prompt hidden states path:", hidden_states_path)
56
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")
57 + obj = example_hidden_states_connector.load_hidden_states(hidden_states_path)
58 + token_ids = obj["token_ids"]
59 + hidden_states = obj["hidden_states"]
60 +
61 + print("Extracted token ids:", token_ids) # Matches prompt token ids
62 + print(
63 + "Extracted hidden states shape:", hidden_states.shape
64 + ) # [prompt_len, num_extracted_layers, hidden_size]
65 + print("Extracted hidden states:", hidden_states)
66
54 - print("Extracted token ids:", token_ids) # Matches prompt token ids
55 - print(
56 - "Extracted hidden states shape:", hidden_states.shape
57 - ) # [prompt len, num_hidden_layers, hidden size]
58 - print("Extracted hidden states:", hidden_states)
67 + example_hidden_states_connector.cleanup_hidden_states(hidden_states_path)
tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py
+2
@@ -120,6 +120,7 @@ def test_extract_hidden_states_with_predictable_dummy_model(
120 },
121 max_model_len=128,
122 enforce_eager=True,
123 + enable_chunked_prefill=False,
124 trust_remote_code=True,
125 load_format="dummy", # Don't try to load real weights
126 )
@@ -184,6 +185,7 @@ def test_extract_hidden_states_qwen35_hybrid_smoke(tmp_path):
185 },
186 max_model_len=256,
187 enforce_eager=True,
188 + enable_chunked_prefill=False,
189 gpu_memory_utilization=0.4,
190 load_format="dummy",
191 )
tests/v1/spec_decode/test_extract_hidden_states.py
+2
@@ -69,6 +69,7 @@ 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,
73 ),
74 attention_config=AttentionConfig(),
75 )
@@ -119,6 +120,7 @@ def test_proposer_initialization_missing_layer_ids():
120 scheduler_config=SchedulerConfig(
121 max_model_len=model_config.max_model_len,
122 is_encoder_decoder=model_config.is_encoder_decoder,
123 + enable_chunked_prefill=False,
124 ),
125 attention_config=AttentionConfig(),
126 )
vllm/config/vllm.py
+11
@@ -734,6 +734,17 @@ class VllmConfig:
734 Right now, this function reads the offloading settings from
735 CacheConfig and configures the KVTransferConfig accordingly.
736 """
737 + # Check if KV connector requires chunked prefill to be disabled.
738 + if (
739 + self.kv_transfer_config is not None
740 + and self.kv_transfer_config.kv_connector == "ExampleHiddenStatesConnector"
741 + and self.scheduler_config.enable_chunked_prefill
742 + ):
743 + raise ValueError(
744 + "ExampleHiddenStatesConnector does not support chunked prefill. "
745 + "Please disable chunked prefill (--no-enable-chunked-prefill)."
746 + )
747 +
748 # KV offloading is only activated when kv_offloading_size is set.
749 if (kv_offloading_size := self.cache_config.kv_offloading_size) is None:
750 return
vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py
+213 -39
@@ -1,11 +1,16 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import fcntl
4 import os
5 +from concurrent.futures import Future, ThreadPoolExecutor
6 from dataclasses import dataclass, field
7 +from functools import partial
8 +from importlib.metadata import version
9 from typing import TYPE_CHECKING, Any
10
7 -import safetensors
11 import torch
12 +from packaging.version import Version
13 +from safetensors.torch import load_file, save_file
14
15 from vllm.config import VllmConfig, get_layers_from_vllm_config
16 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
@@ -14,6 +19,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
19 KVConnectorRole,
20 SupportsHMA,
21 )
22 +from vllm.forward_context import get_forward_context
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
@@ -36,6 +42,39 @@ def extract_from_kv_cache(
42 return kv_cache[slot_mapping // block_size, slot_mapping % block_size][:num_tokens]
43
44
45 +def load_hidden_states(path: str) -> dict[str, torch.Tensor]:
46 + """Load hidden states written by ExampleHiddenStatesConnector.
47 +
48 + Blocks (without polling) until the async write is complete by
49 + acquiring a shared flock on the companion lock file. The kernel
50 + puts the caller to sleep until the writer releases its exclusive lock.
51 +
52 + Args:
53 + path: The file path returned in kv_transfer_params["hidden_states_path"].
54 +
55 + Returns:
56 + Dict with "hidden_states" and "token_ids" tensors.
57 + """
58 + lock_path = path + ".lock"
59 + with open(lock_path) as lf:
60 + fcntl.flock(lf, fcntl.LOCK_SH) # sleeps until writer releases LOCK_EX
61 + data = load_file(path, device="cpu")
62 + return data
63 +
64 +
65 +def cleanup_hidden_states(path: str, keep_hidden_states: bool = False) -> None:
66 + """Clean up hidden states file and lock file after loading.
67 +
68 + If keep_hidden_states is True, only removes the lock file
69 + and keeps the hidden states file.
70 + """
71 + lock_path = path + ".lock"
72 + if os.path.exists(lock_path):
73 + os.remove(lock_path)
74 + if not keep_hidden_states and os.path.exists(path):
75 + os.remove(path)
76 +
77 +
78 @dataclass
79 class ReqMeta:
80 # Request ID
@@ -112,6 +151,13 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
151 logger.info(self._kv_transfer_config)
152 logger.info("Shared storage path is %s", self._storage_path)
153
154 + if Version(version("safetensors")) < Version("0.8.0"):
155 + logger.warning(
156 + "safetensors < 0.8.0 holds the GIL during save_file, which "
157 + "serializes the writer thread pool and hurts throughput. "
158 + "Upgrade to safetensors >= 0.8.0 for better performance."
159 + )
160 +
161 assert self._vllm_config.speculative_config is not None, (
162 "ExampleHiddenStatesConnector only works when using "
163 "'extract_hidden_states' speculative method"
@@ -125,17 +171,97 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
171 self._active_requests: dict[str, NewRequestData] = {}
172 self._req_blocks: dict[str, list[int]] = {}
173
174 + # Async write infrastructure (worker-side).
175 + # Dedicated CUDA stream for DtoH copies so they don't block
176 + # the default stream (model forward). Thread pool for disk writes.
177 + self._copy_stream: torch.cuda.Stream | None = None # lazy init
178 + self._executor = ThreadPoolExecutor(
179 + max_workers=self._kv_transfer_config.get_from_extra_config(
180 + "num_writer_threads", 8
181 + ),
182 + thread_name_prefix="vllm-hs-save",
183 + )
184 + # Whether to use a filesystem lock when writing files to shared storage.
185 + # This is necessary for online transfer clients to avoid incomplete reads,
186 + # but can be disabled for offline tasks that run tasks in batches to completion
187 + self.use_lock = self._kv_transfer_config.get_from_extra_config(
188 + "use_synchronization_lock", True
189 + )
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 + ] = []
195 + # req_id → in-flight disk-write Future for that req_id.
196 + self._req_futures: dict[str, Future] = {}
197 + # req_id → CUDA event marking completion of the DtoH copy. Once
198 + # this event is complete the request is considered "done sending"
199 + # by get_finished; clients block on the per-file flock to wait for
200 + # the disk write itself.
201 + self._req_copy_events: dict[str, torch.cuda.Event] = {}
202 + # req_ids reported as finished-generating by the scheduler,
203 + # accumulated across get_finished calls.
204 + self._accumulated_finished_req_ids: set[str] = set()
205 +
206 + def _get_copy_stream(self) -> torch.cuda.Stream:
207 + """Lazily create the copy stream (CUDA must be initialized)."""
208 + if self._copy_stream is None:
209 + self._copy_stream = torch.cuda.Stream()
210 + return self._copy_stream
211 +
212 # ==============================
213 # Worker-side methods
214 # ==============================
215 def start_load_kv(self, *args, **kwargs: Any) -> None:
132 - pass # Empty implementation of abstract method
216 + pass # Store-only connector — nothing to load
217
218 def wait_for_layer_load(self, layer_name: str) -> None:
135 - pass # Empty implementation of abstract method
219 + pass # Store-only connector — nothing to load
220
221 def wait_for_save(self):
138 - pass # Empty implementation of abstract method
222 + """Submit pending async copies to the thread pool for disk write.
223 +
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.
234 + """
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()
257 +
258 + def _on_write_done(self, req_id: str, future: Future) -> None:
259 + """Surface any exception from the disk-write thread and drop the
260 + completed future from the in-flight tracking dict."""
261 + self._req_futures.pop(req_id, None)
262 + exc = future.exception()
263 + if exc is not None:
264 + logger.error("Hidden-states write failed for req_id=%s: %r", req_id, exc)
265
266 def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
267 from vllm.model_executor.models.extract_hidden_states import (
@@ -151,6 +277,26 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
277 f"Expected 1 CacheOnlyAttentionLayer, got {len(self.cache_layers)}"
278 )
279
280 + @staticmethod
281 + def _write_tensors(
282 + tensors: dict[str, torch.Tensor],
283 + event: torch.cuda.Event,
284 + filename: str,
285 + lock_fd: int | None,
286 + ) -> None:
287 + """Thread worker: wait for async DtoH copy, write to disk, release lock.
288 +
289 + ``lock_fd`` is an open file descriptor on the companion ``.lock``
290 + file with ``LOCK_EX`` already held. Closing it releases the lock,
291 + which unblocks any client sleeping on ``LOCK_SH``.
292 + """
293 + try:
294 + event.synchronize()
295 + save_file(tensors, filename)
296 + finally:
297 + if lock_fd is not None:
298 + os.close(lock_fd) # releases LOCK_EX
299 +
300 def save_kv_layer(
301 self,
302 layer_name: str,
@@ -161,6 +307,10 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
307 """Start saving the KV cache of the layer from vLLM's paged buffer
308 to the connector.
309
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.
313 +
314 Args:
315 layer_name (str): the name of the layer.
316 kv_layer (torch.Tensor): the paged KV buffer of the current
@@ -184,21 +334,46 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
334
335 os.makedirs(self._storage_path, exist_ok=True)
336
187 - slot_mapping = attn_metadata.slot_mapping
337 + copy_stream = self._get_copy_stream()
338 +
339 + # Ensure the copy stream sees all prior writes on the default stream.
340 + ready_event = torch.cuda.Event()
341 + ready_event.record()
342 + copy_stream.wait_event(ready_event)
343 +
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]
191 - req_slot_mapping = slot_mapping[offset : offset + num_tokens]
192 - offset += num_tokens
193 -
194 - hidden_states = extract_from_kv_cache(
195 - kv_layer, req_slot_mapping, num_tokens
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"
369 )
370 tensors = {
198 - "hidden_states": hidden_states.detach().cpu(),
199 - "token_ids": request.token_ids.detach().cpu(),
371 + "hidden_states": pinned_hs,
372 + "token_ids": request.token_ids.clone(),
373 }
201 - safetensors.torch.save_file(tensors, request.filename)
374 + self._pending_copies.append(
375 + (tensors, copy_done, request.filename, request.req_id)
376 + )
377
378 # ==============================
379 # Scheduler-side methods
@@ -258,31 +433,6 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
433 self._active_requests[new_req.req_id] = new_req
434 self._req_blocks[new_req.req_id] = list(new_req.block_ids[0])
435
261 - cached_reqs = scheduler_output.scheduled_cached_reqs
262 - for i, req_id in enumerate(cached_reqs.req_ids):
263 - if req_id not in self._active_requests:
264 - continue
265 -
266 - new_block_ids = cached_reqs.new_block_ids[i]
267 -
268 - cached_req = self._active_requests[req_id]
269 - req_block_ids = self._req_blocks[req_id]
270 -
271 - if new_block_ids is None:
272 - continue
273 -
274 - block_ids = new_block_ids[0]
275 -
276 - req_block_ids.extend(block_ids)
277 - filename = os.path.join(self._storage_path, f"{req_id}.safetensors")
278 -
279 - meta.add_request(
280 - req_id=req_id,
281 - filename=filename,
282 - token_ids=cached_req.prompt_token_ids or [],
283 - new_req=False,
284 - )
285 -
436 return meta
437
438 def request_finished(
@@ -309,7 +459,31 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA):
459 _ = self._active_requests.pop(req_id, None)
460 _ = self._req_blocks.pop(req_id, None)
461
312 - return False, {"hidden_states_path": req_filename}
462 + return True, {"hidden_states_path": req_filename}
463 +
464 + def get_finished(
465 + self, finished_req_ids: set[str]
466 + ) -> 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.
475 + """
476 + self._accumulated_finished_req_ids.update(finished_req_ids)
477 +
478 + done_sending: set[str] = set()
479 + for req_id in list(self._accumulated_finished_req_ids):
480 + event = self._req_copy_events.get(req_id)
481 + if event is None or event.query():
482 + self._req_copy_events.pop(req_id, None)
483 + done_sending.add(req_id)
484 + self._accumulated_finished_req_ids.discard(req_id)
485 +
486 + return done_sending or None, None
487
488 def request_finished_all_groups(
489 self,
vllm/v1/spec_decode/extract_hidden_states.py
+19 -12
@@ -12,8 +12,10 @@ from vllm.config import CUDAGraphMode, VllmConfig, get_layers_from_vllm_config
12 from vllm.forward_context import set_forward_context
13 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
14 from vllm.model_executor.model_loader import get_model
15 +from vllm.utils.platform_utils import is_pin_memory_available
16 from vllm.v1.attention.backend import AttentionMetadataBuilder, CommonAttentionMetadata
17 from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
18 +from vllm.v1.utils import CpuGpuBuffer
19 from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
20 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
21
@@ -50,6 +52,14 @@ class ExtractHiddenStatesProposer:
52 vllm_config.scheduler_config.max_num_batched_tokens + max_batch_size
53 )
54
55 + self.backup_next_token_ids = CpuGpuBuffer(
56 + max_batch_size,
57 + dtype=torch.int32,
58 + pin_memory=is_pin_memory_available(),
59 + device=device,
60 + with_numpy=True,
61 + )
62 +
63 self.hf_config = vllm_config.speculative_config.draft_model_config.hf_config
64 layer_ids = getattr(self.hf_config, "eagle_aux_hidden_state_layer_ids", None)
65 if not layer_ids:
@@ -303,18 +313,15 @@ class ExtractHiddenStatesProposer:
313 (if valid and not discarded) or a backup token from the request state.
314 """
315 num_reqs = gpu_input_batch.num_reqs
306 - device = sampled_token_ids.device
307 -
308 - # Compute backup tokens for discarded / invalid requests
309 - seq_lens_list = (gpu_input_batch.num_tokens_no_spec[:num_reqs] - 1).tolist()
310 - backup_tokens_gpu = torch.tensor(
311 - [
312 - requests[gpu_input_batch.req_ids[i]].get_token_id(seq_lens_list[i])
313 - for i in range(num_reqs)
314 - ],
315 - dtype=torch.int32,
316 - device=device,
317 - )
316 +
317 + # Precompute backup token IDs for discarded requests.
318 + num_reqs = gpu_input_batch.num_reqs
319 + for i in range(num_reqs):
320 + self.backup_next_token_ids.np[i] = requests[
321 + gpu_input_batch.req_ids[i]
322 + ].get_token_id(gpu_input_batch.num_tokens_no_spec[i] - 1)
323 + self.backup_next_token_ids.copy_to_gpu(num_reqs)
324 + backup_tokens_gpu = self.backup_next_token_ids.gpu[:num_reqs]
325
326 assert discard_request_mask.dtype == torch.bool
327