vLLM×AMD GPUでEAGLE-3推論を高速化

大規模言語モデル（LLM）の推論は、自己回帰的なデコーディングによって制約を受けることが増えています。プリフィルが高度に最適化されていても、デコードフェーズではトークンを 1 ステップずつ生成する必要があり、各ステップで通常、対象となる大規模モデル全体を実行する必要があります。Kimi-K2.5 や MiniMax-M2.5 のような、エキスパートを混合した大規模モデルやアテンション計算が重いモデルの場合、この逐次的な処理パターンはサービススループットを制限し、リアルタイムアプリケーションのレイテンシを増加させます。

スペキュレティブ・デコーディング（Speculative Decoding）はこのボトルネックに対処する最も実用的な手法の一つです。これは対象モデルの出力分布を正確に維持しつつ、デコーディング効率を向上させるロスレスな LLM 推論加速技術です。この手法では、より小型または軽量なドラフトモデルを使って複数の未来トークンを提案し、元の対象モデルがそれらのトークンを単一の順方向パスで検証します。ドラフトモデルが予測したトークンが対象モデルも生成するものだった場合、それらをまとめて受け入れることが可能になり、高コストの対象モデルによるデコード反復回数を削減できます。

一般的なスペキュレティブ・デコーディングの手法には、小型ドラフトモデルの使用、マルチトークン予測（MTP）、Medusa スタイルの多头予測、そして EAGLE-3 や DFlash、最近導入された DSpark といった特徴レベルでのドラフト法があります。既存のスペキュレティブ・デコーディング手法の中で、EAGLE-3 は優れたドラフト品質、高い受け入れ率、一貫して競争力のある推論速度向上効果により、特に注目されています。

本ブログでは、AMD Instinct GPU 上で EAGLE-3 のライフサイクル全体を vLLM を基盤に構築し、AMD Quark チームがエンドツーエンドで提供するプロセスをご紹介します。具体的には以下の 3 つのステップです。

1. EAGLE-3 ドラフトモデルのトレーニング：vLLM がターゲットモデルをサービングしてオンポリシーデータを合成し、トレーニング時の隠れ状態を抽出、ループ内での受容評価を実行します。
2. AMD Quark による量子化：ターゲットとドラフトの両方に対して、MXFP4 と FP8 の Day-0 サポートを提供します。
3. ROCm/vLLM を活用した推論加速：AMD Instinct™ MI355X GPU 上で Kimi-K2.5 および MiniMax-M2.5 の推論を InferenceX でベンチマークします。なお、トレーニングセクションの例として使用している MiniMax-M3 EAGLE-3 ドラフトも、同じパイプラインでトレーニングされています。

なぜ Speculative Decoding と EAGLE-3 が重要なのか

標準的な自己回帰的デコーディングでは、ターゲットモデルの 1 ステップごとに 1 トークンを生成します。もしモデルが 1,000 トークンの出力を必要とする場合、事前計算（prefill）後に推論エンジンはおおよそ 1,000 回のターゲットモデルによるデコード反復を実行することになります。これはコストが高い処理です。なぜなら、各デコード反復でモデルの重み、アテンション状態、スケジューラー、そして KV キャッシュ機構すべてにアクセスする必要があるからです。

Speculative Decoding はこのプロセスを変えます：

• ドラフトモデルが複数の候補となる次トークンを提案します。
• ターゲットモデルはそれらの候補を 1 回のパスで検証します。
• ターゲットモデルの分布と一致するトークンは受け入れられます。
• 生成は、最初に拒否されたトークンから、または受け入れられたブロックの末尾から続行されます。

重要なのは「受容率」です。これは、ターゲットモデルがドラフトトークンを何個受け入れることができるかを示す指標です。

ドラフトモデルが次の 4 つのトークンを提案したと想像してください。もしターゲットモデルがこれら 4 つすべてに合意すれば、1 回の検証ステップで 4 つをまとめて受容できます。一方、3 番目のトークンで不一致が生じた場合、最初の 2 つだけが受容され、残りのトークンは再生成する必要があります。

一度に多くのトークンを受容できれば、ターゲットモデルが実行しなければならない高コストなデコードステップが減り、スループットは向上します。重要なのは、ドラフトされたトークンのすべてが発行される前にターゲットモデルによって検証される点です。つまり、推論速度を上げつつも、最終的な出力品質は一切損なわれないのです。

図 1：スペキュレーティブ・デコーディングの仕組み。ドラフトモデルがγ個のトークンを提案し、ターゲットモデルがこれらを単一のパスで検証します。一致する接頭辞（α）が受容され、α+1 個のトークンが発行されます。その後、最初に不一致となったトークンからドラフトを再開します。発行されるすべてのトークンは検証済みであるため、この高速化はロスレスです。

EAGLE は過去数年間で継続的に進化を遂げてきました。当初の EAGLE では特徴レベルでの推測型デコーディングを実現し、EAGLE-2 でドラフトの品質と受容率を向上させました。そして EAGLE-3 では、ターゲットモデルからの多層の特徴を活用することで、さらに精度と高速化を達成しています。

関連性の低い小規模言語モデルに依存するのではなく、ターゲットモデルと密接に連携したドラフトモジュールを訓練します。トレーニング時のテスト手法を採用し、ターゲットモデルの低レベルから中レベル、高レベルまでの意味特徴を組み合わせることで、検証器がより受け入れやすい候補をドラフトモデルが提案できるようになります。

実運用における推論において重要なのは、シンプルです。EAGLE-3 は、ターゲットモデルの出力挙動を保証する検証プロセスを通じて、生成スループットを向上させることができます。

AMD Quark MXFP4: メインストリーム LLM 向けの Day-0 量子化

MXFP4 は、Open Compute Project（OCP）が策定したマイクロスケーリング 4 ビット浮動小数点フォーマットです。4 ビットの要素を小さなブロックにグループ化し、スケールファクターを共有することで、メモリ使用量は INT4 に近い水準に抑えつつ、はるかに優れた数値的振る舞いを維持します。

AMD Instinct MI350 シリーズ GPU（MI350X/MI355X）はネイティブの FP4 行列演算をサポートしており、MXFP4 の重みをそのままハードウェアにマッピングできます。これにより、大規模な混合专家（MoE）デコーディングにおいてボトルネックとなるメモリ帯域幅と容量への負荷を大幅に軽減します。

AMD Quark は AMD が提供するモデル量子化ツールキットです。Quark チームは、主要な大規模言語モデル（LLM）に対して、Hugging Face で「day-0」の MXFP4 量子化チェックポイントを提供しています。具体例としては amd/Kimi-K2.5-MXFP4 や amd/MiniMax-M3-MXFP4 などがあります。

ここでいう「day-0」とは、主要モデルがリリースされた際に、サードパーティによる量子化の対応を待たずに、ROCm や vLLM でそのまま動作するハードウェア最適化版（MXFP4 および FP8）をすぐに提供することを意味します。これらのチェックポイントは、EAGLE-3 のドラフトトレーニングおよび推論における目標モデルとして、すぐに直接使用可能です。

これらのチェックポイントは、vLLM 上で ROCm を介して直接利用されます。具体的には、VLLM_ROCM_USE_AITER_FP4_ASM_GEMM=1 という環境変数を設定することで FP4 ASM GEMM パスを経由し、AITER の MoE カーネルも活用されます。これにより、ユーザーは MXFP4 によるメモリ削減効果と、本番レベルのスループットを同時に享受できます。

スペキュレーティブ・ディコーディング（推論の事前予測）はロスレスです。生成されたドラフトトークンはすべて、提供される目標モデルによって検証されるため、目標モデルが出力する分布には一切影響を与えません。（図 2 参照）

*図 2: AMD Quark の day-0 量子化フロー — 主要な LLM がリリースされると、Quark チームは Hugging Face に MXFP4 または FP8 のビルド版を公開し、ROCm 上の vLLM がこれを直接提供します。*

vLLM を用いた EAGLE-3 ドラフトモデルのトレーニング

スペキュレーティブ・ディコーディングを高速化するには、高受容率のドラフトモデルが不可欠であり、その訓練はモデリングの問題であると同時にシステム設計上の課題でもあります。当社のパイプラインでは、vLLM は推論エンジンとしてだけでなく、訓練の中核としても機能します。ここでは、AMD Instinct GPU で訓練した MiniMax-M3 EAGLE-3 ドラフトを事例として解説します。（なお、以下の推論結果に登場する Kimi-K2.5 や MiniMax-M2.5 の EAGLE-3 ドラフトは、当社が訓練したものではなく、Hugging Face から公開されているコミュニティ製のドラフトです。）（Figure 3）

