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>

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={

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 `<shared_storage_path>/<request_id>.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. |
 

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)

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 == (

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(),
     )

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

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: