vLLMにおける投機的デコーディング：EagleとMedusaでLLM推論スループットを3倍に

大規模言語モデル（LLM）はAIアプリケーションに革命をもたらしましたが、本番環境でのデプロイは、推論のレイテンシーとスループットによってボトルネックになることがよくあります。各トークンが順次生成される自己回帰デコーディングは、特に大規模モデルや長いシーケンスの場合、本質的に低速です。このドキュメントでは、このボトルネックを軽減するための強力な手法である投機的デコーディング（speculative decoding）について、vLLMでの実装とMedusaやEagleのような高度なバリアントに焦点を当て、権威ある詳細なガイドを提供します。基盤となるアーキテクチャ原則、実用的なvLLM構成、パフォーマンスベンチマーク、および重要な本番環境での考慮事項について探求します。

LLM推論における自己回帰のボトルネック

LLMはトークンごとにテキストを生成します。各トークンについて、モデル全体の順伝播が実行され、かなりの計算リソースを消費します。このシーケンシャルな性質は、総推論時間がアウトプットシーケンス長に線形にスケールすることを意味し、リアルタイムで高スループットなアプリケーションを困難にします。

典型的なLLM推論リクエストを考えてみましょう。

1. プロンプトエンコーディング: 入力プロンプトはトークン化され、数値埋め込みに変換されます。
2. 最初のトークン生成: LLMはプロンプト埋め込みを処理し、最初の出力トークンを予測します。これには完全な順伝播が含まれます。
3. 後続トークン生成: 新しく生成されたトークンが入力シーケンスに追加され、プロセスが繰り返されます。後続の各トークンは、先行するすべてのトークンを条件として、別の完全な順伝播を必要とします。

この反復プロセスは、高品質な出力を保証する一方で、計算コストが高くなります。投機的デコーディングのようなパフォーマンス最適化手法の主な目標は、この厳密なシーケンシャルな依存関係を打破し、可能な限り並列処理を導入することです。

投機的デコーディング：基本原則

投機的デコーディングは、より小さく高速な「ドラフト」モデルを利用して複数の将来のトークンを並行して予測し、それらをより大きく正確な「ターゲット」モデルが単一の効率的な順伝播で検証することで、LLM推論を高速化します。

その核心となるアイデアはシンプルかつ深遠です。

1. ドラフト生成: 小さく計算コストの低いドラフトモデル（例：ターゲットモデルの小型版、または専用に構築されたドラフトモデル）が、k個の候補トークンのシーケンスを迅速に生成します。これは将来のトークンに関する投機的な推測です。
2. ターゲット検証: 主要な高品質LLMであるターゲットモデルは、元のプロンプトとk個の候補トークンを単一のバッチで処理します。これにより、シーケンス内の各位置の確率分布を予測し、ドラフトモデルの予測を効果的に検証します。
3. 承認と拒否: 各候補トークンについて、ターゲットモデルの予測トークンがドラフトモデルのトークンと一致する場合、候補は承認されます。不一致が発生するか、すべてのk個の候補が検証されるまでプロセスは続行されます。
4. 再サンプリング: 不一致が発生した場合、その位置のターゲットモデルの予測が使用され、残りの未検証の候補は破棄されます。その後、プロセスは最後に承認されたトークンから再開され、新しい候補が生成されます。

このメカニズムにより、ターゲットモデルはトークンを1つずつ生成するのではなく、複数のトークンを並行して検証することで「先読み」することができます。効率の向上は、検証ステップごとに承認されるトークンの平均数（「承認率」）に直接比例します。

アーキテクチャの概要：ドラフトモデルとターゲットモデルの相互作用

一般的な設定では、2つの異なるモデルがGPUメモリにロードされます。

• ターゲットモデル: 大規模で高品質なLLM（例：Llama-3-70B）。このモデルが最終的な正確な出力を担当します。
• ドラフトモデル: より小さく高速なモデル（例：Llama-3-8B、Eagle-7B）。このモデルの主な役割は、もっともらしいトークンシーケンスを迅速に生成することです。

