8月7日、Aleksa Gordicが「Inside vLLM: Anatomy of a High-Throughput LLM Inference System」と題した記事を公開した。この記事では、高スループットLLM推論システムであるvLLMの内部アーキテクチャを、スケジューリングからマルチGPU実行まで体系的に解説している。
LLM推論の高速化に取り組むエンジニアにとって、vLLMは現時点で最も広く使われているオープンソースの推論エンジンだ。2023年にUC Berkeleyから発表されて以来、PagedAttentionや連続バッチ処理(Continuous Batching)といった技術を中心に急速に普及した。本記事はそのvLLMのV1エンジン(commit 42172ad)を対象に、内部構造を詳細に解説したものだ。(※当該commitは元記事執筆時点のスナップショットであり、現在のmainブランチとは差異がある場合がある)
V0からV1への移行にあたっては、単なるリファクタリングにとどまらない設計上の変化がある。最も大きな点のひとつが、PrefillとDecodeを同一ステップ内で混在処理できるようになったことだ。後述するPagedAttentionやChunked Prefillといった機能と組み合わせることで、スループットとレイテンシの両立が従来より現実的になっている。
PagedAttentionの仕組み
vLLMのスループットを支える中核技術がPagedAttentionだ。PagedAttentionの提案論文(Kwon et al., 2023)に基づく手法で、OSの仮想メモリ管理(ページング)に着想を得て、KVキャッシュをブロック単位で管理する。
デフォルトのブロックサイズは16トークン。17トークンのPrefillリクエストが来た場合、ceil(17/16) = 2ブロックが割り当てられる。ブロック割り当て(allocate_slots)の手順は次の通りだ。
- 必要なブロック数を計算する
- プールに空きがなければ、低優先度リクエストをプリエンプション(再計算による退避)して確保する
free_block_queueから先頭nブロックを取得し、req_to_blocks(リクエストIDとブロックリストのマッピング)に格納する
フォワードパスでは、全シーケンスが1本の「スーパーシーケンス」にフラット化・連結されて処理される。ポジションインデックスとアテンションマスクにより、各シーケンスは自身のトークンのみにアテンションを向ける仕組みになっている。これがパディングなしでの連続バッチ処理を可能にしている。
高度な機能群
基本的なエンジンフローの上に、以下の高度な機能が実装されている。
Chunked Prefill(チャンク分割Prefill)
長いプロンプトのPrefillをチャンク単位に分割して処理する機能。Prefillがデコード処理を長時間ブロックする問題(TTFT:Time To First Token、最初のトークンが返るまでの時間、の悪化)を緩和する。
Prefix Caching(プレフィックスキャッシング)
同じプロンプト冒頭を持つリクエストで、KVキャッシュを再利用する機能。システムプロンプトを共有するシナリオで大幅な計算削減につながる。
Speculative Decoding(投機的デコード)
小さなドラフトモデルで複数トークンを先読みし、大きなメインモデルで並列検証する手法。メモリ帯域バウンドなDecodeフェーズのスループット改善に有効だ。
Guided Decoding(ガイドデコード)
JSONスキーマや正規表現に従った出力を強制する機能。構造化出力マネージャが担当する。
Disaggregated P/D(Prefill/Decode分離)
PrefillとDecodeを異なるノードに分離して処理するアーキテクチャ。大規模サービングでのリソース最適化を狙う。
エンジンの中核:スケジューラとKVキャッシュ管理
vLLMのLLMエンジンは「オフライン推論」の基盤となるコンポーネントだ。以下のシンプルなコードがその出発点として示されている。
from vllm import LLM, SamplingParams
prompts = [
"Hello, my name is",
"The president of the United States is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
def main():
llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
outputs = llm.generate(prompts, sampling_params)
if __name__ == "__main__":
main()
エンジンの主要コンポーネントは以下の通りだ。
- Processor:生の入力をトークナイズし
EngineCoreRequestに変換する - Engine Core:スケジューラ・モデル実行器・構造化出力マネージャを束ねるコア
- Output Processor:内部出力を
RequestOutputに変換してユーザーに返す
エンジンコアの中でも特に重要なのがスケジューラだ。スケジューラは内部にKVキャッシュマネージャを持ち、free_block_queue(空きKVキャッシュブロックのプール)を管理している。VRAMのサイズとブロックサイズに応じて、このプールは数十万ブロック規模になることもある。
推論の各ステップで、スケジューラは2種類のリクエストを処理する。
- Prefillリクエスト:プロンプト全トークンに対するフォワードパス。計算バウンド。
- Decodeリクエスト:直近の1トークンのみに対するフォワードパス。KVキャッシュが活用される分、メモリ帯域バウンドになる。
V1エンジンでは、同一ステップ内でPrefillとDecodeを混在させて処理できる。V0エンジンはどちらか一方しか処理できなかったため、これは大きな改善点だ。
マルチGPU・サービング層
エンジンは単一GPUのUniProcExecutorから、複数GPUに対応するMultiProcExecutorへとスケールアウトできる。並列化方式としてDP(データ並列)、TP(テンソル並列)、PP(パイプライン並列)、EP(エキスパート並列)をサポートしている。
サービング層では非同期エンジンが中心となり、各ステップ後に新旧リクエストを統合して処理する連続バッチ処理が有効になる。クライアントはHTTPまたはgRPC経由でリクエストを送り込む構成だ。
CUDAグラフによるレイテンシ最適化
エンジン初期化時、--enforce-eagerオプションが指定されていない場合、各バッチサイズに対してダミー実行を行いCUDAグラフをキャプチャする。CUDAグラフはGPU処理のシーケンス全体をDAGとして記録しておき、推論時に再生(replay)することでカーネル起動オーバーヘッドを削減し、レイテンシを改善する。
本記事はvLLMの解説シリーズ第1弾であり、今後はサブシステムごとにより深い解説が続く予定だ。スケジューラの詳細、ベンチマーク手法(ルーフラインモデル:演算強度とメモリ帯域の比からハードウェア性能上限を可視化する分析手法)、自動チューニングについても扱われる予定となっている。
詳細はInside vLLM: Anatomy of a High-Throughput LLM Inference Systemを参照していただきたい。