投機的デコーディングでLLM推論を高速化:vLLMとllama.cpp実践ガイド

AI tutorial - IT technology blog
AI tutorial - IT technology blog

クイックスタート:5分でトークン生成速度を2倍にする

インタラクティブなアプリケーションで大規模言語モデル(LLM)を配信する際、常に直面する壁がメモリ帯域幅です。自己回帰(Autoregressive)デコーディングでは、1トークンを生成するたびに高帯域幅メモリ(HBM)からすべての重みを読み出す必要があります。2基のA100上で動作する70Bパラメータモデルの場合、このメモリバウンドな処理によって生成速度が著しく低下し、バッチサイズ1ではわずか毎秒18〜22トークン程度にとどまることがよくあります。

投機的デコーディング(Speculative Decoding)はこのボトルネックを解消します。軽量で小さなドラフトモデルが複数の候補トークンを素早く生成し、大型のターゲットモデルが1回のフォワードパスですべての候補を同時に検証します。計算コストは1トークンを生成する場合とほぼ同等でありながら、一度に3〜5トークンを獲得できます。

ここでは、vLLMとllama.cppの両方でパイプラインを構築する手順を説明します。

選択肢A:vLLMによる高速セットアップ

すでにNVIDIA GPUクラスター上でvLLMを実行している場合、投機的デコーディングを有効にするには、アライメントされた小型ドラフトモデルを指定する3つのコマンドラインフラグを追加するだけです。

pip install vllm

# Llama-3.1-8B-InstructをドラフトモデルとしてLlama-3.1-70B-Instructを配信
vllm serve meta-llama/Meta-Llama-3.1-70B-Instruct \
    --speculative-model meta-llama/Meta-Llama-3.1-8B-Instruct \
    --num-speculative-tokens 5 \
    --tensor-parallel-size 4 \
    --port 8000

重みのロードが完了したら、テストリクエストを送信してエンジンを検証します:

curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "meta-llama/Meta-Llama-3.1-70B-Instruct",
    "messages": [{"role": "user", "content": "JSONLファイルを安全にパースするPython関数を作成してください。"}],
    "temperature": 0.1
  }'

選択肢B:llama.cppによる省リソースセットアップ

Apple Silicon、ローカルワークステーション、またはエッジノード環境では、llama.cppの-md(model draft)フラグを使用することで、追加設定なしでネイティブな投機的サンプリングを利用できます。

# Metal / CUDAサポートを有効にしてllama.cppをビルド
git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp && cmake -B build && cmake --build build --config Release -j

# Qwen2.5-1.5BをドラフトモデルとしてQwen2.5-14Bを実行
./build/bin/llama-speculative \
    -m models/qwen2.5-14b-instruct-q4_k_m.gguf \
    -md models/qwen2.5-1.5b-instruct-q4_k_m.gguf \
    -p "TCPのフロー制御を3つの箇条書きで説明してください:" \
    -ngl 99 -ngld 99 \
    --draft-max 6

M2 Pro Mac(32GB統合メモリ)において、この設定により生成速度は毎秒17.8トークンから毎秒33.4トークンへと急上昇し、コードを一切変更することなく即座に87%の高速化を達成しました。

投機的デコーディングの仕組み

なぜ出力品質を損なうことなくこの手法が機能するのでしょうか? その鍵は「演算強度(Arithmetic Intensity)」にあります。

メモリ帯域幅のボトルネック

標準的なトークン生成は深刻なメモリバウンドに陥っています。1トークンを出力するたびに、GPUは何十ギガバイトもの重みをオンチップSRAMに読み込み、わずかな行列ベクトル積を計算した後にその重みを破棄します。その結果、演算コアは最大80%の時間アイドリング状態になります。

複数のトークンをバッチ処理することで状況は一変します。最新のGPUでは、6トークンを含むプロンプトの処理にかかる時間は、単一トークンの処理とほぼ同等です。投機的デコーディングはこの余剰演算能力を巧みに活用します。

2段階の実行ループ

デコーディングサイクルは、綿密に連携する3つのステージで動作します:

  1. ドラフトステップ: 小型モデル(1.5Bや8Bなど)が自己回帰的にK回(通常4〜6トークン)実行されます。ドラフトモデルはサイズが小さいため、メモリ転送は数ミリ秒で完了します。
  2. 並列検証: 大型ターゲットモデルが、プロンプトとK個の候補トークンすべてを1回の並列フォワードパスで一括処理します。
  3. 採択と修正: ターゲットモデルは、ドラフトトークンの確率と自身の計算結果を比較します。不一致が見つかるまで有効なトークンを順番に採択します。最初の不一致トークンに達した時点で、ターゲットモデルは自身の確率分布から直接、代替トークンをサンプリングします。

これにより出力内容は変わるのでしょうか? 答えは「いいえ」です。投機的デコーディングは修正棄却サンプリング(Modified Rejection Sampling)を採用しています。数学的に、出力分布はビット単位でターゲットモデル本来の確率分布と完全に一致します。推論の深さを全く損なうことなく、エンドツーエンドのレイテンシを半減させることができます。

本番環境向けアドバンスド設定

単一のプロンプトで投機的デコーディングを実行するのは簡単です。しかし、本番環境で数千件の同時リクエストを処理しながらスループットを維持するには、綿密なサイジングとチューニングが欠かせません。私たちは過去6か月間、まさにこのアーキテクチャを用いて1日あたり1,400万件以上の検索・コーディングリクエストを処理してきました。

1. Prompt LookupによるVRAMゼロのドラフティング

必ずしも第2のニューラルネットワークが必要なわけではありません。GPUメモリがすでに逼迫している場合、Prompt Lookup DecodingやN-gramキャッシングのようなヒューリスティックな手法を使えば、追加のVRAM消費ゼロで大幅なパフォーマンス向上が得られます。

from vllm import LLM, SamplingParams

sampling_params = SamplingParams(temperature=0.0, max_tokens=256)

# ドラフトモデルをロードせずにn-gram投機的デコーディングを設定
llm = LLM(
    model="mistralai/Mistral-7B-Instruct-v0.3",
    speculative_model="[ngram]",
    num_speculative_tokens=5,
    ngram_prompt_lookup_max=4,
    trust_remote_code=True
)

prompts = ["カスタマーサービスの対話履歴を要約してください:シカゴの悪天候により注文 #9823 に遅延が発生..."]
outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    print(output.outputs[0].text)

ドキュメントの要約、コードのリファクタリング、構造化データの抽出などでは、出力トークンがプロンプト内の単語を繰り返し使用することがよくあります。これらのワークロードにおいて、N-gram投機的ルックアップはメモリ使用量を一切増やすことなく、58%〜64%の採択率を達成します。

2. 投機先読みトークン数(K)の調整

--num-speculative-tokensフラグは、ターゲットモデルによる検証前にドラフトモデルが生成する候補トークン数を指定します。この値の選定は最適化問題となります:

  • K = 3〜4: 自由形式の文章、クリエイティブライティング、複雑なマルチステップ推論など、ドラフトの採択率が50%を下回るタスクに最適です。
  • K = 5〜6: 構造化データの生成(JSONペイロード、SQL構文、スキーマ検証)におけるスイートスポットです。トークンの予測可能性が高いため、採択率を75%以上に維持できます。
  • K > 8: 収穫逓減が発生します。ドラフトエンジンが2番目のトークンで外した場合、3〜8番目のトークンの検証はターゲットGPUのサイクルを無駄に消費するだけになります。

3. マルチGPUノードにおけるVRAM割り当て

ドラフトモデルは物理VRAMをプライマリモデルと共有します。FP16の8Bドラフトモデルをロードすると約16GBのVRAMが必要となり、KVキャッシュの割り当て容量が削られ、最大同時実行数の上限が低下します。

4基のA100(80GB)ノードでLlama-3.1-70Bを配信する際にKVキャッシュのスラッシングを防ぐには、ドラフトモデルをFP8に量子化するか、トポロジ全体に適切に配置します:

vllm serve meta-llama/Meta-Llama-3.1-70B-Instruct \
    --tensor-parallel-size 4 \
    --speculative-model neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8 \
    --speculative-draft-tensor-parallel-size 1 \
    --gpu-memory-utilization 0.90 \
    --max-num-seqs 64

本番運用の現場知見

トークナイザーの厳格な一致を徹底

ドラフトモデルとターゲットモデルは、完全に同一のボキャブラリおよびトークナイズ方式を共有している必要があります。Llama-3のターゲットとQwenのドラフトモデルを組み合わせると、トークンIDテーブルが異なるため即座にクラッシュします。必ず同一アーキテクチャ系統のモデルペアを選択してください:

  • ターゲット:Llama-3.1-70B → ドラフト:Llama-3.1-8B または Llama-3.2-3B
  • ターゲット:Qwen2.5-72B → ドラフト:Qwen2.5-1.5B または Qwen2.5-7B
  • ターゲット:DeepSeek-Coder-33B → ドラフト:DeepSeek-Coder-1.3B

Temperatureは0.4以下に維持

投機的デコーディングは、トークン分布がシャープな場合に最大のパフォーマンス向上を発揮します。temperature = 0.0(Greedyデコーディング)では、ターゲットモデルとドラフトモデルの出力が最も高頻度で一致し、72%以上の採択率を維持できます。

Temperatureが0.8や1.0に近づくにつれて採択率は25%前後まで低下します。その水準では、ドラフト生成のオーバーヘッドが検証による高速化効果を完全に打ち消し、ベースラインの自己回帰デコーディングよりも全体のレイテンシが悪化します。

Prometheusで採択率を監視

vLLMのメトリクスをPrometheusに公開し、vllm:spec_decode_draft_acceptance_rateにアラートを設定します。正常な本番クラスタでは、このメトリクスは0.55〜0.75の間で推移します。もし採択率が0.40を下回った場合は、投機的デコーディングを動的に無効化するかN-gramルックアップに切り替え、サーバーリソースの浪費を防ぎます。

Share: