vLLM 隱藏狀態提取系統

vLLM 已引入原生的隱藏狀態提取系統（透過 PR #33736 包含於 vllm>=v0.18.0），以促進投機解碼模型的訓練與開發。此系統允許使用者在不犧牲 vLLM 核心效能最佳化的前提下，提取 token 序列的內部中間表示。

解決隱藏狀態提取瓶頸

提取隱藏狀態對於投機解碼中訓練草稿模型至關重要——在此情境下，小模型預測 token，較大的「驗證模型」進行確認——因為向草稿模型提供驗證模型的內部狀態可以提升對齊度與品質。過去，研究人員依賴兩種次佳方法：

  • 使用 transformers 套件： 此方法會失去 vLLM 的分散式支援與效能最佳化，且因 transformers 與 vLLM 隱藏狀態不匹配，可能產生錯誤。
  • 修補 vLLM 內部： 手動呼叫內部 API 會帶來高維護負擔，且需要停用前綴快取、自動批次處理以及非同步伺服器等關鍵功能。

技術設計與實作

vLLM 隱藏狀態提取系統的設計旨在避免在標準推論的「熱路徑」上產生額外負擔，同時管理大型隱藏狀態張量的龐大記憶體需求。以 hidden size 為 4096 的 Qwen3-8B 模型為例，提取四層、長度為 8k token 的序列大約需要 268 MB 的資料。

「Dummy Model」架構

為了在不引入新執行時開銷的情況下實作此功能，vLLM 利用既有基礎設施：

  1. Eagle-3 Plumbing： vLLM 已具備將隱藏狀態從驗證模型傳遞至草稿模型以進行 Eagle-3 投機解碼的管線。
  2. Dummy Draft Model： 系統會建立一個 dummy 草稿模型來接收這些隱藏狀態。該模型不執行注意力運算，而是使用 dummy 注意力層直接將隱藏狀態插入自身的 KV 快取。
  3. KV Connector API： vLLM 使用其可擴充的 KV Connector API——最初設計用於 Prefill/Decode Disaggregation 中的 KV 快取提取——將 dummy 草稿模型的 KV 快取（現在包含隱藏狀態）儲存至磁碟或傳輸至其他程序。

透過將隱藏狀態視為 KV 快取資料，vLLM 能以相同的分頁記憶體系統管理它們，確保與前綴快取與分塊預填的相容性。

使用方式與限制

使用者可以透過 Python API 或以特定設定啟動 vLLM 伺服器來提取隱藏狀態。伺服器同時需要 --speculative_config 以設定 dummy 草稿模型，與 --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"
	}
}'

Key Constraints:

  • 輸出格式： 伺服器回傳一個 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