vLLM 隐状态提取系统

vLLM 已在 vllm>=v0.18.0（通过 PR #33736）中引入了原生的隐藏状态提取系统，以便于投机解码模型的训练和开发。该系统允许用户在不牺牲 vLLM 核心性能优化的前提下，提取 token 序列的内部中间表示。

解决隐藏状态提取瓶颈

提取隐藏状态对于投机解码中的草稿模型训练至关重要——在该场景中，小模型预测的 token 需要由更大的“验证器”模型进行确认——因为向草稿模型提供验证器的内部状态可以提升对齐度和质量。此前，研究者主要依赖两种次优方法：

  • 使用 transformers 库： 该方法会失去 vLLM 的分布式支持和性能优化，并且由于 transformers 与 vLLM 隐藏状态之间的不匹配，可能会引入 bug。
  • Patch vLLM 内部： 手动调用内部 API 会带来高维护成本，并且需要关闭前缀缓存、自动批处理和异步服务器等关键特性。

技术设计与实现

vLLM 隐状态提取系统的设计目标是在不增加标准推理“热路径”开销的情况下，管理大规模隐藏状态张量的显著内存需求。以隐藏维度为 4096 的 Qwen3-8B 模型为例，提取四层、长度为 8k token 的序列大约需要 268 MB 数据。

“Dummy Model” 架构

为在不引入额外运行时开销的前提下实现该功能，vLLM 利用了已有的基础设施：

  1. Eagle-3 管道： vLLM 已经具备将验证器模型的隐藏状态传递给草稿模型以进行 Eagle-3 投机解码的管道。
  2. Dummy Draft Model（虚拟草稿模型）： 系统创建一个虚拟草稿模型来接收这些隐藏状态。该模型不执行注意力计算，而是使用一个虚拟注意力层直接将隐藏状态写入自身的 KV 缓存。
  3. KV Connector API（KV 连接器 API）： vLLM 使用其可扩展的 KV Connector API——最初用于 Prefill/Decode Disaggregation 中的 KV 缓存提取——将虚拟草稿模型的 KV 缓存（此时已包含隐藏状态）保存到磁盘或转移到其他进程。

通过将隐藏状态视作 KV 缓存数据，vLLM 能够使用同一分页内存系统进行管理，确保与前缀缓存和分块预填充的兼容性。

使用方式与限制

用户可以通过 Python API 或以特定配置启动 vLLM 服务器来提取隐藏状态。服务器需要同时提供 --speculative_config（用于设置虚拟草稿模型）和 --kv_transfer_config（用于定义连接器）。

示例服务器命令：

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

关键约束条件：

  • 输出格式： 服务器返回一个 kv_transfer_params 字典，其中包含指向 .safetensors 文件的 hidden_states_path，该文件内保存 token_ids 与 hidden_states。
  • 范围： 仅保存提示 token 及其隐藏状态。团队建议使用 v1/completions 接口并将 max_tokens=1。
  • 并行性： 系统支持 --tensor-parallel-size 与 --data-parallel-size，可用于单节点多 GPU 部署。

未来路线图

  • Speculators 集成： speculators 库（v0.5.0）已通过 PR #353 更新，以使用此原生 vLLM 系统，实现草稿模型的在线训练。
  • 异步写入： 当前的 ExampleHiddenStatesConnector 使用阻塞写入；后续更新将实现异步写入以提升性能。
  • 设备间传输： 为摆脱基于磁盘的存储，vLLM 正在开发直接在设备之间传输隐藏状态的连接器，并支持多节点环境。

Sources