AMD GPU上のvLLMで学ぶ投機的デコーディング

導入

自己回帰型の生成では、モデルは1トークンを出力するたびに、それを系列へ追加して次のトークンを計算します。この方式は単純で安定していますが、長い応答ではターゲットモデルのデコードを何度も順番に実行しなければならず、遅延とスループットの制約になります。vLLMの新しい実践ガイドは、ROCm上のAMD Instinct MI300XおよびMI355Xを対象に、投機的デコーディングの仕組みと調整方法を整理しています。

ドラフトと検証

投機的デコーディングは元の大規模モデルを置き換えません。元のモデルはターゲットモデルとして残り、最終的な出力を決めます。まず軽量なドラフトコンポーネントが複数の将来トークンを提案し、その後ターゲットモデルが1回の検証処理で候補列を左から順に確認します。

候補が連続して受理されれば、複数トークンをまとめて確定できます。途中で拒否が発生した場合、その後ろの候補は破棄され、ターゲットモデルが次の進行を決めるトークンを出力します。したがって、効果の本質は候補を作ることではなく、1回の検証でどれだけ有効な出力を進められるかにあります。

vLLMで扱われる方式

記事では、方式を大きく3つのグループに整理しています。

• ネイティブMTP：ターゲットモデルのアーキテクチャに組み込まれた補助予測経路を使い、候補を逐次生成します。
• 分離型MTPドラフター：対応する別チェックポイントを使い、ターゲットの活性値や共有KVキャッシュを利用して候補を作ります。
• ターゲット条件付きネットワーク：EAGLE-3、DFlash、DSparkが含まれます。EAGLE-3は隠れ状態から自己回帰的に提案し、DFlashは候補ブロックを並列に作り、DSparkは軽量な因果補正と信頼度に基づくプレフィックス選択を加えます。

これはターゲットモデルの分類ではなく、ドラフト側の分類です。1つのターゲットモデルがネイティブMTPと、別途学習されたEAGLE-3、DFlash、DSparkを併用できる場合もあります。ドラフターが受け取る情報も方式ごとに異なり、隠れ表現、複数層の状態、KVキャッシュ、またはそれらを組み合わせた特徴が使われます。

評価で見るべき点

提案するトークン数を増やせば、常に速くなるわけではありません。候補が多く拒否されれば、ドラフト生成と検証のコストが増え、まとめて確定する利点が小さくなります。記事の測定でも、出力トークンのスループットはドラフト方式、提案長、モデル系列、ドラフトチェックポイント、ワークロード、受理挙動によって変化します。

そのため、実運用では受理率や確定トークン数を観測し、通常の自己回帰デコードを基準に方式ごとの比較を行うことが重要です。固定設定をそのまま適用するより、対象モデルと実際のリクエスト特性に合わせて段階的に調整する方が安全です。

意義

AMD GPU上のvLLMにおける投機的デコーディングは、単一の高速化手法ではなく、モデル適合性、候補生成コスト、並列性を組み合わせて選ぶシステム最適化です。最終的な判断材料は「何トークン提案できるか」ではなく、「ターゲットモデルが安定して何トークン受理できるか」です。

出典：vLLM Blog