vLLM 隠れ状態抽出システム

vLLM は、投機的デコードモデルの訓練と開発を支援するために、ネイティブな隠れ状態抽出システム（vllm>=v0.18.0 に含まれ、PR #33736 で導入）を導入しました。このシステムにより、vLLM のコアパフォーマンス最適化を犠牲にすることなく、トークンシーケンスの内部中間表現を抽出できます。

隠れ状態抽出のボトルネック解消

隠れ状態の抽出は、投機的デコードにおけるドラフトモデルの訓練において重要です（小さなモデルがトークンを予測し、より大きな「検証者」モデルがそれを確認する）。検証者の内部状態をドラフトモデルに提供することで、整合性と品質が向上します。従来、研究者は以下の2つの最適でない方法に頼っていました。

  • transformers ライブラリの使用: この方法は vLLM の分散サポートとパフォーマンス最適化を失い、transformers と vLLM の隠れ状態の不一致によりバグが発生する可能性があります。
  • vLLM の内部をパッチする: 内部 API を手動で呼び出すことは保守負担が大きく、プレフィックスキャッシュ、オートバッチング、非同期サーバーなどの重要機能を無効にする必要があります。

技術設計と実装

vLLM の隠れ状態抽出システムは、標準推論の「ホットパス」へのオーバーヘッドを回避しつつ、大規模な隠れ状態テンソルの膨大なメモリ要件を管理するよう設計されています。隠れサイズが 4096 の Qwen3-8B モデルで、8k トークンのシーケンスに対して4層を抽出するには、約 268 MB のデータが必要です。

「ダミーモデル」アーキテクチャ

新たなランタイムオーバーヘッドを導入せずにこれを実装するため、vLLM は既存のインフラストラクチャを活用します。

  1. Eagle-3 配管: vLLM はすでに、検証者モデルからドラフトモデルへ隠れ状態を移動させる Eagle-3 投機的デコード用の配管を備えています。
  2. ダミードラフトモデル: システムはこれらの隠れ状態を受け取るダミードラフトモデルを作成します。注意演算を実行する代わりに、ダミー注意層を使用して隠れ状態を直接自身の KV キャッシュに挿入します。
  3. KV Connector API: vLLM は拡張可能な KV Connector API を利用します（もともと Prefill/Decode 分離における KV キャッシュ抽出用に設計）。これにより、ダミードラフトモデルの KV キャッシュ（現在は隠れ状態を含む）をディスクに保存したり、別プロセスへ転送したりできます。

隠れ状態を KV キャッシュデータとして扱うことで、vLLM は同じページングメモリシステムで管理でき、プレフィックスキャッシュやチャンク化されたプリフィルとの互換性が確保されます。

使用方法と制限事項

ユーザーは Python API を介して、または特定の設定で vLLM サーバーを起動することで隠れ状態を抽出できます。サーバーは、ダミードラフトモデルを設定するための --speculative_config と、コネクタを定義するための --kv_transfer_config の両方が必要です。

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:

  • 出力形式: サーバーは kv_transfer_params 辞書を返し、その中に .safetensors ファイルへのパスである hidden_states_path が含まれ、token_ids と hidden_states が格納されています。
  • 範囲: プロンプトトークンとその隠れ状態のみが保存されます。チームは max_tokens=1 の v1/completions エンドポイントの使用を推奨しています。
  • 並列性: システムはシングルノードのマルチGPU展開向けに --tensor-parallel-size と --data-parallel-size をサポートしています。

今後のロードマップ

  • Speculators の統合: speculators ライブラリ（v0.5.0）は PR #353 によりこのネイティブ vLLM システムを使用するよう更新され、ドラフトモデルのオンライン訓練が可能になりました。
  • 非同期書き込み: 現在の ExampleHiddenStatesConnector はブロッキング書き込みを使用していますが、将来的なアップデートでパフォーマンス向上のために非同期書き込みが実装される予定です。
  • デバイス間転送: ディスクベースのストレージを超えるため、vLLM はデバイス間で隠れ状態を直接転送するコネクタを開発中で、マルチノード環境への対応も含まれます。

Sources