*Figure 3: vLLM を中核とした EAGLE-3 パイプライン。単一の vLLM-on-ROCm ランタイムが、5 つのステージにわたる一連のプロセスを統括します。AMD Quark MXFP4/FP8 ターゲットモデルを用いてオンポリシーデータを合成する（Stage 1）、ターゲットの隠れ状態を FSDP2 ベースの冷間開始ドラフト訓練器へストリーミングし、ループ内のサービス評価を実行して最良のチェックポイントを選定する（Stage 2）。その後、ドラフトモデルをエクスポート（Stage 3）し、EAGLE-3 のスペキュレーティブ・ディコーディング用にデプロイ（Stage 4）、さらに受容長と GPU ごとのスループットを評価（Stage 5）します。

オンポリシーデータ合成は vLLM によって提供されます。EAGLE-3 のドラフトモデルは、ターゲットの分布内にあるデータから最も効果的に学習します。そこで AMD Quark MXFP4 ターゲットを vLLM-ROCm サーバーとして構築し、そこからオンポリシーな応答を生成します。チャット用には /v1/chat/completions エンドポイント（正確なサービングチャットテンプレートを使用）と、非チャットや分布外での堅牢性を確保するための生テキスト用 /v1/completions エンドポイント（テンプレートをバイパス）の両方に対応しています。後でサービス提供時に使用するのと同じエンジンとテンプレートでデータを生成することで、トレーニングとサービングの一貫性を保つことができます。

隠れ状態の抽出も vLLM が担います。EAGLE-3 は、無関係な小規模モデルではなく、ターゲット内部の特徴——低・中・高レベルの隠れ状態と fc_norm ——を条件としてドラフトを生成します。vLLM の隠れ状態抽出フックにより、実行中のターゲットエンジンからこれらの補助レイヤーを直接取得できます。サポートされているモードは 3 つあり、相互に切り替え可能です。1 つ目はオンラインモードで、ターゲットがトレーニング環境と同一ノード上に配置される場合です。2 つ目はオフラインモードで、隠れ状態をディスクにダンプします。3 つ目はストリーミングモードで、生 vLLM サーバーからディスクを経由せずに隠れ状態をトレーニング側にストリームします。このストリーミング機能こそが、単一ノード上で 420B MXFP4 MoE ターゲットのトレーニングを実用的なものにしている要因です。

コールドスタート FSDP2 トレーニングでは、単層構成の EAGLE-3 ドラフトヘッドをゼロから学習させます。トレーニング時のテスト（TTT）損失と位置減衰重み付けを用い、FSDP2 環境下で実行します。検証器として AMD Quark MXFP4 ターゲットを使用するため、ドラフトモデルはデプロイ時に直面する活性化空間に対して正確に学習を行うことになります。

vLLM を用いたループ内評価を再実施します。トレーニング中の損失値は実際のトークン受容率を過大評価するため、定期的に現在のチェックポイントをエクスポートし、vLLM によるスペキュレーティブ・デコーディングでサーブして真の受容長を測定します。その後、このサーブ指標に基づいて最良のチェックポイントを選択します。本番環境で稼働するエンジンと、ドラフトモデルを選定するエンジンは同一です。

vLLM でのエクスポートとデプロイでは、選択されたドラフトモデルを Hugging Face フォーマットに変換し、vLLM 対応のドラフトディレクトリに統合します。その後、vLLM-ROCm 環境で EAGLE スペキュレーティブ・デコーディングを実行してデプロイします。これは次節で詳細を測定するパスと完全に一致しています。

vLLM がデータ生成、隠れ状態の抽出、ループ内評価、そしてサービングのすべてを担っているため、ドラフトモデルは本番環境のエンジンに対してエンドツーエンドでトレーニングおよび検証されます。これが、トレーニング中に測定した受容率が本番環境でも有効に機能する理由です。

SPEED-Bench におけるドラフト品質：11 のドメインと長いコンテキスト

トレーニング済みの MiniMax-M3 EAGLE-3 ドラフトモデルを、スペキュレーティブ・デコーディング用のマルチドメインベンチマークである SPEED-Bench で評価します。評価指標は受容長（AL）です。これはターゲット検証ステップごとに生成されるトークンの平均数を示し、数値が高いほど性能が良いことを意味します。AL が 1 の場合は速度向上がない状態を指します。

ドメイン別の受容長（SPEED-Bench 定性評価）：

11 のドメイン全体での平均 AL は 2.77、つまりターゲットステップあたり約 2.8 トークンが採用されています。これは構造化された技術コンテンツにおいて最も顕著で、コーディング（3.16）、数学（3.12）、RAG（3.11）、多言語（3.07）の順です。一方で、あらゆるドラフトモデルにとって予測が最も困難なオープンエンドな文章作成やロールプレイでも AL 2.0〜2.9 を維持できています。

さらに重要なのは、プロンプト長を 1K トークンから 32K トークンに増やしても採用長がほぼ横ばい（2.64 から 2.63）である点です。つまり、長いコンテキストでも速度向上効果が減衰しないということです。

推測トークンを 3 つ使用した場合、1 番目、2 番目、3 番目のドラフト位置の採用率はそれぞれ約 75%、55%、41% です（累積値）。これらの結果は、vLLM を中核とした私たちのレシピが実を結んだものです。具体的には、ターゲットモデルを通じて生成されたオンポリシーデータ、ターゲットモデル自身の特徴量に基づく隠れ状態の教師信号、AMD Quark MXFP4 検証器に特化したコールドスタート学習、そして実際のサービス環境での採用率に基づいたチェックポイント選定です。（Figure 4）

*Figure 4: SPEED-Bench における MiniMax-M3 EAGLE-3 の採用長は、1K から 32K のコンテキストにわたってほぼ横ばい（1K で 2.64、32K で 2.63）であり、長いプロンプトでも速度向上効果が減衰しません。*

One Team, End to End: The AMD Quark Advantage

この取り組みの目立った特徴は、ターゲットモデルとドラフトモデルの両方に対するスタック全体を「AMD Quark」という単一のチームが所有している点です。

ターゲットモデル：day-0 対応の MXFP4/FP8 量子化および ROCm/vLLM での展開。

ドラフトモデル：EAGLE-3 のトレーニング、FP8/MXFP4 量子化、ROCm/vLLM での展開。

エンドツーエンド統合：オンポリシーデータ合成、隠れ状態の抽出、サービス評価、エクスポート、そして推測的デコーディングによるサービングまで、すべて vLLM を通じて連携され、一体で検証されています。

このワンストップ体制により、AMD Instinct ユーザーは、量子化されたターゲットモデル、高い採択率を持つ対応するドラフトモデル、そしてチューニング済みの vLLM 推測的デコーディング展開を、単一のチームから受け取ることができます。これは、CUDA に依存したツールを手動で組み立てる必要がなく、新モデルに対する day-0 サポートも含まれています。

アクセラレーション結果

以下のドラフト結果セクションでは、1K/1K ワークロード（ISL=1024、OSL=1024）の結果のみを記載します。スピードアップは、EAGLE-3 のスループットを対応する推測的デコーディングなしのベースラインのスループットで割ることで計算されています。

Kimi-K2.5 の結果は、AMD Instinct MI355X（TP=4）を使用し、ランダムプロンプトに対して num_prompts=10 x concurrency、num_warmups=2 x concurrency で実行され、各セルあたり 10 シードずつ使用しています。Kimi のチャートでは、BF16 と FP8 のドラフトパスを併記しており、BF16 の vLLM v0.19.0 スイープでは MML=2248、FP8 のスイープでは MML=2304 を使用しています。ここでいう MML（max-model-len）は最大コンテキスト長であり、vLLM モデルが単一のリクエストで処理できるトークンの総数（プロンプト＋生成出力）を指します。

Kimi K2.5 EAGLE-3：BF16 および AMD Quark FP8 ドラフト

Docker イメージ：BF16 スイープでは vllm/vllm-openai-rocm:v0.19.0（MML=2248）を、FP8 スイープでは vllm/vllm-openai-rocm:nightly-fb1ac806c55a6dc96fe92261b80c8550e9c39d2f（MML=2304）を使用します。

対象モデル：amd/Kimi-K2.5-MXFP4。BF16 ドラフトモデルは lightseekorg/kimi-k2.5-eagle3、FP8 ドラフトモデルは amd/kimi-k2.5-eagle3-fp8 です。後者は AMD Quark の FP8 メタデータで量子化されており、BF16 版のターゲット言語モデル（LM）ヘッドと共有しています。

この構成では、FP8 ドラフトパスは vLLM の RowWiseTorchFP8ScaledMMLinearKernel を経由します。具体的には、hipBLASLt の行ごとのスケーリングされた FP8 GEMM に対する torch._scaled_mm を使用し、AITER の事前シャッフル済み FP8 パスとは異なります。一方、ターゲットの MXFP4 モデルは、VLLM_ROCM_USE_AITER_FP4_ASM_GEMM=1 を設定することで、ROCm の FP4 ASM パスを利用します。