投機的デコーディングによる推論ループは次のように進行します。

パフォーマンス向上の鍵は、ステップEとFの並列性にあります。ターゲットモデルでのk回のシーケンシャルな順伝播の代わりに、k個のトークンを検証するために1回の順伝播のみが実行されます。

高度な投機的デコーディング：ツリーベースの検証（Medusa & Eagle）

標準的な投機的デコーディングは大きな利点をもたらしますが、その効率は線形的な承認率によって制限されます。ドラフトモデルがターゲットと完全に一致しない場合、不一致が早期に発生し、実効的なkが低下する可能性があります。MedusaやEagleのようなツリーベースの投機的デコーディング手法は、候補トークンのツリーを生成することでこの問題に対処し、より堅牢な並列検証を可能にします。

Medusa: ツリーベースの投機のためのマルチヘッドデコーディング

Medusa（Multi-Head Decoding）は、単一のドラフトモデルに複数の予測ヘッドを追加することで、投機的デコーディングを強化します。ドラフトモデルは、単一の次のトークンを予測する代わりに、異なる将来の位置で複数の可能な次のトークンを同時に予測します。

Medusaの仕組み:

1. マルチヘッドドラフトモデル: 標準的なLLMは、複数の出力ヘッドを持つようにファインチューニングまたは適応されます。各ヘッドh_iは、現在のトークンに対する相対的な位置iのトークンを予測するように訓練されます。例えば、h_0は次のトークンを予測し、h_1はt_1が与えられた場合にその次のトークンを予測します。
2. ツリー生成: 現在のトークンが与えられると、マルチヘッドドラフトモデルは候補トークンの「ツリー」を生成します。各ヘッドがブランチに貢献します。例えば、h_0はt_1を予測し、h_1はt_2を予測します。これにより、可能な継続の小さなローカル検索ツリーが作成されます。
3. 並列検証: ターゲットモデルは、生成されたツリー内のすべてのパスを単一の順伝播で検証します。これは、ターゲットモデルに元のプロンプトとツリーからのすべての候補シーケンスをフィードすることで行われます。
4. 最適なパス選択: ターゲットモデルの出力確率を使用して、生成されたツリー内で最も長い有効なプレフィックスを特定します。これにより、一部のブランチがすぐに分岐しても、より多くのトークンを受け入れることができます。

Medusaは、単一のターゲットモデルパスで検証できるトークンの数を大幅に増やし、より高い実効承認率とより大きなスループットにつながります。

Eagle: Medusaスタイルのデコーディングに最適化されたドラフトモデル

Eagleは、Medusaのコンセプトをさらに発展させ、マルチヘッド投機的デコーディング専用の、高度に最適化されたドラフトモデルを設計しています。Medusaは通常、既存のLLMを複数のヘッドで適応させますが、Eagleは効率的なマルチヘッドドラフトモデルとしてゼロから設計されています。

Eagleの主な特徴:

• 最適化されたアーキテクチャ: Eagleモデルは、汎用LLMよりも小さく効率的であり、投機的なトークンツリーを生成するタスクのために特別に設計されています。
• 投機のために事前学習済み: これらは、より大規模なターゲットモデルによって承認される可能性が高い高品質の投機的なトークンシーケンスを生成することに焦点を当てて、訓練またはファインチューニングされることがよくあります。
• 高い承認率: その特殊な設計により、Eagleモデルは、汎用的な小型LLMをドラフトモデルとして使用する場合と比較して、より高い承認率を達成する傾向があります。

vLLM（または他の推論エンジン）でEagleをドラフトモデルとして使用すると、そのタスクのために特別に構築されているため、標準的な小型LLMを使用する場合よりも優れたパフォーマンスが得られます。

N-gramドラフト投機

