vLLM Hidden States Extraction System

vLLM Hidden States Extraction System

vLLM has introduced a native hidden states extraction system (included in vllm>=v0.18.0 via PR #33736) to facilitate the training and development of speculative decoding models. This system allows users to extract internal intermediate representations of token sequences without sacrificing vLLM's core performance optimizations.

Solving the Hidden States Extraction Bottleneck

Extracting hidden states is critical for training draft models in speculative decoding—where a small model predicts tokens that a larger "verifier" model confirms—because providing draft models with the verifier's internal states can improve alignment and quality. Previously, researchers relied on two suboptimal methods:

  • Using the transformers library: This approach loses vLLM's distributed support and performance optimizations and can introduce bugs due to mismatches between transformers and vLLM hidden states.
  • Patching vLLM internals: Manually calling internal APIs creates a high maintenance burden and requires disabling key features like prefix caching, auto-batching, and the async server.

Technical Design and Implementation

The vLLM hidden states extraction system is designed to avoid overhead on the "hot path" of standard inference while managing the significant memory requirements of large hidden state tensors. For a Qwen3-8B model with a hidden size of 4096, extracting four layers for a sequence of 8k tokens requires approximately 268 MB of data.

The "Dummy Model" Architecture

To implement this without introducing new runtime overhead, vLLM leverages existing infrastructure:

  1. Eagle-3 Plumbing: vLLM already possesses the plumbing to move hidden states from a verifier model to a draft model for Eagle-3 speculative decoding.
  2. Dummy Draft Model: The system creates a dummy draft model that receives these hidden states. Instead of performing attention operations, this model uses a dummy attention layer to insert the hidden states directly into its own KV cache.
  3. KV Connector API: vLLM utilizes its extensible KV Connector API—originally designed for KV cache extraction in Prefill/Decode Disaggregation—to save the dummy draft model's KV cache (now containing the hidden states) to disk or transfer it to another process.

By treating hidden states as KV cache data, vLLM can manage them using the same paged memory system, ensuring compatibility with prefix caching and chunked prefill.

Usage and Limitations

Users can extract hidden states via the Python API or by launching the vLLM server with specific configurations. The server requires both a --speculative_config to set up the dummy draft model and a --kv_transfer_config to define the connector.

Example Server Command:

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": [3, 18, 33, 36]
		}
	}
}' --kv_transfer_config '{
	"kv_connector": "ExampleHiddenStatesConnector", 
	"kv_role": "kv_producer", 
	"kv_connector_extra_config": {
		"shared_storage_path": "/tmp/hidden_states"
	}
}'

Key Constraints:

  • Output Format: The server returns a kv_transfer_params dictionary containing a hidden_states_path pointing to a .safetensors file containing token_ids and hidden_states.
  • **Scope: Only prompt tokens and their hidden states are saved. The team recommends using the v1/completions endpoint with max_tokens=1.
  • Parallelism: The system supports --tensor-parallel-size and --data-parallel-size for single-node multi-GPU deployments.

Future Roadmap

  • Speculators Integration: The speculators library (v0.5.0) has been updated via PR #353 to use this native vLLM system, enabling online training of draft models.
  • Asynchronous Writes: The current ExampleHiddenStatesConnector uses blocking writes; future updates will implement asynchronous writes to improve performance.
  • Device-to-Device Transfer: To move beyond disk-based storage, vLLM is developing connectors that transfer hidden states directly between devices, including support for multi-node environments.

Sources