菜单
主题
vLLM Logo

从 vLLM 中提取隐藏状态

8 分钟阅读
Fynn Schmitt-Ulms
#投机采样（speculative-decoding）

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

动机

隐藏状态是模型对 Token 序列的内部中间表征。它们能深入洞察模型的内部状态，并在投机采样中被大量使用。

投机采样回顾

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

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

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

  1. 使用 transformers 进行隐藏状态生成。这种方法可行，但有两个主要缺点：(A) 丢失了 vLLM 的所有性能优化（如大模型/分布式支持等）。(B) 引入了由 transformers 和 vLLM 隐藏状态之间微小差异导致的潜在 Bug。
  2. 对 vLLM 进行深度修改和补丁（Patching）。这通常需要手动配置 vLLM 的核心组件并直接调用内部 API。随着 vLLM 内部结构的不断更新，这导致了巨大的维护负担。同时，这意味着必须禁用许多 vLLM 特性（如前缀缓存、自动批处理、异步服务器等）。以前版本的 Speculators（<0.5.0）就是这样处理隐藏状态生成的。

这两种方法都有各自的劣势，随着投机采样变得越来越流行，需要更优秀、性能更高的解决方案。

设计考量

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

首先，系统应以高性能方式返回隐藏状态。模型隐藏状态可能非常大。对于具有 4096 hidden_size 的 Qwen3-8B 模型，提取的隐藏状态形状为 [seq_len, num_layers_to_extract, 4096]。对于 8k Token 的序列、4 层且使用 FP16 精度，数据总量高达 268 MB。因此，序列化隐藏状态并将其直接放入请求的响应体中是不切实际的。

由于隐藏状态占用空间巨大，即使只是在 VRAM 中临时存储它们也不是一件简单的事。必须为所有并发请求预分配并管理内存，包括处理分块预填充（chunked prefill）、请求抢占等，以避免 OOM 错误。

由于此功能仅在用户需要从 vLLM 获取隐藏状态时才适用（这在大多数部署场景中并非必要），因此必须确保该功能不会对 vLLM 的“热路径”带来任何新的开销（运行时或认知上的）。在实践中，这意味着限制变更范围，并在尽可能多的地方重用现有特性。

最后，最终用户对隐藏状态的使用、存储或传输有多种需求。例如，在“离线”Speculator 训练中，隐藏状态是为整个数据集生成的，并在训练开始前缓存到磁盘。而“在线”训练则是在训练过程中实时生成隐藏状态，并需要将隐藏状态高效地传输到每个训练进程，最好无需先写入磁盘。为了支持这些不同场景，隐藏状态提取系统需要具备灵活性和可扩展性。

设计洞察

基于上述要求，通过若干设计洞察，我们实现了隐藏状态提取系统。这些洞察总结如下。

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

结合这些思路（图 1），我们可以通过以下方式提取隐藏状态：

  1. 创建一个虚拟草稿模型，该模型通过现有的 Eagle-3 模型链路接收来自 vLLM 的验证器隐藏状态。
  2. 此虚拟模型拥有一个带有自有 KV 缓存的虚拟注意力层。虚拟模型不执行注意力计算，而是直接将隐藏状态输入插入其 KV 缓存。
  3. 然后，自定义的 KV 连接器将虚拟草稿模型的 KV 缓存数据（现在存储了我们的隐藏状态）保存到磁盘，或以其他方式传输。

这满足了所有设计需求：利用现有的 Eagle-3 路径将隐藏状态输送至草稿模型，并提供了一种高效的隐藏状态提取方法，该方法通过 KV 连接器 API 足够灵活，可以处理各种下游使用场景。由于草稿模型将隐藏状态存储在虚拟注意力层中，vLLM 能够为其分配 VRAM。此外，vLLM 使用与 KV 缓存相同的分页内存系统来管理隐藏状态，这支持了前缀缓存、分块预填充、高效批处理等功能。

Figure 1: Design diagram for hidden states extraction system.
图 1：隐藏状态提取系统的设计图。

使用与限制

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 连接器，该连接器旨在仅从草稿模型层提取隐藏状态。在撰写本文时，仅存在“ExampleHiddenStatesConnector”（一种写入磁盘的简单实现），但很快会添加性能更高的连接器。请注意，为了使系统正常工作，必须同时使用这两个组件。

一旦 vLLM 运行，对服务器的任何请求都将返回一个包含 "hidden_states_path" 的 "kv_transfer_params" 字典。该路径指向一个包含隐藏状态和 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]
}

备注

  • 这适用于单节点多 GPU 部署的 --tensor-parallel-size 和 --data-parallel-size 参数。
  • 仅提示词（Prompt）Token 及其隐藏状态会被保存。因此，我们建议在调用 v1/completions 端点时使用 max_tokens=1 的采样参数。

后续工作

  • 集成到 vLLM 的 speculators 项目：speculators 库专为投机采样算法的高效训练而设计。最近合并的 speculators PR #353 更新了 speculators，使其使用新的 vLLM 原生隐藏状态提取系统，并启用了草稿模型的在线训练。此功能将包含在 speculators v0.5.0 中。
  • 性能改进：隐藏状态专用 KV 连接器（"ExampleHiddenStatesConnector"）的初始实现尚未优化，并且包含阻塞式的隐藏状态写入操作。目前正在积极努力在该连接器中启用异步写入。
  • 设备到设备连接器：ExampleHiddenStatesConnector 将隐藏状态直接写入磁盘，随后可被训练进程使用。这种方法简单且提供了一个很好的测试实现，但在较大的训练负载下扩展性不佳。后续工作将包括开发更先进的隐藏状态连接器，直接在设备之间（包括多节点环境）传输隐藏状态。