*Figure 5: Kimi-K2.5 EAGLE-3 output throughput (tok/s/GPU) at 1K/1K on AMD Instinct MI355X (TP=4). Both the BF16 and AMD Quark FP8 draft paths beat the no-speculative baseline (1.69x-1.90x and 1.76x-2.00x respectively); the gain is largest at low concurrency.*

MiniMax M2.5 BF16 EAGLE-3

Docker image: vllm/vllm-openai-rocm:nightly-4eafc729285e459a5fc96efd6f7b313b155cad48

対象モデル：MiniMaxAI/MiniMax-M2.5。ドラフトモデル：thoughtworks/MiniMax-M2.5-Eagle3。設定は BF16 ドラフトパスで num_speculative_tokens=3、draft_tensor_parallel_size=1 です。以下の数値は 1K/1K のランダムプロンプトを使用し、TP=4 および EP を有効化して、並行度ごとにシードを 5 つずつ設定した結果です。MiniMax の再現実験と同じく、ベースラインと EAGLE-3 を比較しています。

*図 6：MiniMax-M2.5 の EAGLE-3 出力スループット（tok/s/GPU、1K/1K、TP=4）。BF16 ドラフトパスは、推測なしのベースラインと比較して 1.38 倍から 1.79 倍の性能を示し、特に並行度が低い環境でその差が顕著です。

Kimi-K2.5 の 1K/1K スイープ全体では、BF16 の EAGLE-3 ドラフトパスは推測なしベースラインに対して 1.69 倍から 1.90 倍のスループットを達成し、AMD Quark の FP8 EAGLE-3 ドラフトパスでは 1.76 倍から 2.00 倍の向上が見られました（図 5）。MiniMax-M2.5 においては、BF16 の EAGLE-3 ドラフトパスが並行度 64 から 4 の範囲でベースラインに対し 1.38 倍から 1.79 倍のスループットを提供します（図 6）。Kimi の FP8 ドラフトパスは、現在統合されている RowWise hipBLASLt FP8 ドラフトカーネルパスを使用しています。

まとめ

EAGLE-3 によるスペキュレーティブ・ディコーディングは、AMD Instinct GPU 上で損失なく大幅なスループット向上を実現します。1K/1K のスイープテストでは、Kimi-K2.5 で 1.69 倍から 2.00 倍、MiniMax-M2.5 では最大 1.79 倍の性能向上を確認しました。

この手法が実用的なエンドツーエンドのソリューションとなる理由は、3 つの要素の組み合わせにあります。まず、ターゲットモデルとドラフトモデルの両方に対して AMD Quark が提供する日付ゼロ（day-0）の MXFP4/FP8 量子化を採用している点です。次に、オンポリシーデータの合成、隠れ状態の抽出、そして実際のサービスにおける受け入れ率に基づいたチェックポイント選定を行う vLLM を中核としたトレーニングパイプラインです。最後に、ROCm と vLLM を活用したスペキュレーティブ・ディコーディングによる推論実行です。

ターゲットモデルとドラフトモデルの量子化、トレーニング、デプロイメントをすべて AMD Quark チームが一元管理しているため、AMD Instinct ユーザーは手作業で組み合わせたようなものではなく、一貫性のある日付ゼロ対応のスペキュレーティブ・ディコーディングスタックを利用できます。

謝辞

本稿は、EAGLE-3 のスペキュレーティブ・ディコーディングに対する ROCm バックエンドのカバー範囲拡大、再現可能なトレーニングと推論の実行レシピ、すぐに使えるドラフトモデルの提供、そして上流プロジェクトへの連携に向けた明確な道筋を示すことで、AMD 上の vLLM エコシステムを強化するものです。

AMD Quark チーム、AMD ROCm と vLLM の貢献者、InferenceX のメンテナーおよびレビュアー、そして EAGLE-3 の研究コミュニティに感謝いたします。特に、InferenceX における EAGLE-3 ベンチマーク統合への貢献に対して、Andy Luo 氏率いるチームの Haichen Zhang さん、Chun Fang さん、Chang Liu さんに心よりお礼申し上げます。

その他のリソース

• EAGLE-3 プロジェクト

• AMD Quark

• vLLM