vLLMの高度な実装ではあまり一般的ではありませんが、より単純な形式の投機的デコーディングも注目に値します。例えば、N-gramドラフト投機は、単純なN-gram言語モデル（または小さな非ニューラルモデル）を使用して次のk個のトークンを予測します。これは非常に高速ですが、ニューラルネットワークと比較してN-gramモデルの予測能力が限られているため、承認率が低くなることがよくあります。これは、より高度なニューラルベースの投機的デコーディングの初期の概念的基盤となりました。

vLLMの統合と設定

高スループットLLMサービングで知られるvLLMは、投機的デコーディングを強力にサポートしています。これにより、複雑さの多くが抽象化され、ユーザーはいくつかのコマンドライン引数またはAPIパラメータで有効にすることができます。

vLLMでの投機的デコーディングの有効化

投機的デコーディングを有効にするには、ターゲットモデルとドラフトモデルの両方を指定する必要があります。vLLMは、両方のロード、管理、および推論のオーケストレーションを処理します。

主要パラメータ:

• --model: ターゲットLLMのパスまたは名前。
• --speculative-model: ドラフトLLMのパスまたは名前。
• --num-speculative-tokens: ドラフトモデルが1回の投機的ステップで生成する最大トークン数。これはkに対応します。値が高いほど、1ステップで検証されるトークンが多くなる可能性がありますが、ドラフトモデルの精度が十分でない場合、早期拒否の可能性も高まります。
• --draft-model-tp-size: ドラフトモデルのテンソル並列処理サイズ。
• --target-model-tp-size: ターゲットモデルのテンソル並列処理サイズ。

例：投機的デコーディングを使用したvLLMの実行

Llama-3-70Bをターゲット、Llama-3-8BをドラフトモデルとしてvLLMを起動する方法をデモンストレーションします。これは、これらのモデル（例：Hugging Face Hubから）へのアクセスと十分なGPUリソースがあることを前提としています。

1. vLLMのインストール:

pip install vllm

2. 投機的デコーディングなしでvLLMを起動（ベースライン）：

python -m vllm.entrypoints.api_server \
    --model meta-llama/Meta-Llama-3-70B-Instruct \
    --tensor-parallel-size 2 \
    --port 8000

注：Llama-3-70Bは通常、フル精度で少なくとも2x A100 80GB GPU、または量子化（例：AWQ、GPTQ）で1x A100 80GBを必要とします。

3. 投機的デコーディング（Llama-3-8Bドラフト）でvLLMを起動：

python -m vllm.entrypoints.api_server \
    --model meta-llama/Meta-Llama-3-70B-Instruct \
    --speculative-model meta-llama/Meta-Llama-3-8B-Instruct \
    --num-speculative-tokens 8 \
    --tensor-parallel-size 2 \
    --draft-model-tp-size 1 \
    --port 8001

この構成では：

• --model meta-llama/Meta-Llama-3-70B-Instruct: ターゲットモデルを指定します。
• --speculative-model meta-llama/Meta-Llama-3-8B-Instruct: ドラフトモデルを指定します。
• --num-speculative-tokens 8: ドラフトモデルは投機的に最大8トークンを生成しようとします。
• --tensor-parallel-size 2: ターゲットモデル（70B）は2つのGPUにシャードされます。
• --draft-model-tp-size 1: ドラフトモデル（8B）は単一のGPUで実行されます。vLLMは、利用可能なGPUの1つにドラフトモデルをインテリジェントに配置します。メモリが許せばターゲットモデルのシャードと共有することも、専用のGPUを使用することも可能です。

4. 投機的デコーディング（Eagle-7Bドラフト）でvLLMを起動：

python -m vllm.entrypoints.api_server \
    --model meta-llama/Meta-Llama-3-70B-Instruct \
    --speculative-model google/gemma-2b-it-eagle \
    --num-speculative-tokens 12 \
    --tensor-parallel-size 2 \
    --draft-model-tp-size 1 \
    --port 8002

ここでは、google/gemma-2b-it-eagleをEagleスタイルのドラフトモデルの例として使用しています。num-speculative-tokensは、Eagleモデルの最適化された性質と期待される高い承認率のため、より高くなる可能性があることに注意してください。

クライアント側の推論例

これらのサーバーとやり取りするには、vLLM Pythonクライアントを使用できます。

from openai import OpenAI

# Client for baseline (no speculative decoding)
client_baseline = OpenAI(api_key="EMPTY", base_url="http://localhost:8000/v1")

# Client for speculative decoding with Llama-3-8B draft
client_spec_llama = OpenAI(api_key="EMPTY", base_url="http://localhost:8001/v1")

# Client for speculative decoding with Eagle-7B draft
client_spec_eagle = OpenAI(api_key="EMPTY", base_url="http://localhost:8002/v1")

def generate_text(client, prompt, max_tokens=128):
    chat_completion = client.chat.completions.create(
        model="meta-llama/Meta-Llama-3-70B-Instruct", # Model name is for API routing, not actual loading
        messages=[
            {"role": "system", "content": "You are a helpful AI assistant."},
            {"role": "user", "content": prompt}
        ],
        max_tokens=max_tokens,
        temperature=0.7,
        stream=False
    )
    return chat_completion.choices[0].message.content

prompt = "Explain the concept of quantum entanglement in simple terms."

print("--- Baseline Inference ---")
response_baseline = generate_text(client_baseline, prompt)
print(response_baseline)

print("\n--- Speculative Inference (Llama-3-8B Draft) ---")
response_spec_llama = generate_text(client_spec_llama, prompt)
print(response_spec_llama)

print("\n--- Speculative Inference (Eagle-7B Draft) ---")
response_spec_eagle = generate_text(client_spec_eagle, prompt)
print(response_spec_eagle)

ベンチマークとパフォーマンス分析

投機的デコーディングの利点を定量化するには、厳密なベンチマークが不可欠です。ここでは、異なる構成間でのスループット、レイテンシー、およびリソース使用率を比較します。

ベンチマーク方法論: vllm.benchmarks.benchmark_throughputを使用して、典型的な本番ワークロードをシミュレートします。このツールを使用すると、さまざまな負荷条件下での1秒あたりのトークン数（スループット）、最初のトークンのレイテンシー（FTL）、および合計レイテンシーを測定できます。

ハードウェア構成:

• GPU: 2x NVIDIA A100 80GB（Llama-3-70Bターゲットモデル用）
• CPU: Intel Xeon Platinum 8380（2.3 GHz、80コア）
• RAM: 512GB
• vLLMバージョン: 0.4.0（またはそれ以降）

ベンチマークシナリオ:

1. ベースライン: Llama-3-70B-Instruct（投機的デコーディングなし）。
2. 投機的（Llama-3-8Bドラフト）: Llama-3-70B-Instruct + Llama-3-8B-Instructドラフト（--num-speculative-tokens 8）。
3. 投機的（Eagle-7Bドラフト）: Llama-3-70B-Instruct + google/gemma-2b-it-eagleドラフト（--num-speculative-tokens 12）。

収集されたメトリクス:

• スループット（tokens/s）: 1秒あたりに生成される平均出力トークン数。
• 最初のトークンのレイテンシー（ms）: 最初の出力トークンを生成するのにかかる時間。
• 合計レイテンシー（128トークン、ms）: 128トークンの完全なシーケンスを生成するのにかかる時間。
• GPUメモリ使用量（GB）: すべてのモデルによって消費される総VRAM。
• 承認率（%）: ターゲットモデルによって承認されたドラフトトークンの平均割合。

ベンチマーク結果

注：これらのメトリクスは例示的なものであり、同様のセットアップで観察される典型的なパフォーマンスに基づいています。実際の数値は、特定のハードウェア、vLLMバージョン、モデルの量子化、およびワークロードの特性によって異なる場合があります。

結果の分析

1. スループット: 投機的デコーディングは、スループットを大幅に向上させます。Llama-3-8Bをドラフトモデルとして使用すると、スループットがほぼ3倍（2.74倍）になり、特殊なEagle-7Bドラフトモデルはさらに3.3倍以上に押し上げます。これは、1秒あたりにより多くのリクエストを処理したり、より長いシーケンスをより速く処理したりすることに直接つながります。
2. レイテンシー:
   • 最初のトークンのレイテンシー（FTL）: 投機的デコーディングでは、FTLがわずかに増加します。これは予想されることであり、初期設定では2つのモデルのロードと調整が必要であり、最初の投機的パスは単一のターゲットモデルパスよりもわずかに時間がかかる可能性があります。ただし、この増加は、長いシーケンスの合計生成時間の文脈では無視できることがよくあります。
   • 合計レイテンシー: 長いシーケンス（例：128トークン）の場合、合計レイテンシーは劇的に減少します。並列検証メカニズムにより、シーケンシャルなトークン生成に費やされる累積時間が大幅に短縮されます。
3. GPUメモリ: 投機的デコーディングは、2つのモデル（ターゲットとドラフト）を同時にロードする必要があるため、本質的に多くのGPUメモリを必要とします。Llama-3-70BはすでにA100 80GBの大部分を消費しており、8Bまたは2Bのドラフトモデルを追加すると、総メモリ要件が増加します。これにより、多くの場合、マルチGPUセットアップまたは積極的な量子化戦略が必要になります。
4. 承認率: 承認率は重要なメトリクスです。承認率が高いほど、より多くのドラフトトークンが検証および承認され、再サンプリングステップが少なくなり、効率が向上します。Eagle-7Bは、Llama-3-8B（72%）と比較して大幅に高い承認率（88%）を示しており、専用に構築されたドラフトモデルの利点を浮き彫りにしています。num-speculative-tokensパラメータは、ドラフトモデルの承認率と組み合わせて調整する必要があります。承認率が低いのに値が高すぎると、逆効果になる可能性があります。

TensorRT-LLMと投機的デコーディング

vLLMは投機的デコーディングのための高レベルでユーザーフレンドリーなインターフェースを提供しますが、NVIDIAのTensorRT-LLMは、特にNVIDIA GPU向けに、さらに深いレベルの最適化を提供します。TensorRT-LLMはLLMを高度に最適化された推論エンジンにコンパイルし、多くの場合、優れたパフォーマンスをもたらします。

TensorRT-LLMは投機的デコーディングもサポートしており、PyTorchや他の汎用バックエンドのみに依存するフレームワークよりも高いスループットと低いレイテンシーを達成することがよくあります。ドラフトモデルとターゲットモデルの順伝播、および検証ロジックの両方で、カスタムCUDAカーネルとハードウェア固有の最適化を活用しています。

本番環境での極端なパフォーマンス要件、特にNVIDIAハードウェアを使用する場合、投機的デコーディングのためにTensorRT-LLMを（直接またはそれをバックエンドとして使用するフレームワークを介して）統合すると、vLLM単独で提供される以上の追加の利点が得られます。vLLM自体は、特定のモデルでTensorRT-LLMと統合でき、vLLMの使いやすさとTensorRTの生のパフォーマンスを組み合わせることができます。

よくある落とし穴と本番環境での問題点

本番環境で投機的デコーディングをデプロイするには、安定性、パフォーマンス、費用対効果を確保するために、いくつかの要因を慎重に検討する必要があります。

1. ドラフトモデルの選択と不一致

• 落とし穴: 小さすぎたり、適切に訓練されていなかったり、ターゲットモデルとアーキテクチャ的に異なっていたりするドラフトモデルを使用すること。これにより、承認率が低くなり、パフォーマンスの向上が打ち消され、頻繁な再サンプリングのためにレイテンシーが増加する可能性があります。
• 問題点: どんな小さなモデルでもドラフトとして盲目的に選択すること。
• 軽減策:
  • アーキテクチャの整合性: 理想的には、ドラフトモデルはターゲットモデルの小型版（例：Llama-3-70Bに対するLlama-3-8B）であるか、投機のために特別に設計されたモデル（例：Eagle）であるべきです。
  • ドメインの整合性: 高い承認率を維持するために、ドラフトモデルがターゲットモデルと同じデータで訓練されているか、同じドメイン向けにファインチューニングされていることを確認してください。
  • ベンチマーク: 特定のターゲットモデルとワークロードで異なるドラフトモデルを徹底的にベンチマークし、ドラフトモデルの速度と承認率の最適なバランスを見つけてください。

