vLLMにおける投機的デコーディング：Medusa、EAGLE、マルチトークン投機による2.5倍の推論速度

大規模言語モデル（LLM）の推論レイテンシーは、リアルタイムアプリケーションにとって重大なボトルネックです。各トークンが順次生成される自己回帰デコーディングは、本質的にスループットを制限します。投機的デコーディング（speculative decoding）は、より小さく高速な「ドラフト」モデルを利用して複数のトークンを同時に提案し、それをより大きな「ターゲット」モデルが並行して検証するという、パラダイムシフトをもたらします。この手法は、LLM推論のメモリ帯域幅に制約される性質を緩和し、計算量の増加と引き換えにレイテンシーを効果的に削減します。

このガイドでは、vLLMにおける投機的デコーディングの実装とアーキテクチャ上の考慮事項について、Medusa、EAGLE、およびマルチトークン投機戦略に焦点を当てて詳しく説明します。これらの手法が、L4やH100のような最新のNVIDIA GPUで、しばしば2.5倍を超える大幅な推論速度向上をどのように達成するかを探ります。

投機的デコーディングのパラダイム

従来の自己回帰デコーディングは、以下のループを含みます。

1. 次のトークンのロジットを計算する。
2. 次のトークンをサンプリングする。
3. シーケンスにトークンを追加する。
4. 繰り返す。

このプロセスは本質的にシーケンシャルです。各ステップでは、ターゲットモデル全体に対する完全な順伝播パスが必要であり、これは通常、大規模なモデルパラメータとKVキャッシュアクセスによりメモリ帯域幅に制約されます。

投機的デコーディングは、ドラフトモデルを導入することでこのシーケンシャルな依存関係を解消します。ワークフローは以下の通りです。

1. ドラフトモデルがk候補トークンのシーケンスを生成します。ドラフトモデルは大幅に小さいため、これは高速な操作です。
2. ターゲットモデルは、元のプロンプトとk候補トークンに対して単一の順伝播パスを実行します。これにより、kトークンの検証が並列化されます。
3. 各候補トークンについて、ターゲットモデルのロジットがドラフトモデルのロジットと比較されます。
4. 受け入れられたトークンは出力に追加されます。トークンが拒否された場合、プロセスは最後に受け入れられたトークンから再開され、次のトークンにはターゲットモデルのロジットが使用されます。

核となるアイデアは、大規模なターゲットモデルの順伝播パスのコストを複数のトークンに償却することです。効率の向上は、検証ステップごとに受け入れられるトークンの数に比例します。

ドラフト-ターゲット検証メカニズム

D をドラフトモデル、T をターゲットモデルとします。 シーケンス x_0, \dots, x_t が与えられた場合：

1. ドラフティング: D は k 個の候補トークン y_1, \dots, y_k を生成します。ここで y_i \sim P_D(y | x_0, \dots, x_t, y_1, \dots, y_{i-1}) です。
2. 検証: T は x_0, \dots, x_t, y_1, \dots, y_k のロジットを単一のバッチで計算します。これにより、各 y_i について P_T(y | x_0, \dots, x_t, y_1, \dots, y_{i-1}) が得られます。
3. 受理/拒否: 各 y_i について：
   • u \sim U(0,1) をサンプリングします。
   • もし u < \min(1, \frac{P_T(y_i | \text{context})}{P_D(y_i | \text{context})}) なら、y_i を受理します。
   • そうでなければ、y_i とそれに続くすべての候補 y_{i+1}, \dots, y_k を拒否します。次のトークンは、最後に受理されたトークンの P_T(y | \text{context}) からサンプリングされます。

このメカニズムにより、投機的デコーディングの出力分布が標準の自己回帰デコーディングの出力分布と同一であることが保証され、モデルの品質が維持されます。

vLLMにおける投機的デコーディング戦略

vLLMは投機的デコーディングを堅牢にサポートしており、ドラフティングプロセスを最適化するための様々な戦略を統合しています。

1. マルチトークン投機（バニラ）

これは基本的なアプローチであり、ドラフトモデルは通常、ターゲットモデルのより小さくファインチューニングされたバージョン、または完全に異なる高速なモデルです。ドラフトモデルはトークンの線形シーケンスを生成します。

アーキテクチャ:

• ドラフトモデル: より小さなLLM（例：Llama-70Bターゲットに対するLlama-7B）。
• ターゲットモデル: フルサイズのLLM。
• vLLM統合: vLLMスケジューラが両方のモデルを管理し、ドラフティングと検証のステップを調整します。両モデルのKVキャッシュは効率的に管理されます。

コード例（vLLM設定）:

from vllm import LLM, SamplingParams

# Initialize the target model
target_model_path = "meta-llama/Llama-2-7b-hf" # Or Llama-3-8B, Mixtral-8x7B, etc.
llm = LLM(
    model=target_model_path,
    tensor_parallel_size=1, # Adjust based on GPU count
    gpu_memory_utilization=0.9,
    # Enable speculative decoding with a draft model
    speculate_model="google/gemma-2b", # A smaller, faster draft model
    num_speculative_tokens=5, # Number of tokens to draft
    max_model_len=2048,
)

# Define sampling parameters
sampling_params = SamplingParams(
    temperature=0.0, # Deterministic for benchmarking
    top_p=1.0,
    max_tokens=128,
)

# Generate text
prompts = [
    "What is the capital of France?",
    "Write a short poem about a cat.",
]

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

# Example of disabling speculative decoding for comparison
# llm_no_spec = LLM(model=target_model_path, gpu_memory_utilization=0.9)
# outputs_no_spec = llm_no_spec.generate(prompts, sampling_params)

2. Medusa: ツリーベースの投機

Medusaは、各ステップで複数の次のトークンを予測し、予測ツリーを形成することで、マルチトークン投機を強化します。単一のドラフトヘッドの代わりに、Medusaはいくつかの「デコーディングヘッド」をターゲットモデルに追加します。各ヘッドは異なる将来の位置のトークンを予測します。

アーキテクチャ:

• Medusaヘッドを持つターゲットモデル: ターゲットモデルは、最終隠れ状態の上にN追加の線形層（デコーディングヘッド）で拡張されます。各ヘッドは、h_i位置のトークンを予測するように訓練されますt+i+1。
• ドラフティング: 変更されたターゲットモデルを単一の順伝播パスで実行すると、N候補トークンが並列に生成されます。
• 検証: ターゲットモデル（Medusaヘッドなし）は、これらのNトークンとそれ以降の予測をツリー状に検証します。

利点:

• 単一モデル: 別途ドラフトモデルをロードする必要がなく、メモリフットプリントと複雑さが軽減されます。
• 並列予測: ドラフティング段階の単一の順伝播パスで複数のトークンが予測されます。

コード例（MedusaのvLLM設定）:

Medusaは、Medusaヘッドで特別にファインチューニングされたモデルを必要とします。vLLMはこれらのモデルを直接サポートしています。

from vllm import LLM, SamplingParams

# Assuming a Medusa-trained model is available
# Example: a Llama-2-7b model fine-tuned with Medusa heads
medusa_model_path = "medusa-llama-2-7b-hf" # Placeholder for a Medusa-trained model
target_model_path = "meta-llama/Llama-2-7b-hf" # The base model for verification

llm_medusa = LLM(
    model=target_model_path, # vLLM uses the base model for verification
    tensor_parallel_size=1,
    gpu_memory_utilization=0.9,
    # Specify the Medusa draft model. This is typically the same base model
    # but vLLM expects a path to a model with Medusa heads.
    # In practice, you'd point this to the Medusa-augmented model.
    speculate_model=medusa_model_path,
    num_speculative_tokens=5, # Number of tokens to draft (corresponds to Medusa heads)
    max_model_len=2048,
)

sampling_params = SamplingParams(temperature=0.0, top_p=1.0, max_tokens=128)
prompts = ["Explain the concept of quantum entanglement."]
outputs = llm_medusa.generate(prompts, sampling_params)

for output in outputs:
    print(f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}")

注: vLLM 0.4.0以降、Medusaのサポートが統合されています。speculate_model引数はベースモデルを指し、vLLMはモデルの設定またはチェックポイントにMedusaヘッドが存在する場合、自動的にそれらを検出してロードします。カスタムMedusaモデルの場合、アーキテクチャが互換性があることを確認してください。

3. EAGLE: 特徴レベルのドラフト再帰

EAGLE（Extending Auto-Regressive Generation with Lookahead Enhancement）は異なるアプローチを取ります。EAGLEは、トークンを直接予測するのではなく、ターゲットモデルの隠れ状態を予測するために小さな「ドラフト」モデルを訓練します。これにより、より堅牢で柔軟なドラフティングプロセスが可能になります。

アーキテクチャ:

• ターゲットモデル: フルサイズのLLM。
• EAGLEドラフトモデル: ターゲットモデルの隠れ状態を予測するために訓練された小さなモデル（例：数層のトランスフォーマー層）。このドラフトモデルは特徴レベルで動作します。
• ドラフティング: EAGLEドラフトモデルは、ターゲットモデルからの現在の隠れ状態を受け取り、次のkトークンの隠れ状態を予測します。これらの予測された隠れ状態は、候補トークンを得るためにターゲットモデルの最終線形層を通過します。
• 検証: 標準のターゲットモデル検証。

利点:

• より強力なドラフト: 隠れ状態を予測することで、特に複雑なシーケンスにおいて、より正確なドラフトにつながる可能性があります。
• 柔軟性: EAGLEドラフトモデルは、ターゲットモデル全体を再訓練することなく、異なるターゲットモデルに容易に適応させることができます。

コード例（EAGLEのvLLM設定）:

EAGLEモデルは通常、別個のチェックポイントです。

from vllm import LLM, SamplingParams

target_model_path = "meta-llama/Llama-2-7b-hf"
eagle_draft_model_path = "google/gemma-2b" # Placeholder for an EAGLE-trained draft model

llm_eagle = LLM(
    model=target_model_path,
    tensor_parallel_size=1,
    gpu_memory_utilization=0.9,
    speculate_model=eagle_draft_model_path, # Specify the EAGLE draft model
    num_speculative_tokens=7, # Number of tokens to draft
    max_model_len=2048,
)

sampling_params = SamplingParams(temperature=0.0, top_p=1.0, max_tokens=128)
prompts = ["Describe the process of photosynthesis in detail."]
outputs = llm_eagle.generate(prompts, sampling_params)

for output in outputs:
    print(f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}")

パフォーマンスベンチマークとトレードオフ

投機的デコーディングの主な利点は、最初のトークンまでの時間（TTFT）の短縮とスループット（トークン/秒）の向上です。実際のゲインはいくつかの要因に依存します。

• ドラフトモデルの品質: より良いドラフトモデルは、より高い受理率につながり、メリットを最大化します。
• num_speculative_tokens (k): 少なすぎると検証のオーバーヘッドが支配的になります。多すぎると受理率が低下し、頻繁な拒否と無駄な計算が発生します。最適なkは通常4から8の間です。
• ハードウェア: メモリ帯域幅と計算量。投機的デコーディングは、メモリ帯域幅（シーケンシャルなKVキャッシュアクセス）を計算量（並列検証）の増加と交換します。計算量対メモリ帯域幅の比率が高いGPU（例：H100）は、より大きな恩恵を受けます。

ベンチマーク比較（トークン/秒）

考察:

• 大幅な高速化: すべての投機的デコーディング手法は、特にターゲットモデルの順伝播パスが高価な大規模モデルにおいて、実質的なゲインをもたらします。
• H100の利点: H100 GPUは、その高い計算能力により、より大きな相対的ゲインを示しており、並列検証の計算集約的な性質を浮き彫りにしています。
• EAGLE/Medusaの優位性: EAGLEとMedusaは、より洗練されたドラフティングメカニズムにより、バニラのマルチトークン投機を上回ることがよくあります。

本番環境での落とし穴とトラブルシューティング

1. 投機的デコーディングにおけるCUDA out of memory:

   • 問題: 投機的デコーディング、特に別個のドラフトモデルを使用する場合、GPUメモリ消費が増加します。ターゲットモデルとドラフトモデル（およびそれらのKVキャッシュ）の両方がVRAMに常駐します。
   • 解決策:
     • gpu_memory_utilizationコンストラクタでLLMを減らす。
     • num_speculative_tokensを減らす。
     • より小さなドラフトモデルを使用する。
     • tensor_parallel_sizeを増やして、より多くのGPUにモデルを分散させる。
     • Medusaを使用している場合、Medusaヘッドが過度に大きくないことを確認する。

2. 出力品質の低下 / 不正確な応答:

   • 問題: 投機的デコーディングは理論的には同じ出力分布を生成することが保証されていますが、実装のバグや誤った設定が問題を引き起こす可能性があります。vLLMの堅牢な実装ではこれは稀です。
   • 解決策:
     • temperatureとtop_pの設定を確認する。投機的デコーディングは、決定論的サンプリング（temperature=0.0、top_p=1.0）で最も効果的です。確率的サンプリングは、完璧に処理されない場合、微妙な問題を露呈することがあります。
     • ドラフトモデルがターゲットモデルと適切に整合していることを確認する。訓練が不十分なドラフトモデルは受理率が低く、推論速度がベースラインと同等かそれ以下に低下します。
     • vLLMのバージョンを確認する。安定した投機的デコーディングのサポートがある最新バージョンを使用していることを確認する。

3. パフォーマンスの改善なし / 推論の遅延:

   • 問題: 投機的デコーディングにはオーバーヘッドがあります。受理率が低すぎる場合や、ドラフトモデルが遅すぎる場合、オーバーヘッドがメリットを上回ることがあります。
   • 解決策:
     • プロファイリング: nvprofまたはNVIDIA Nsight Systemsを使用してGPU使用率をプロファイリングする。GPU使用率の低い期間や過剰なメモリ転送を探す。
     • ドラフトモデルの選択: ドラフトモデルがターゲットモデルよりも大幅に小さく、高速であることを確認する。ターゲットモデルの1/10程度のサイズのドラフトモデルが良い出発点です。
     • num_speculative_tokensのチューニング: num_speculative_tokensを試す。4〜5から始めて、モデルとハードウェアに最適なスイートスポットを見つけるために増減させる。kが高すぎると多くの拒否につながり、kが低すぎるとターゲットモデルのコストを十分に償却できません。
     • バッチサイズ: 投機的デコーディングは、より大きなバッチサイズから恩恵を受け、GPUをより有効に活用できます。ワークロードに十分な同時リクエストがあることを確認する。
     • モデルの整合性: ドラフトモデルの予測が常に不十分な場合（受理率が低い場合）、ターゲットモデルにとって良い「教師」ではない可能性があります。ターゲットモデルの出力に似たデータでドラフトモデルをファインチューニングすることを検討する。

4. Medusaモデルのロード時にKeyError: 'medusa_num_heads'または類似のエラーが発生する:

   • 問題: vLLMはMedusaに対して特定の構成キーまたはモデルアーキテクチャを期待します。モデルのチェックポイントがこれに準拠していない場合、ロードに失敗する可能性があります。
   • 解決策:
     • Medusaモデルが、必要なconfig.jsonエントリ（例：medusa_num_heads、medusa_start_idx）で正しく訓練され、保存されていることを確認する。
     • speculate_modelパスが、vLLMがヘッドのロードを処理する場合、正しいMedusa拡張モデルまたはそのベースモデルを指していることを確認する。
     • 正確なMedusaモデル形式については、vLLMのドキュメントを参照する。

よくある質問

Q1: 投機的デコーディングは生成されるテキストの品質を低下させますか？

A1: いいえ。投機的デコーディングは、標準の自己回帰デコーディングとまったく同じ分布からサンプルを生成することが数学的に保証されています。知覚される品質の低下は、設定ミス、バグ、または投機的デコーディングのコアアルゴリズムとは無関係な問題に起因する可能性が高いです。

Q2: 最適なnum_speculative_tokens (k) は何ですか？

A2: 最適なkは、特定のターゲットモデル、ドラフトモデル、およびハードウェアに大きく依存します。一般的に、4〜8トークンの値が最適なバランスを提供します。kが低すぎるとターゲットモデルのコストを十分に償却できず、kが高すぎると頻繁な拒否と無駄な計算につながります。経験的なチューニングが推奨されます。

Q3: どんな小さなモデルでもドラフトモデルとして使用できますか？

A3: どんな小さなモデルでも使用することはできますが、ドラフトモデルとしての有効性は、ターゲットモデルの次のトークンを正確に予測する能力に依存します。ターゲットモデルの蒸留版またはファインチューニング版、あるいはドラフティング用に特別に設計されたモデル（EAGLEなど）は、はるかに高い受理率とそれによる高速化をもたらします。

Q4: 投機的デコーディングは常に標準の自己回帰デコーディングよりも高速ですか？

A4: 常にそうとは限りません。ドラフトモデルが遅すぎる場合、受理率が低すぎる場合、または2つのモデルを管理するオーバーヘッドがメリットを上回る場合、投機的デコーディングは遅くなる可能性があります。これは、非常に小さなターゲットモデルや、メモリ帯域幅が主要なボトルネックではないハードウェアで特に当てはまります。最新のGPU上の大規模モデル（例：H100上のLlama-70B）の場合、高速化は常に顕著です。

Q5: 投機的デコーディングはバッチ処理とどのように相互作用しますか？

A5: 投機的デコーディングはバッチ処理と非常に相補的です。リクエストのバッチを処理する場合、ターゲットモデルは異なるリクエストに対して複数の投機的シーケンスを並行して検証できます。これにより、GPUの使用率と全体的なスループットがさらに向上します。vLLMのPagedAttentionおよび連続バッチ処理メカニズムは、投機的デコーディングと効率的に連携するように設計されています。