vLLMのAMD GPUにおける推測的デコード：パフォーマンスと手法
まとめ
推測的デコードをvLLMでAMD Instinct MI300X/MI355X GPUに導入することで、高速なドラフトモデルが複数の未来トークンを提案し、ターゲットモデルが1回の順伝播で検証する仕組みが可能になります。これにより、特定のモデル・手法の組み合わせでは最大約2.9倍のスループット向上が達成され、元のモデルの出力挙動は維持されます。
vLLMにおける推測的デコードの仕組み
推測的デコードは、元の（ターゲット）LLMの前にドラフト・検証ステージを追加します。ドラフトコンポーネントが候補トークンのシーケンスを生成し、ターゲットモデルはその全候補ブロックを1回の順伝播で評価します。ターゲットモデルによって受け入れられたトークンは確定され、最初に拒否されたトークンはターゲットモデル自身の出力で置き換えられます。生成が終了するまでこのプロセスを繰り返します。
主な特性
- ターゲットモデルが真実の出力源であり、検証前に出力トークンは生成されません。
- 1回のターゲットモデルの順伝播で複数のドラフトトークンを確定できるため、高コストなターゲットモデルの順伝播回数を削減できます。
- 受け入れは左から右へ評価され、拒否が発生すると残りの候補ブロックは中止されます。
評価されたドラフト手法
vLLMは5つの具体的なドラフトアプローチをサポートしており、それぞれがドラフトネットワークがターゲットモデルから情報をどのように受け取るか、およびトークン生成が逐次的か並列的かで異なります。
| 手法 | ドラフトコンポーネント | ターゲットモデルの情報の使用 | トークン生成スタイル |
|---|---|---|---|
| ネイティブMTP | モデルネイティブな補助予測パス | ターゲットモデルの隠れ表現（または前のMTPステップ） | 逐次的（繰り返しMTP呼び出し） |
| Gemma 4 MTP | ターゲットとペアになった別途のMTPチェックポイント | ターゲットの活性化 + 共有KVキャッシュ | 逐次的 |
| EAGLE‑3 | 専用の自己回帰的推測器 | ターゲット層の初期・中間・終盤の融合隠れ状態 | 逐次的、各ドラフトトークンが次のものに依存 |
| DFlash | 専用の並列推測器 | ターゲット隠れ状態を各ドラフト層で使用するKVペアに投影 | 1回の順伝播ですべての位置を同時に予測 |
| DSpark | DFlashのバックボーン + 軽量Markovヘッド | DFlashと同一のターゲットコンテキスト | 並列バックボーン + 逐次的なトークン選択の補正 |
ネイティブMTP
ターゲットモデル内に実装され、補助ヘッドを使って固定数の未来トークンを予測します。ドラフトトークンは1つずつ生成され、各ステップで前のMTP出力の隠れ状態を消費します。num_speculative_tokensがネイティブ深度を超えると、vLLMは追加の順伝播でMTPパスを再利用します。
Gemma 4 MTP
別途のチェックポイント（アシスタントモデル）がターゲットのKVキャッシュと活性化を共有しており、すでに計算されたコンテキストを再利用できます。ドラフトトークンは依然として逐次的に生成されますが、ドラフトコンポーネントがターゲットテンソルを再利用するため、メモリオーバーヘッドは小さいです。
EAGLE‑3
EAGLE‑3は、ターゲット層の3つの隠れ状態の融合表現を受け取る推測器を訓練します。推測器は自己回帰的に動作：最初のドラフトトークンは融合ターゲット特徴とサンプリングされたトークン埋め込みを組み合わせ、以降のトークンは前のドラフト出力に依存します。これにより、提案長に比例した逐次的なドラフト作業が得られます。
DFlash
DFlashは、未来の位置ブロック全体を並列で予測します。ターゲットによってすでに検証されたアンカートークンがブロックの開始となり、残りのマスクされた位置が同時に埋められます。ターゲット隠れ状態は、各ドラフト層がアクセスできる追加のKVペアに投影され、並列処理中にターゲットコンテキストを維持します。
DSpark
DSparkはDFlashのバックボーンに、並列バックボーンの後に軽量Markovヘッドを追加して、トークンごとの依存性を導入します。バックボーンはすべての位置のベースロジットを生成し、Markovヘッドは以前に選択されたドラフトトークンを使用して各位置を調整することで、並列の高速性と逐次的な整合性のハイブリッドを実現します。
vLLMでの推測的デコードの有効化
推測的デコードは--speculative-configフラグで設定されます。JSONペイロードは、手法、オプションのドラフトチェックポイント、およびnum_speculative_tokens（提案長）を指定します。
# ネイティブMTP（追加チェックポイント不要）
vllm serve <target-model> \
--speculative-config '{"method": "mtp", "num_speculative_tokens": 4}'
# Gemma 4 MTP（アシスタントチェックポイントが必要）
vllm serve <target-model> \
--speculative-config '{"method": "mtp", "model": "google/gemma-4-26B-A4B-it-assistant", "num_speculative_tokens": 4}'
# EAGLE‑3
vllm serve <target-model> \
--speculative-config '{"method": "eagle3", "model": "RedHatAI/<target>-speculator.eagle3", "num_speculative_tokens": 3}'
# DFlash（並列ブロック）
vllm serve <target-model> \
--speculative-config '{"method": "dflash", "model": "z-lab/<target>-DFlash", "num_speculative_tokens": 15}'
# DSpark（並列 + Markov補正）
vllm serve <target-model> \
--speculative-config '{"method": "dspark", "model": "RedHatAI/<target>-speculator.dspark", "num_speculative_tokens": 7}'
メモリの考慮事項 – ネイティブMTPはターゲットモデルと重みを共有するため、GPUメモリの追加はほとんどありません。他のすべての手法は別途のドラフトチェックポイントをロードするため、必要な余裕はチェックポイントのサイズ、精度（FP16/INT4）、テンソル並列度、ランタイムバッファに依存します。
ベンチマーク手法
著者たちは、AMD Instinct MI300XおよびMI355X GPU上でROCm™スタックを使用して出力トークンスループット（トークン/秒）を測定しました。ベンチマークにはタスクベースのデータセット（GSM8K、MATH500、HumanEval、MBPP）を使用し、現実的な受け入れパターンを反映しました。各ターゲット・手法ペアについて、num_speculative_tokens（N）をスイープし、以下の値を記録しました：
- スループット比（推測的 / ベースライン）
- 平均受け入れ長（MAL） – 1回の検証ラウンドあたりの平均受け入れドラフトトークン数
- 受け入れ率（AR） – 検証に耐えた提案トークンの割合
すべての測定値は、同じハードウェアおよびソフトウェアスタック上で実行された非推測的自己回帰ベースラインに対して相対的です。
主なパフォーマンスの発見
観測された最高のスピードアップ
| ターゲットモデル | 手法 | データセット | 最適N | スループット比 |
|---|---|---|---|---|
gemma-4-26B-A4B-it |
DFlash | MATH500 | 7 | 2.87× |
gemma-4-26B-A4B-it |
Gemma 4 MTP | MATH500 | 5 | 2.74× |
gemma-4-26B-A4B-it |
EAGLE‑3 | GSM8K | 5 | 2.27× |
gemma-4-31B-it |
DFlash | MATH500 | 7 | 2.41× |
Qwen3-8B |
DSpark | GSM8K | 7 | 1.63× |
Qwen3-6-35B-A3B |
DFlash | MATH500 | 7 | 2.06× |
Kimi-K2.5 |
DFlash | GSM8K | 7 | 2.37× |
手法ごとのトレンド
- 逐次的ドラフト（ネイティブMTP、Gemma 4 MTP、EAGLE‑3） は、やや小さい提案長（N ≈ 3–5）までスループットが向上します。それ以上になると、追加の逐次的ドラフト作業がターゲット順伝播の節約を上回り、スループットは plateau または低下します。
- 並列的ドラフト（DFlash、DSpark） は、大きなN（N ≈ 7–11）で最も高い比を達成することが多いです。後方の位置では受け入れ率が低下しますが、並列コストは低いため、全体的なスループットは依然として上昇します。
- DSpark は軽量Markovヘッドを追加しており、同じNに対して純粋なDFlashより数パーセント低いスループットになることが一般的です。これは追加の逐次的補正ステップによるものです。
- モデルファミリーの影響 – Qwen 3.5/3.6モデルに組み込まれたネイティブMTPは、同じモデルサイズでDFlashを上回りますが、Gemma 4では並列手法が優勢です。
- ワークロード依存性 – コード中心のデータセット（HumanEval、MBPP）は、後方のトークンが予測しにくくなるため、短い提案を好む傾向があります。一方、数学中心のデータセット（GSM8K、MATH500）は長い提案を許容します。
受け入れ挙動
すべての実験において、最初のドラフトトークンの受け入れ率は90％以上です。受け入れ率は位置とともに単調に低下します。DFlashでN = 15の場合、15番目のトークンの受け入れ率は5％未満になることがあります。平均受け入れ長（MAL）はスループットと強く相関しており、MALが高いほどターゲット順伝播回数が減り、スピードアップが大きくなります。
実用的なチューニングガイド
- 小さな値から始める – どの手法でも、
num_speculative_tokens = 1（追加コストなし）から始め、正しさを確認してください。 - Nをスイープする – 代表的なワークロード上で、スループット、MAL、ARを測定しながらNを増やします（例：1–3–5–7–11–15）。
- 位置ごとのARを観察する – 位置k以降で受け入れ率が急激に低下する場合、無駄なドラフト作業を避けるためにNをkに減らしてください。
- メモリ予算の管理 – ドラフトチェックポイント用に十分なGPUメモリを確保してください。並列手法（DFlash/DSpark）は通常、最も多くのメモリを必要とします。
- ハードウェア固有のチューニング – AMD Instinct GPUはROCm最適化されたアテンションバックエンド（
triton_attn）により恩恵を受けます。並列手法はMI300X/MI355Xの高いメモリ帯域幅を活用します。 - バッチサイズとトークン制限 – 大きな
max-num-batched-tokensと長いコンテキストウィンドウは、特に並列手法においてドラフトコストのアモリタイズを改善します。
ドラフトチェックポイントの入手先
| 公開者 | 対応する手法 | 例のチェックポイント |
|---|---|---|
| Gemma 4 MTP | google/gemma-4-26B-A4B-it-assistant, google/gemma-4-31B-it-assistant |
|
| LightSeek | EAGLE‑3, EAGLE‑3.1 | lightseekorg/kimi-k2.5-eagle3-mla |
| Red Hat AI | EAGLE‑3, DFlash, DSpark | RedHatAI/gemma-4-26B-A4B-it-speculator.eagle3, RedHatAI/gemma-4-31B-it-speculator.dspark |
| Z-Lab | DFlash | z-lab/gemma-4-26B-A4B-it-DFlash, z-lab/Qwen3.8B-DFlash-b16 |
| DeepSeek AI | EAGLE‑3, DFlash, DSpark | deepseek-ai/eagle3_qwen3_8b_ttt7, deepseek-ai/dflash_qwen3_8b_block7 |
| Inferact | EAGLE‑3, DSpark | Inferact/MiniMax-M3-EAGLE3, Inferact/Kimi-K3-DSpark |
新規推測器の訓練（高レベルワークフロー）
- 目的のワークロード（チャット、コード、数学など）に適した代表的なプロンプトを収集。
- 推論時に使用する正確なトークナイザー、チャットテンプレート、サンプリング設定でターゲットモデルの応答を生成。
- 隠れ状態抽出モードを選択 – オンライン（リアルタイム）、オフライン（事前に保存）、ハイブリッド（最初のエポックでキャッシュ）。
- 必要なターゲット層を抽出（例：EAGLE‑3では初期・中間・終盤層；DFlash/DSparkでは選択されたすべての層）。
- ターゲットと同じ語彙、隠れサイズ、トークン埋め込みを持つ推測器を訓練。手法固有のヘッド（並列ブロック、Markov補正など）を含める。
- 受け入れ率、MAL、エンドツーエンドスループットを検証。受け入れ率が低い場合は、プロンプトの混合や訓練ハイパーパラメータを反復調整。
- チェックポイントをパッケージ化し、
--speculative-configを使用してターゲットモデルとともに配信。
今後の方向性
- 学習されていない推測（例：n-gram接尾辞予測）は、非常に繰り返しの多いコード編集ワークロードにおいて学習された推測器を補完できる可能性があります。
- より広範な並行性の研究 – 複数ユーザーのバッチ負荷、異なるバッチサイズ、異なるサンプリング温度での推測的デコードの測定。
- 推測器データの影響 – プロンプトの多様性やドメイン特化の訓練データが、コード、数学、チャット、多言語タスクにおける受け入れに与える影響を体系的に分析。
- より深いプロファイリング – AMD GPU上でドラフト生成、ターゲット検証、KVキャッシュ再利用、グラフ実行オーバーヘッドを分離し、カーネルレベルの最適化を支援。
謝辞
ハードウェアアクセスと協力について、AMDのHongxia YangおよびPeng Sun、およびEmbedded LLMのPin Siang Tan、Jun Kang Chow、Ye Hur Cheongに感謝します。
免責事項
測定は2つのAMD Instinct構成で実施されました：
- ハードウェア1：8 × MI300X GPU（gfx942）とデュアルEPYC 9654 96コアCPU。
- ハードウェア2：8 × MI355X GPU（gfx950）とデュアルEPYC 9575F 64コアCPU（MiniMax-M3-MXFP8用）。 ソフトウェアスタック：Ubuntu 22.04.5 LTS、ROCm 7.2.53211、vLLM 0.23.1rc1.dev1120+g0f0f28b53、PyTorch 2.11.0、Transformers 5.13.1、Python 3.12.13。ドライバ、カーネルバージョン、またはハードウェア構成が異なると結果が異なる可能性があります。
Sources
関連
- Dispatch
- Dispatch
- Dispatch
- Dispatch
- Dispatch