From ebc6ef971a71b1a43ec728fae52237524b3056ca Mon Sep 17 00:00:00 2001 From: Fynn Schmitt-Ulms Date: Thu, 11 Jun 2026 09:44:45 -0400 Subject: [PATCH] Hidden states extraction improvements (#43805) Signed-off-by: Fynn Schmitt-Ulms Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../benchmark_hidden_state_extraction.py | 2 - .../extract_hidden_states.md | 50 ++- .../extract_hidden_states_offline.py | 29 +- .../test_extraction.py | 235 +++++++++--- .../spec_decode/test_extract_hidden_states.py | 2 - vllm/config/vllm.py | 11 - .../v1/example_hidden_states_connector.py | 358 ++++++++++-------- 7 files changed, 454 insertions(+), 233 deletions(-) diff --git a/benchmarks/benchmark_hidden_state_extraction.py b/benchmarks/benchmark_hidden_state_extraction.py index 6056fcdd07..f0a35a0cf1 100644 --- a/benchmarks/benchmark_hidden_state_extraction.py +++ b/benchmarks/benchmark_hidden_state_extraction.py @@ -92,7 +92,6 @@ def run_baseline( llm = LLM( model=model, enable_prefix_caching=False, - enable_chunked_prefill=False, **extra_args, ) sampling_params = SamplingParams(max_tokens=1) @@ -194,7 +193,6 @@ async def _run_extraction_async( engine_args = AsyncEngineArgs( model=model, enable_prefix_caching=False, - enable_chunked_prefill=False, max_num_batched_tokens=40960, max_model_len=40960, speculative_config={ diff --git a/docs/features/speculative_decoding/extract_hidden_states.md b/docs/features/speculative_decoding/extract_hidden_states.md index 2184a71f48..b7df376d9f 100644 --- a/docs/features/speculative_decoding/extract_hidden_states.md +++ b/docs/features/speculative_decoding/extract_hidden_states.md @@ -19,7 +19,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import ( with tempfile.TemporaryDirectory() as tmpdir: llm = LLM( model="Qwen/Qwen3-8B", - enable_chunked_prefill=False, speculative_config={ "method": "extract_hidden_states", "num_speculative_tokens": 1, @@ -59,17 +58,58 @@ For improved performance, it is recommended to use a RAM-mounted file system suc ```bash vllm serve Qwen/Qwen3-8B \ --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]}}}' \ - --kv_transfer_config '{"kv_connector": "ExampleHiddenStatesConnector", "kv_role": "kv_producer", "kv_connector_extra_config": {"shared_storage_path": "/dev/shm/hidden_states"}}' \ - --no-enable-chunked-prefill + --kv_transfer_config '{"kv_connector": "ExampleHiddenStatesConnector", "kv_role": "kv_producer", "kv_connector_extra_config": {"shared_storage_path": "/dev/shm/hidden_states"}}' +``` + +## Per-Request Options + +Both offline and online modes support per-request options via `kv_transfer_params`: + +| Parameter | Default | Description | +| --- | --- | --- | +| `hidden_states_path` | Auto-generated | Custom file path for saving hidden states. If not set, files are saved to `/.safetensors`. Requires `allow_custom_save_path` to be enabled in the server config. | +| `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. | + +### Offline usage + +Pass per-request options via `extra_args` on `SamplingParams`: + +```python +SamplingParams( + max_tokens=32, + extra_args={ + "kv_transfer_params": { + "hidden_states_path": "/tmp/my_output.safetensors", + "include_output_tokens": True, + } + }, +) +``` + +### Online usage + +Pass `kv_transfer_params` as a top-level field in the API request: + +```json +{ + "model": "Qwen/Qwen3-8B", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 32, + "kv_transfer_params": { + "hidden_states_path": "/tmp/my_output.safetensors", + "include_output_tokens": true + } +} ``` ## Configuration -The `kv_connector_extra_config` dict accepts these options: +The `kv_connector_extra_config` dict accepts these server-level options: | Parameter | Default | Description | | --- | --- | --- | -| `shared_storage_path` | `/tmp` | Directory where hidden state files are saved | +| `shared_storage_path` | `/tmp` | Directory where hidden state files are saved (used when `hidden_states_path` is not set per-request) | +| `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. | | `num_writer_threads` | `8` | Thread pool size for async disk writes | | `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. | diff --git a/examples/features/speculative_decoding/extract_hidden_states_offline.py b/examples/features/speculative_decoding/extract_hidden_states_offline.py index f8909566f4..5db315a043 100644 --- a/examples/features/speculative_decoding/extract_hidden_states_offline.py +++ b/examples/features/speculative_decoding/extract_hidden_states_offline.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os import tempfile from vllm import LLM, SamplingParams @@ -18,7 +19,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import ( with tempfile.TemporaryDirectory() as tmpdirname: llm = LLM( model="Qwen/Qwen3-8B", # Your target model - enable_chunked_prefill=False, # required speculative_config={ "method": "extract_hidden_states", "num_speculative_tokens": 1, @@ -38,13 +38,30 @@ with tempfile.TemporaryDirectory() as tmpdirname: kv_role="kv_producer", kv_connector_extra_config={ "shared_storage_path": tmpdirname, + "allow_custom_save_path": True, }, ), ) prompts = ["Generate a sentence with hidden states", "Write a python function"] - sampling_params = SamplingParams(max_tokens=1) - outputs = llm.generate(prompts, sampling_params) + + # One request uses defaults, the other uses a custom save path and + # includes output token hidden states via per-request kv_transfer_params. + sampling_params_list = [ + SamplingParams(max_tokens=1), + SamplingParams( + max_tokens=10, + extra_args={ + "kv_transfer_params": { + "hidden_states_path": os.path.join( + tmpdirname, "custom_output.safetensors" + ), + "include_output_tokens": True, + } + }, + ), + ] + outputs = llm.generate(prompts, sampling_params_list) for output in outputs: print("\nPrompt:", output.prompt) @@ -52,16 +69,16 @@ with tempfile.TemporaryDirectory() as tmpdirname: hidden_states_path = output.kv_transfer_params.get("hidden_states_path") assert hidden_states_path is not None - print("Prompt hidden states path:", hidden_states_path) + print("Hidden states path:", hidden_states_path) obj = example_hidden_states_connector.load_hidden_states(hidden_states_path) token_ids = obj["token_ids"] hidden_states = obj["hidden_states"] - print("Extracted token ids:", token_ids) # Matches prompt token ids + print("Extracted token ids:", token_ids) print( "Extracted hidden states shape:", hidden_states.shape - ) # [prompt_len, num_extracted_layers, hidden_size] + ) # [num_tokens, num_extracted_layers, hidden_size] print("Extracted hidden states:", hidden_states) example_hidden_states_connector.cleanup_hidden_states(hidden_states_path) diff --git a/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py b/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py index 5cc19247f5..390519fb55 100644 --- a/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py +++ b/tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py @@ -1,44 +1,40 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import gc import os +import tempfile import pytest import torch -from safetensors import safe_open +from tests.utils import create_new_process_for_each_test, multi_gpu_test from vllm import LLM, ModelRegistry, SamplingParams +from vllm.distributed.kv_transfer.kv_connector.v1 import ( + example_hidden_states_connector, +) def get_and_check_output(output, expected_shape): assert output.kv_transfer_params is not None hidden_states_path = output.kv_transfer_params.get("hidden_states_path") assert hidden_states_path is not None - assert os.path.exists(hidden_states_path) - # Load and verify the saved tensors - with safe_open(hidden_states_path, "pt") as f: - # Check that token_ids and hidden_states are present - tensor_names = f.keys() - assert "token_ids" in tensor_names - assert "hidden_states" in tensor_names + obj = example_hidden_states_connector.load_hidden_states(hidden_states_path) + token_ids = obj["token_ids"] + hidden_states = obj["hidden_states"] - token_ids = f.get_tensor("token_ids") - hidden_states = f.get_tensor("hidden_states") + prompt_token_ids = output.prompt_token_ids + assert torch.equal(token_ids, torch.tensor(prompt_token_ids)) - prompt_token_ids = output.prompt_token_ids - assert torch.equal(token_ids, torch.tensor(prompt_token_ids)) + assert hidden_states.shape == expected_shape - assert hidden_states.shape == expected_shape - - # Verify hidden_states are not all zeros (i.e., they were actually computed) - assert not torch.allclose(hidden_states, torch.zeros_like(hidden_states)) + # Verify hidden_states are not all zeros (i.e., they were actually computed) + assert not torch.allclose(hidden_states, torch.zeros_like(hidden_states)) return token_ids, hidden_states -@pytest.fixture(scope="module") +@pytest.fixture def predictable_llama_config_path(tmp_path_factory): """Create a minimal LlamaConfig for PredictableLlamaForCausalLM.""" from transformers import LlamaConfig, LlamaTokenizerFast @@ -53,7 +49,7 @@ def predictable_llama_config_path(tmp_path_factory): num_hidden_layers=24, # Enough layers to test various layer_ids num_attention_heads=4, num_key_value_heads=4, - max_position_embeddings=128, + max_position_embeddings=1024, architectures=["PredictableLlamaForCausalLM"], ) @@ -85,24 +81,25 @@ def register_predictable_model(): def test_extract_hidden_states_with_predictable_dummy_model( predictable_llama_config_path, tmp_path, monkeypatch ): - """Comprehensive test using a predictable dummy model with synthetic weights. + """Test hidden-state extraction with a predictable dummy model. - The PredictableLlamaForCausalLM outputs deterministic hidden states where - each layer produces values equal to (layer_index). This test verifies: - 1. Hidden states are correctly extracted from requested layers - 2. Values match the expected predictable pattern - 3. Layer ordering is preserved correctly (non-sequential layer IDs) - 4. Multiple prompts of different lengths produce consistent layer values + Tests 3 scenarios: + + 1. **Basic extraction**: non-sequential layer ordering, multiple prompts + of varying length — verifies correct layer association and + deterministic values. + 2. **Chunked prefill**: max_num_batched_tokens=128 with ~500-token + prompts so each is split across multiple scheduler iterations — + verifies hidden states are reassembled correctly. + 3. **Per-request options**: custom hidden_states_path and + include_output_tokens — verifies per-request kv_transfer_params + plumbing. """ - # Force fork so the engine worker inherits the autouse fixture's - # ModelRegistry.register_model("PredictableLlamaForCausalLM", ...). - # Spawn (the CI default) starts a fresh Python process that wouldn't - # see the registration. monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "fork") - # Test with non-sequential layer ordering to verify correct association layer_ids = [5, 2, 10] num_layers = len(layer_ids) + max_num_batched_tokens = 128 llm = LLM( model=predictable_llama_config_path, @@ -116,16 +113,21 @@ def test_extract_hidden_states_with_predictable_dummy_model( kv_transfer_config={ "kv_connector": "ExampleHiddenStatesConnector", "kv_role": "kv_producer", - "kv_connector_extra_config": {"shared_storage_path": tmp_path}, + "kv_connector_extra_config": { + "shared_storage_path": tmp_path, + "allow_custom_save_path": True, + }, }, - max_model_len=128, + max_model_len=1024, + max_num_batched_tokens=max_num_batched_tokens, enforce_eager=True, - enable_chunked_prefill=False, trust_remote_code=True, - load_format="dummy", # Don't try to load real weights + load_format="dummy", ) - # Test with multiple prompts of different lengths + hidden_size = llm.llm_engine.model_config.get_hidden_size() + + # --- Scenario 1: basic extraction with non-sequential layers ---------- prompts = [ "Short", "Medium length", @@ -133,15 +135,10 @@ def test_extract_hidden_states_with_predictable_dummy_model( "Much longer prompt with many tokens", # repeated prompt ] sampling_params = SamplingParams(max_tokens=1, temperature=0.0) - hidden_size = llm.llm_engine.model_config.get_hidden_size() outputs = llm.generate(prompts, sampling_params) - del llm - gc.collect() assert len(outputs) == len(prompts) - for output in outputs: - # hidden_states shape is [prompt_len, num_hidden_layers, hidden_size] expected_shape = ( len(output.prompt_token_ids), num_layers, @@ -156,12 +153,100 @@ def test_extract_hidden_states_with_predictable_dummy_model( torch.full_like(layer_hidden, layer_id), atol=1e-5, ), ( - f"Layer {layer_id} at position {idx} should output {float(layer_id)}, " - f"but got mean={layer_hidden.mean():.3f}, " - f"min={layer_hidden.min():.3f}, max={layer_hidden.max():.3f}" + f"Layer {layer_id} at position {idx} should output " + f"{float(layer_id)}, but got mean=" + f"{layer_hidden.mean():.3f}, min=" + f"{layer_hidden.min():.3f}, max={layer_hidden.max():.3f}" ) + # --- Scenario 2: chunked prefill with long prompts -------------------- + long_prompt = " ".join(["word"] * 500) + chunked_prompts = [ + long_prompt, + long_prompt + " extra tokens here", + "Short", + ] + outputs = llm.generate(chunked_prompts, sampling_params) + assert len(outputs) == len(chunked_prompts) + for output in outputs: + prompt_len = len(output.prompt_token_ids) + expected_shape = (prompt_len, num_layers, hidden_size) + _token_ids, hidden_states = get_and_check_output(output, expected_shape) + + for idx, layer_id in enumerate(layer_ids): + layer_hidden = hidden_states[:, idx, :] + assert torch.allclose( + layer_hidden, + torch.full_like(layer_hidden, layer_id), + atol=1e-5, + ), ( + f"Layer {layer_id} at position {idx} should output " + f"{float(layer_id)}, but got mean=" + f"{layer_hidden.mean():.3f}, min=" + f"{layer_hidden.min():.3f}, max=" + f"{layer_hidden.max():.3f}. " + f"prompt_len={prompt_len}, " + f"max_num_batched_tokens={max_num_batched_tokens}" + ) + + # --- Scenario 3: per-request options ---------------------------------- + max_tokens = 5 + custom_path = os.path.join(tmp_path, "subdir", "custom.safetensors") + + sampling_params_list = [ + SamplingParams(max_tokens=max_tokens, temperature=0.0), + SamplingParams( + max_tokens=max_tokens, + temperature=0.0, + extra_args={ + "kv_transfer_params": { + "hidden_states_path": custom_path, + "include_output_tokens": True, + } + }, + ), + ] + per_req_prompts = ["Short", "Medium length"] + outputs = llm.generate(per_req_prompts, sampling_params_list) + + # First output: prompt-only hidden states, default path + out0 = outputs[0] + path0 = out0.kv_transfer_params["hidden_states_path"] + assert path0 != custom_path + obj0 = example_hidden_states_connector.load_hidden_states(path0) + assert torch.equal(obj0["token_ids"], torch.tensor(out0.prompt_token_ids)) + assert obj0["hidden_states"].shape == ( + len(out0.prompt_token_ids), + num_layers, + hidden_size, + ) + example_hidden_states_connector.cleanup_hidden_states(path0) + + # Second output: prompt + output tokens, custom path + out1 = outputs[1] + assert out1.kv_transfer_params["hidden_states_path"] == custom_path + obj1 = example_hidden_states_connector.load_hidden_states(custom_path) + token_ids = obj1["token_ids"] + hidden_states = obj1["hidden_states"] + # The final output token was never an input to the model, so its hidden + # state is not in the cache — hence the -1. + total_tokens = len(out1.prompt_token_ids) + len(out1.outputs[0].token_ids) - 1 + assert token_ids.shape[0] == total_tokens + assert hidden_states.shape == (total_tokens, num_layers, hidden_size) + + # Verify predictable layer values hold for all tokens (prompt + output) + for idx, layer_id in enumerate(layer_ids): + layer_hidden = hidden_states[:, idx, :] + assert torch.allclose( + layer_hidden, + torch.full_like(layer_hidden, layer_id), + atol=1e-5, + ) + example_hidden_states_connector.cleanup_hidden_states(custom_path) + + +@create_new_process_for_each_test() def test_extract_hidden_states_qwen35_hybrid_smoke(tmp_path): """Smoke test for Qwen3.5 hybrid (mamba + full-attention) models. Uses load_format="dummy" to just check shape/plumbing. @@ -185,7 +270,6 @@ def test_extract_hidden_states_qwen35_hybrid_smoke(tmp_path): }, max_model_len=256, enforce_eager=True, - enable_chunked_prefill=False, gpu_memory_utilization=0.4, load_format="dummy", ) @@ -193,19 +277,68 @@ def test_extract_hidden_states_qwen35_hybrid_smoke(tmp_path): prompts = ["Hello world", "Test prompt with several tokens"] sampling_params = SamplingParams(max_tokens=1, temperature=0.0) outputs = llm.generate(prompts, sampling_params) - del llm - gc.collect() assert len(outputs) == len(prompts) for output in outputs: assert output.kv_transfer_params is not None hidden_states_path = output.kv_transfer_params.get("hidden_states_path") assert hidden_states_path is not None - assert os.path.exists(hidden_states_path) - with safe_open(hidden_states_path, "pt") as f: - token_ids = f.get_tensor("token_ids") - hidden_states = f.get_tensor("hidden_states") + obj = example_hidden_states_connector.load_hidden_states(hidden_states_path) + token_ids = obj["token_ids"] + hidden_states = obj["hidden_states"] + + assert torch.equal(token_ids, torch.tensor(output.prompt_token_ids)) + assert hidden_states.shape == ( + len(output.prompt_token_ids), + len(layer_ids), + hidden_size, + ) + + +@pytest.mark.timeout(60) +@multi_gpu_test(num_gpus=2) +@create_new_process_for_each_test() +def test_extract_hidden_states_tp2(): + """Test that hidden states extraction works with tensor_parallel_size=2.""" + tmp_dir = tempfile.mkdtemp() + layer_ids = [5, 11, 17] + hidden_size = 1024 # Qwen/Qwen3-0.6B hidden_size + + llm = LLM( + model="Qwen/Qwen3-0.6B", + tensor_parallel_size=2, + speculative_config={ + "method": "extract_hidden_states", + "num_speculative_tokens": 1, + "draft_model_config": { + "hf_config": {"eagle_aux_hidden_state_layer_ids": layer_ids} + }, + }, + kv_transfer_config={ + "kv_connector": "ExampleHiddenStatesConnector", + "kv_role": "kv_producer", + "kv_connector_extra_config": {"shared_storage_path": tmp_dir}, + }, + max_model_len=256, + enforce_eager=True, + gpu_memory_utilization=0.4, + load_format="dummy", + ) + + prompts = ["Hello world", "Test prompt with several tokens"] + sampling_params = SamplingParams(max_tokens=1, temperature=0.0) + outputs = llm.generate(prompts, sampling_params) + + assert len(outputs) == len(prompts) + for output in outputs: + assert output.kv_transfer_params is not None + hidden_states_path = output.kv_transfer_params.get("hidden_states_path") + assert hidden_states_path is not None + + obj = example_hidden_states_connector.load_hidden_states(hidden_states_path) + token_ids = obj["token_ids"] + hidden_states = obj["hidden_states"] assert torch.equal(token_ids, torch.tensor(output.prompt_token_ids)) assert hidden_states.shape == ( diff --git a/tests/v1/spec_decode/test_extract_hidden_states.py b/tests/v1/spec_decode/test_extract_hidden_states.py index b568d0b204..2a67257b09 100644 --- a/tests/v1/spec_decode/test_extract_hidden_states.py +++ b/tests/v1/spec_decode/test_extract_hidden_states.py @@ -69,7 +69,6 @@ def _create_proposer( scheduler_config=SchedulerConfig( max_model_len=model_config.max_model_len, is_encoder_decoder=model_config.is_encoder_decoder, - enable_chunked_prefill=False, ), attention_config=AttentionConfig(), ) @@ -120,7 +119,6 @@ def test_proposer_initialization_missing_layer_ids(): scheduler_config=SchedulerConfig( max_model_len=model_config.max_model_len, is_encoder_decoder=model_config.is_encoder_decoder, - enable_chunked_prefill=False, ), attention_config=AttentionConfig(), ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index a1a3420945..86a2f4d09e 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -756,17 +756,6 @@ class VllmConfig: Right now, this function reads the offloading settings from CacheConfig and configures the KVTransferConfig accordingly. """ - # Check if KV connector requires chunked prefill to be disabled. - if ( - self.kv_transfer_config is not None - and self.kv_transfer_config.kv_connector == "ExampleHiddenStatesConnector" - and self.scheduler_config.enable_chunked_prefill - ): - raise ValueError( - "ExampleHiddenStatesConnector does not support chunked prefill. " - "Please disable chunked prefill (--no-enable-chunked-prefill)." - ) - # KV offloading is only activated when kv_offloading_size is set. if (kv_offloading_size := self.cache_config.kv_offloading_size) is None: return diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py index 3e4e675085..696d3f7fb4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/example_hidden_states_connector.py @@ -19,10 +19,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorRole, SupportsHMA, ) -from vllm.forward_context import get_forward_context +from vllm.distributed.parallel_state import get_tensor_model_parallel_rank from vllm.logger import init_logger from vllm.v1.attention.backend import AttentionMetadata -from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput +from vllm.v1.core.sched.output import SchedulerOutput if TYPE_CHECKING: 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: @dataclass -class ReqMeta: - # Request ID +class PendingSave: req_id: str - # Request filename filename: str - # Request tokens token_ids: torch.Tensor - # Whether this request is a new request or partially computed already - new_req: bool - - @staticmethod - def make_meta( - req_id: str, - filename: str, - token_ids: list[int], - new_req: bool, - ) -> "ReqMeta": - return ReqMeta( - req_id=req_id, - filename=filename, - token_ids=torch.tensor(token_ids), - new_req=new_req, - ) + block_ids: list[int] @dataclass class ExampleHiddenStatesConnectorMetadata(KVConnectorMetadata): - requests: list[ReqMeta] = field(default_factory=list) - - def add_request( - self, - req_id: str, - filename: str, - token_ids: list[int], - new_req: bool = True, - ) -> None: - self.requests.append(ReqMeta.make_meta(req_id, filename, token_ids, new_req)) + pending_saves: list[PendingSave] = field(default_factory=list) + # req_id → filename for newly scheduled requests — the worker pre-creates + # lock files for these so the lock exists before the client receives the + # output path. + new_req_filenames: dict[str, str] = field(default_factory=dict) class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): @@ -167,9 +144,16 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): getattr(spec_config, "eagle_aux_hidden_state_layer_ids", []) ) + # Scheduler-side state + self._pending_saves: dict[str, PendingSave] = {} self._request_filenames: dict[str, str] = {} - self._active_requests: dict[str, NewRequestData] = {} - self._req_blocks: dict[str, list[int]] = {} + + # Worker-side state (set by register_kv_caches). + self._kv_cache: torch.Tensor | None = None + self._hs_group_idx: int = 0 + # Only TP rank 0 writes hidden states to disk; other TP ranks no-op. + # Set in register_kv_caches (after distributed init). + self._is_tp_rank_zero: bool = True # Async write infrastructure (worker-side). # Dedicated CUDA stream for DtoH copies so they don't block @@ -184,14 +168,23 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): # Whether to use a filesystem lock when writing files to shared storage. # This is necessary for online transfer clients to avoid incomplete reads, # but can be disabled for offline tasks that run tasks in batches to completion + self.allow_custom_save_path = self._kv_transfer_config.get_from_extra_config( + "allow_custom_save_path", False + ) + if self.allow_custom_save_path: + logger.warning( + "allow_custom_save_path is enabled. API clients can write " + "hidden states to arbitrary paths on the server filesystem. " + "Only enable this with trusted clients." + ) self.use_lock = self._kv_transfer_config.get_from_extra_config( "use_synchronization_lock", True ) - # (tensors_dict, copy_done_event, filename, req_id) queued by - # save_kv_layer, submitted to thread pool by wait_for_save. - self._pending_copies: list[ - tuple[dict[str, torch.Tensor], torch.cuda.Event, str, str] - ] = [] + # req_id → open fd on the .lock file with LOCK_EX held. + # Pre-created in wait_for_save when a request first arrives, + # consumed by _submit_async_write which passes the fd to the + # thread pool worker for release after writing. + self._lock_fds: dict[str, int] = {} # req_id → in-flight disk-write Future for that req_id. self._req_futures: dict[str, Future] = {} # req_id → CUDA event marking completion of the DtoH copy. Once @@ -218,42 +211,28 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): def wait_for_layer_load(self, layer_name: str) -> None: pass # Store-only connector — nothing to load - def wait_for_save(self): - """Submit pending async copies to the thread pool for disk write. + def wait_for_save(self) -> None: + """Pre-create lock files for newly arrived requests. - For each pending write we acquire an exclusive flock on a - companion ``.lock`` file **before** submitting to the thread pool. - The thread worker releases the lock after the data file is fully - written. Clients call :func:`load_hidden_states` which takes a - shared flock — the kernel sleeps the client until the writer is - done. Because ``wait_for_save`` runs before the worker returns - output to the scheduler, the lock file is guaranteed to exist - (and be held) by the time the client receives the path. - - The lock can be disabled via the "use_synchronization_lock" extra config. + This runs on the worker BEFORE the scheduler returns the output + path to the client, guaranteeing that the lock file exists (and + LOCK_EX is held) by the time the client tries to open it. """ - for tensors, event, filename, req_id in self._pending_copies: - prior = self._req_futures.get(req_id) - assert prior is None, "Found another KV transfer request with same req_id!" - - lock_fd = None - if self.use_lock: - # Create/open the lock file and acquire an exclusive lock. - # The lock is held by this fd; the thread worker will close - # the fd after writing, which releases the lock. - lock_path = filename + ".lock" - lock_fd = os.open( - lock_path, os.O_CREAT | os.O_WRONLY | os.O_TRUNC, 0o644 - ) - fcntl.flock(lock_fd, fcntl.LOCK_EX) - - future = self._executor.submit( - self._write_tensors, tensors, event, filename, lock_fd - ) - self._req_copy_events[req_id] = event - self._req_futures[req_id] = future - future.add_done_callback(partial(self._on_write_done, req_id)) - self._pending_copies.clear() + if not self._is_tp_rank_zero: + return + if not self.use_lock or not self.has_connector_metadata(): + return + metadata = self._get_connector_metadata() + if not isinstance(metadata, ExampleHiddenStatesConnectorMetadata): + return + for req_id, filename in metadata.new_req_filenames.items(): + if req_id in self._lock_fds: + continue + lock_path = filename + ".lock" + os.makedirs(os.path.dirname(lock_path), exist_ok=True) + lock_fd = os.open(lock_path, os.O_CREAT | os.O_WRONLY | os.O_TRUNC, 0o644) + fcntl.flock(lock_fd, fcntl.LOCK_EX) + self._lock_fds[req_id] = lock_fd def _on_write_done(self, req_id: str, future: Future) -> None: """Surface any exception from the disk-write thread and drop the @@ -264,6 +243,9 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): logger.error("Hidden-states write failed for req_id=%s: %r", req_id, exc) def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): + # Delay tp rank0 initialization until after distributed init + self._is_tp_rank_zero = get_tensor_model_parallel_rank() == 0 + from vllm.model_executor.models.extract_hidden_states import ( CacheOnlyAttentionLayer, ) @@ -276,6 +258,14 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): assert len(self.cache_layers) == 1, ( f"Expected 1 CacheOnlyAttentionLayer, got {len(self.cache_layers)}" ) + self._kv_cache = kv_caches[self.cache_layers[0]] + + # Find the KV cache group index for hidden states + if self._kv_cache_config is not None: + for i, group in enumerate(self._kv_cache_config.kv_cache_groups): + if self.cache_layers[0] in group.layer_names: + self._hs_group_idx = i + break @staticmethod def _write_tensors( @@ -304,35 +294,33 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): attn_metadata: AttentionMetadata, **kwargs: Any, ) -> None: - """Start saving the KV cache of the layer from vLLM's paged buffer - to the connector. + # Hidden states are already cached by CacheOnlyAttentionLayer during + # forward. Extraction happens in get_finished once all tokens are done. + pass - Launches an async DtoH copy on a dedicated CUDA stream. The - actual disk write is deferred to wait_for_save() which submits - it to a thread pool. + def _submit_async_write( + self, + pending: PendingSave, + ) -> None: + """Extract hidden states from KV cache and submit async DtoH + disk write. - Args: - layer_name (str): the name of the layer. - kv_layer (torch.Tensor): the paged KV buffer of the current - layer in vLLM. - attn_metadata (AttentionMetadata): the attention metadata. - **kwargs: additional arguments for the save operation. + Called from get_finished for each request that has finished generating. """ - if layer_name not in self.cache_layers: + if not self._is_tp_rank_zero: return + assert self._kv_cache is not None - from vllm.model_executor.models.extract_hidden_states import ( - CacheOnlyAttentionMetadata, + # Compute slot mapping from block_ids + block_ids_t = torch.tensor(pending.block_ids, dtype=torch.long) + num_blocks = block_ids_t.shape[0] + block_offsets = torch.arange(0, self._block_size, dtype=torch.long) + slot_mapping = ( + block_offsets.reshape((1, self._block_size)) + + block_ids_t.reshape((num_blocks, 1)) * self._block_size ) + slot_mapping = slot_mapping.flatten() - assert isinstance(attn_metadata, CacheOnlyAttentionMetadata), ( - "ExampleHiddenStatesConnector only supports CacheOnlyAttentionBackend" - ) - - connector_metadata = self._get_connector_metadata() - assert isinstance(connector_metadata, ExampleHiddenStatesConnectorMetadata) - - os.makedirs(self._storage_path, exist_ok=True) + num_tokens = pending.token_ids.shape[0] copy_stream = self._get_copy_stream() @@ -341,39 +329,56 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): ready_event.record() copy_stream.wait_event(ready_event) - slot_mapping = get_forward_context().slot_mapping[layer_name] # type: ignore - offset = 0 - for request in connector_metadata.requests: - num_tokens = request.token_ids.shape[0] - with torch.cuda.stream(copy_stream): - req_slot_mapping_gpu = slot_mapping[offset : offset + num_tokens] - assert req_slot_mapping_gpu.device == kv_layer.device - offset += num_tokens - - hidden_states_gpu = extract_from_kv_cache( - kv_layer, req_slot_mapping_gpu, num_tokens - ) - # Async DtoH copy into pinned host memory. - pinned_hs = torch.empty_like( - hidden_states_gpu, device="cpu", pin_memory=True - ) - pinned_hs.copy_(hidden_states_gpu, non_blocking=True) - - # Record completion of this copy on the copy stream. - copy_done = torch.cuda.Event() - copy_done.record(copy_stream) - - # token_ids is already on CPU (created in ReqMeta.make_meta). - assert not request.token_ids.is_cuda, ( - "Expected token_ids on CPU, got CUDA tensor" + with torch.cuda.stream(copy_stream): + # Move the CPU slot_mapping to GPU on the copy stream so the + # implicit H2D inside fancy indexing doesn't sync the default + # stream. + slot_mapping_gpu = slot_mapping.to( + device=self._kv_cache.device, non_blocking=True ) - tensors = { - "hidden_states": pinned_hs, - "token_ids": request.token_ids.clone(), - } - self._pending_copies.append( - (tensors, copy_done, request.filename, request.req_id) + hidden_states_gpu = extract_from_kv_cache( + self._kv_cache, slot_mapping_gpu, num_tokens ) + # Async DtoH copy into pinned host memory. + pinned_hs = torch.empty_like( + hidden_states_gpu, device="cpu", pin_memory=True + ) + pinned_hs.copy_(hidden_states_gpu, non_blocking=True) + + # Record completion of this copy on the copy stream. + copy_done = torch.cuda.Event() + copy_done.record(copy_stream) + + # token_ids is already on CPU (created in request_finished). + assert not pending.token_ids.is_cuda, ( + "Expected token_ids on CPU, got CUDA tensor" + ) + tensors = { + "hidden_states": pinned_hs, + "token_ids": pending.token_ids.clone(), + } + + # Submit to thread pool for disk write. + prior = self._req_futures.get(pending.req_id) + assert prior is None, "Found another KV transfer request with same req_id!" + + os.makedirs(os.path.dirname(pending.filename), exist_ok=True) + + # Use the pre-created lock fd from wait_for_save (already holds + # LOCK_EX). Falls back to creating one here if use_lock is True + # but no pre-created fd exists (shouldn't happen in normal flow). + lock_fd = self._lock_fds.pop(pending.req_id, None) + if lock_fd is None and self.use_lock: + lock_path = pending.filename + ".lock" + lock_fd = os.open(lock_path, os.O_CREAT | os.O_WRONLY | os.O_TRUNC, 0o644) + fcntl.flock(lock_fd, fcntl.LOCK_EX) + + future = self._executor.submit( + self._write_tensors, tensors, copy_done, pending.filename, lock_fd + ) + self._req_copy_events[pending.req_id] = copy_done + self._req_futures[pending.req_id] = future + future.add_done_callback(partial(self._on_write_done, pending.req_id)) # ============================== # Scheduler-side methods @@ -421,17 +426,34 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): scheduler_output (SchedulerOutput): the scheduler output object. """ meta = ExampleHiddenStatesConnectorMetadata() + + # Transfer pending saves into metadata (scheduler → worker bridge) + meta.pending_saves = list(self._pending_saves.values()) + self._pending_saves.clear() + + # Resolve save paths for new requests and tell the worker so it can + # pre-create lock files before the client receives the output path. for new_req in scheduler_output.scheduled_new_reqs: - token_ids = new_req.prompt_token_ids or [] - filename = os.path.join(self._storage_path, f"{new_req.req_id}.safetensors") - meta.add_request( - new_req.req_id, - filename=filename, - token_ids=token_ids, + default_path = os.path.join( + self._storage_path, f"{new_req.req_id}.safetensors" ) + kv_params = ( + new_req.sampling_params.extra_args.get("kv_transfer_params") + if new_req.sampling_params and new_req.sampling_params.extra_args + else None + ) or {} + custom_path = kv_params.get("hidden_states_path") + if custom_path is not None and not self.allow_custom_save_path: + logger.warning( + "Request %s provided hidden_states_path but " + "allow_custom_save_path is disabled. Ignoring " + "custom path and using default.", + new_req.req_id, + ) + custom_path = None + filename = custom_path or default_path self._request_filenames[new_req.req_id] = filename - self._active_requests[new_req.req_id] = new_req - self._req_blocks[new_req.req_id] = list(new_req.block_ids[0]) + meta.new_req_filenames[new_req.req_id] = filename return meta @@ -444,35 +466,54 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): Called exactly once when a request has finished, before its blocks are freed. - The connector may assumes responsibility for freeing the blocks - asynchronously by returning True. - - Returns: - True if the request is being saved/sent asynchronously and blocks - should not be freed until the request_id is returned from - get_finished(). - Optional KVTransferParams to be included in the request outputs - returned by the engine. + Returns True to delay block freeing until get_finished extracts + the hidden states from the KV cache. """ req_id = request.request_id - req_filename = self._request_filenames.pop(req_id, None) - _ = self._active_requests.pop(req_id, None) - _ = self._req_blocks.pop(req_id, None) - - return True, {"hidden_states_path": req_filename} + filename = self._request_filenames.pop(req_id) + kv_params = request.kv_transfer_params or {} + if kv_params.get("include_output_tokens", False): + # Exclude the final token — it was the model's output, never an + # input to a forward pass, so its hidden state is not in the cache. + token_ids = torch.tensor(list(request.all_token_ids)[:-1]) + elif request.prompt_token_ids is not None: + token_ids = torch.tensor(request.prompt_token_ids) + else: + logger.warning( + "Request %s has no prompt_token_ids (prompt_embeds only). " + "Saved token_ids will be empty.", + req_id, + ) + token_ids = torch.tensor([], dtype=torch.long) + self._pending_saves[req_id] = PendingSave( + req_id=req_id, + filename=filename, + token_ids=token_ids, + block_ids=list(block_ids), + ) + return True, {"hidden_states_path": filename} def get_finished( self, finished_req_ids: set[str] ) -> tuple[set[str] | None, set[str] | None]: - """Poll DtoH-copy completion for requests that finished generating. + """Extract hidden states and poll DtoH-copy completion. - The scheduler passes finished_req_ids to tell the worker which - requests are done generating. We accumulate these across calls - and return a request as "finished sending" once its DtoH copy - event is complete (or if it never had a pending copy). The - subsequent disk write may still be in flight; clients block on - the per-file flock to wait for it. + On the worker side, connector metadata carries pending saves from the + scheduler. For each one we extract from the KV cache and launch an + async DtoH copy + thread-pool disk write. + + We then poll accumulated finished req_ids: a request is "done sending" + once its DtoH copy event is complete. The subsequent disk write may + still be in flight; clients block on the per-file flock to wait for it. """ + # Extract and submit async writes for newly finished requests + if self.has_connector_metadata(): + connector_metadata = self._get_connector_metadata() + if isinstance(connector_metadata, ExampleHiddenStatesConnectorMetadata): + for pending in connector_metadata.pending_saves: + self._submit_async_write(pending) + + # Poll for completed DtoH copies self._accumulated_finished_req_ids.update(finished_req_ids) done_sending: set[str] = set() @@ -482,6 +523,11 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): self._req_copy_events.pop(req_id, None) done_sending.add(req_id) self._accumulated_finished_req_ids.discard(req_id) + # Clean up any leftover lock fds (e.g. aborted requests + # that never went through _submit_async_write). + lock_fd = self._lock_fds.pop(req_id, None) + if lock_fd is not None: + os.close(lock_fd) return done_sending or None, None @@ -490,7 +536,7 @@ class ExampleHiddenStatesConnector(KVConnectorBase_V1, SupportsHMA): request: "Request", block_ids: tuple[list[int], ...], ) -> tuple[bool, dict[str, Any] | None]: - return self.request_finished(request, block_ids[0]) + return self.request_finished(request, block_ids[self._hs_group_idx]) @classmethod def get_required_kvcache_layout(cls, vllm_config: "VllmConfig") -> str | None: