EAGLE-3 スペキュラティブ デコーディング on AMD Instinct GPUs

EAGLE-3 スペキュラティブ デコーディング on AMD Instinct GPUs

vLLM と AMD Quark チームは、AMD Instinct GPU 上での EAGLE-3 スペキュラティブ デコーディングのエンドツーエンド パイプラインを導入し、大規模な mixture-of-experts (MoE) と注意重視モデルに対する損失のない推論高速化を可能にしました。この統合により、AMD Instinct MI355X GPU 上で Kimi-K2.5 に対して 1.69x から 2.00x のスループット向上、MiniMax-M2.5 に対して 1.38x から 1.79x のスループット向上が達成されます。

スペキュラティブ デコーディングと EAGLE-3 アプローチ

スペキュラティブ デコーディングは、高コストなターゲットモデルのデコードイテレーション数を削減する損失のない加速技術です。トークンを一つずつ生成する代わりに、軽量なドラフトモデルが複数の将来のトークンを提案し、ターゲットモデルが単一のフォワードパスでそれを検証します。ターゲットモデルの分布と一致するトークンは受け入れられ、最初に拒否されたトークンから生成が続行されます。

EAGLE-3 は、ターゲットモデルのマルチレイヤー特徴を活用することで、以前のバージョンを改善します。無関係な小規模言語モデルを使用する代わりに、EAGLE-3 はターゲットモデルと密接に整合したドラフトモジュールを訓練し、低レベル、中レベル、高レベルのセマンティック特徴を利用してドラフトの受容率と全体的な推論速度を向上させます。

AMD Quark MXFP4 量子化

大規模な MoE モデルのメモリ帯域幅と容量の制約に対処するため、AMD Quark チームは MXFP4 (OCP Microscaling 4-bit floating-point) 形式を使用して主流の LLM に対して「day-0」量子化を提供します。

MXFP4 実装の主要な技術的詳細には以下が含まれます：

  • ハードウェア アクセラレーション: AMD Instinct MI350シリーズ GPU (MI350X/MI355X) はネイティブな FP4 マトリックス アクセラレーションを提供します。
  • メモリ効率: MXFP4 は共有スケールファクターを持つ小さなブロックに 4 ビット要素をグループ化し、メモリフットプリントを削減しながら INT4 より優れた数値動作を維持します。
  • vLLM 統合: これらのチェックポイントは、FP4 ASM GEMM パス (VLLM_ROCM_USE_AITER_FP4_ASM_GEMM=1) と AITER MoE カーネルを介して ROCm 上の vLLM で消費されます。

vLLM中心の EAGLE-3 トレーニング パイプライン

高い受容率を持つドラフトモデルのトレーニングは、5つのステージにわたるトレーニングループの中心として vLLM が機能するシステム問題として扱われます。

  1. オンポリシー データ合成: vLLM は、トレーニングとデプロイメントの間の一貫性を確保するために、正確なサービング チャット テンプレートを使用して AMD Quark MXFP4 ターゲットを提供し、応答を生成します。
  2. 隠れ状態抽出: vLLM の隠れ状態抽出フックは、実行中のターゲットエンジンから直接、補助レイヤー（低レベル、中レベル、高レベルの隠れ状態および fc_norm）を公開します。これにより、オンライン、オフライン、ストリーミング モードがサポートされ、後者はシングルノードで 420B MXFP4 MoE ターゲットのトレーニングを可能にします。
  3. コールドスタート FSDP2 トレーニング: シングルレイヤーの EAGLE-3 ドラフトヘッドは、FSDP2 の下でトレーニングタイムテスト (TTT) 損失とポジションデカイウェイティングを使用してスクラッチからトレーニングされます。
  4. ループ内 サーブ-エバリュエーション: 現在のチェックポイントは定期的にエクスポートされ、vLLM スペキュラティブ デコーディングの下で提供され、真の受容長を測定します。これにより、選択されたチェックポイントが本番エンジンに最適化されることを保証します。
  5. エクスポートとデプロイメント: 最終ドラフトは Hugging Face 形式でエクスポートされ、vLLM-ROCm を介してデプロイされます。

ドラフトの品質とコンテキスト スケーリング

SPEED-Bench の評価によると、MiniMax-M3 EAGLE-3 ドラフトは 11 ドメイン全体で平均受容長 (AL) が 2.77 です。最高の受容率は構造化された技術コンテンツで見られます：

  • コーディング: 3.16 AL
  • 数学: 3.12 AL
  • RAG: 3.11 AL
  • マルチリンガル: 3.07 AL

注目すべきことに、プロンプト長が 1K から 32K トークンに増加しても（2.64 から 2.63）、受容長は安定しており、長いコンテキストでもスピードアップが低下しないことを示しています。

AMD Instinct MI355X 上のパフォーマンス ベンチマーク

AMD Instinct MI355X (TP=4) で 1K 入力および 1K 出力トークン (1K/1K) を使用して実施されたベンチマークは、非スペキュラティブ ベースラインと比較して顕著なスループット向上を示しています：

Kimi-K2.5 結果

  • BF16 ドラフト パス: スループットが 1.69x から 1.90x 向上。
  • AMD Quark FP8 ドラフト パス: スループットが 1.76x から 2.00x 向上。

MiniMax-M2.5 結果

  • BF16 ドラフト パス: スループットが 1.38x から 1.79x 向上。

両モデルにおいて、パフォーマンスの向上は低い並行度レベルで最も顕著です。

Sources