从 vLLM 中提取隐藏状态

PR #33736（包含在 vllm>=v0.18.0 中）为 vLLM 引入了一套全新的隐藏状态（hidden states）提取系统。本文探讨了该功能的开发动机、设计思路、使用方式及未来方向，以及它在 vLLM 的 Speculators（一个用于创建和训练投机采样模型的库）中的应用。

动机

隐藏状态是模型对 Token 序列内部的中间表示。它们能够洞察模型的内部状态，并在投机采样中发挥着重要作用。

投机采样回顾

投机采样通常将“验证器”模型（即您试图部署的大型 LLM）与一个小型“草稿”模型相结合。草稿模型生成草稿 Token，验证器模型随后并行进行验证。这可以显著加快解码速度（根据不同方法可提升 2-5 倍），尤其是在小批量场景下，此时模型性能受限于内存带宽。

研究人员发现，向草稿模型提供来自验证器模型的内部隐藏状态，可以改善草稿的对齐效果和整体质量。因此，像 Eagle-3、P-Eagle、DFlash 等方法应运而生，这些方法需要将验证器模型多个层级的隐藏状态作为输入。

由于草稿模型需要将隐藏状态作为输入，训练它们需要访问海量的隐藏状态和验证器输出数据集。大多数投机采样库（如 Speculators）过去通过以下两种方式之一解决此问题：

1. 使用 transformers 库生成隐藏状态。这种方法可行，但有两个主要缺点：(A) 失去了 vLLM 的所有性能优化（如大模型支持、分布式支持等）。(B) 引入了一类潜在的错误，即 transformer 与 vLLM 生成的隐藏状态之间可能存在细微偏差。
2. 对 vLLM 进行深度修改和打补丁。这通常需要手动设置 vLLM 的核心组件并直接调用内部 API。随着 vLLM 内部架构的不断升级，这带来了巨大的维护压力。此外，这也意味着许多 vLLM 功能（如前缀缓存、自动批处理、异步服务等）必须被禁用。旧版本 Speculators（<0.5.0）就是这样处理隐藏状态生成的。

这两种方法各有利弊。随着投机采样的日益普及，业界迫切需要更高效的解决方案。

设计考量

在将隐藏状态提取功能直接集成到 vLLM 时，我们考虑了多项核心需求。

首先，该系统必须高效地返回隐藏状态。模型隐藏状态可能非常庞大。对于 hidden_size 为 4096 的 Qwen3-8B 模型，提取的隐藏状态形状为 [seq_len, num_layers_to_extract, 4096]。对于一个 8k 长度、4 层提取、FP16 精度的序列，数据量高达 268 MB。因此，序列化这些隐藏状态并将其直接包含在 HTTP 响应主体中是不现实的。

此外，由于隐藏状态占用空间巨大，即使是暂时存储在显存 (VRAM) 中也非易事。系统必须预先分配并管理所有并发请求的内存，同时还要处理分块预填充（Chunked Prefill）、请求抢占等问题，以防止显存溢出 (OOM)。

由于此功能仅在用户明确需要隐藏状态时才开启（在大多数部署场景下并非必需），因此确保该功能不会对 vLLM 的“热路径”性能造成任何运行负担或认知负担至关重要。在实践中，这意味着我们将修改范围限制在最小，并尽可能复用现有的功能。

最后，最终用户对隐藏状态的使用、存储和传输方式需求多样。例如，在“离线”投机模型训练中，需要为整个数据集生成隐藏状态并在训练开始前缓存到磁盘。而“在线”训练则需要在训练过程中动态生成隐藏状态，并需要高效地将这些状态传输到各个训练进程，最好无需先写入磁盘。为了支持这些场景，隐藏状态提取系统必须具备高度的灵活性和可扩展性。

设计见解

基于上述需求，我们通过以下核心设计思路实现了隐藏状态提取系统，总结如下：

1. vLLM 支持使用 Eagle-3（及类似）投机采样模型运行推理，它们使用验证器模型的隐藏状态作为输入。因此，vLLM 内部已经具备了将隐藏状态从验证器模型移动到草稿模型的底层链路。
2. vLLM 拥有一个可扩展的 KV Connector API，用于高效提取 KV 缓存数据，该 API 已用于预填充/解码分离（Prefill/Decode Disaggregation）等功能。现有的 API 实现支持通过 Nixl 传输 KV 缓存数据、写入磁盘、存储在共享内存中等。该 API 还支持异步传输 KV 缓存状态，并确保在传输完成前 KV 缓存块不会被释放。
3. 隐藏状态与 KV 缓存数据的映射逻辑相同：每个 Token 对应一个隐藏状态值，该值仅在其前缀序列上下文中有效。
4. vLLM 支持为投机采样草稿模型设置独立的 KV 缓存配置/大小。

结合上述理念（见图 1），我们可以通过以下方式提取隐藏状态：

1. 创建一个“虚拟草稿模型”，它利用现有的 Eagle-3 模型底层路径接收来自 vLLM 的验证器隐藏状态。
2. 此虚拟模型包含一个具备独立 KV 缓存的“虚拟注意力层”。它不运行实际的注意力机制，而是直接将隐藏状态输入插入到其 KV 缓存中。
3. 随后，自定义的 KV Connector 将虚拟草稿模型的 KV 缓存数据（现在存储了我们的隐藏状态）保存到磁盘或通过其他方式传输。

该设计完全满足所有需求：利用已有的 Eagle-3 通路将隐藏状态管道式传输至草稿模型，并通过 KV Connector API 提供了一种灵活、高效且可扩展的提取方法，以应对不同的下游应用。由于草稿模型在虚拟注意力层中存储隐藏状态，vLLM 会自动为其分配显存。同时，vLLM 使用与 KV 缓存相同的分页内存系统来管理这些隐藏状态，从而支持前缀缓存、分块预填充、高效批处理等功能。

使用方法与限制

examples/offline_inference/extract_hidden_states.py 展示了如何使用 Python API 提取隐藏状态。该系统也支持在 vLLM 服务器端使用，可通过以下命令启动：

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"
	}
}'

此命令设置了该系统的两个核心组件：--speculative_config 指示 vLLM 使用伪“extract_hidden_states”投机方法，以构建虚拟草稿模型，并支持指定提取哪些层的隐藏状态。第二个组件是 --kv_transfer_config，它配置了一个自定义的 KV Connector，专门用于提取草稿模型层的隐藏状态。目前，仅存在“ExampleHiddenStatesConnector”（一种简单的写盘实现），但后续会添加性能更高的连接器。请注意，这两个组件必须配套使用，系统才能正常工作。

一旦 vLLM 运行，对服务器的所有请求都会返回一个 "kv_transfer_params" 字典，其中包含 "hidden_states_path"。该路径指向一个包含隐藏状态和 Token ID 的 safetensors 文件。保存目录可通过上述配置中的 "shared_storage_path" 字段指定。

# `/tmp/hidden_states/{req_id}.safetensors`
{
	"token_ids": [prompt_seq_len],
	"hidden_states": [prompt_seq_len, num_hidden_layers, hidden_size]
}

备注

• 该功能支持 --tensor-parallel-size 和 --data-parallel-size 参数，适用于单机多 GPU 部署。
• 系统仅保存 Prompt Token 及其隐藏状态。因此，我们建议在调用 v1/completions 接口时，将采样参数设为 max_tokens=1。

后续工作

• 集成到 vLLM 的 speculators 项目中：speculators 库旨在高效训练投机采样算法。近期合并的 speculators PR #353 已更新为使用 vLLM 原生的隐藏状态提取系统，并启用了草稿模型的在线训练。此功能将包含在 speculators v0.5.0 中。
• 性能改进：目前的隐藏状态专用 KV Connector（"ExampleHiddenStatesConnector"）尚未经过深度优化，包含阻塞式写入。我们正积极研发该连接器的异步写入功能。
• 设备间连接器：当前的 ExampleHiddenStatesConnector 直接将隐藏状态写入磁盘，供后续训练进程读取。虽然简单易用且适合测试，但在大规模训练负载下扩展性较差。未来我们将开发更先进的隐藏状态连接器，支持直接在不同设备间（包括多节点环境）传输隐藏状态。