2. GPUメモリのオーバーヘッド

• 落とし穴: 2つのモデルを同時に実行することによるGPUメモリの合計要件を過小評価すること。これにより、メモリ不足（OOM）エラー、バッチサイズの縮小、またはより高価なハードウェアの必要性が生じる可能性があります。
• 問題点: ドラフトモデルのメモリフットプリントが無視できると仮定すること。
• 軽減策:
  • 量子化: メモリフットプリントを削減するために、ターゲットモデルとドラフトモデルの両方に量子化（例：AWQ、GPTQ、FP8）を適用します。vLLMはさまざまな量子化スキームをサポートしています。
  • モデルシャーディング（テンソル並列処理）: ターゲットモデルを複数のGPUに分散させます。ドラフトモデルは通常、単一のGPUで実行でき、メモリが許せばターゲットモデルのシャードと共有することも可能です。
  • 動的メモリ割り当て: 開発およびテスト中にGPUメモリ使用量を注意深く監視します。必要に応じてバッチサイズを調整するか、より大きなGPUを検討してください。

3. num-speculative-tokensのチューニング

• 落とし穴: 経験的な検証なしにnum-speculative-tokensを高く設定しすぎたり、低く設定しすぎたりすること。
• 問題点: 値が高すぎると、多くのトークンが拒否され、計算が無駄になります。値が低すぎると、並列処理を十分に活用できない可能性があります。
• 軽減策:
  • 経験的チューニング: このパラメータは、ドラフトモデルの品質とターゲットモデルの特性に大きく依存します。特定のセットアップでスループットを最大化する最適な値を見つけるために、異なる値（例：4、8、12、16）をベンチマークしてください。
  • 承認率の監視: 適切なnum-speculative-tokens値は、高い承認率（例：70-80%以上）をもたらし、同時に大幅な高速化を提供します。

4. 動的バッチ処理との相互作用

• 落とし穴: 投機的デコーディングは、特にリクエストのプロンプトと出力の長さが異なる場合に、動的バッチ処理を複雑にします。
• 問題点: 最適でないバッチ処理戦略は、投機的デコーディングの効果を低下させる可能性があります。
• 軽減策:
  • vLLMのPagedAttentionと動的バッチ処理は、投機的デコーディングと連携するように設計されています。フレームワークの最適化を信頼してください。ただし、非常に多様なリクエストパターンは、システムに課題をもたらす可能性があることに注意してください。
  • 可能な限り同様のリクエストをグループ化するために、リクエストキューイングとスケジューリング戦略を検討し、バッチ利用率を最大化します。

5. コールドスタートレイテンシー

• 落とし穴: 1つのモデルではなく2つのモデルをロードするため、コールドスタートレイテンシーが増加すること。
• 問題点: サーバーレスまたはオートスケーリング環境での初期ロード時間を考慮しないこと。
• 軽減策:
  • プリウォーミング: 推論エンドポイントのプリウォーミング戦略を実装します。
  • 永続インスタンス: 高トラフィックサービスの場合、コールドスタートを避けるために永続インスタンスを維持します。
  • 最適化されたロード: ロード時間を最小限に抑えるために、モデルが高速ストレージ（例：NVMe SSD）に保存されていることを確認します。

6. 監視と可観測性

• 落とし穴: 本番環境で投機的デコーディングのパフォーマンスを監視するための特定のメトリクスが不足していること。
• 問題点: 全体的なスループットのみを監視し、投機的プロセスの効率に関する洞察を見逃すこと。
• 軽減策:
  • 主要メトリクス: 承認率を監視する