7月3日、Larry Li、Chun Fang、Chao Li、Xikai Meng、Andy Luo、Haichen Zhang、Bowen Bao、Spandan Tiwari、Ashish Sirasaoが「Accelerating Large-Scale LLM Inference on AMD Instinct MI350X/MI355X with Eagle3 and AMD Quark」と題した記事を公開した。この記事では、AMD Instinct MI355X GPU上でEagle3投機的デコードとAMD Quark FP8量子化を組み合わせ、Kimi-K2.5やMiniMax-M2.5といった大規模LLMの推論スループットを大幅に改善する技術的な取り組みについて詳しく紹介されている。
なぜ今、投機的デコードなのか
LLM推論のボトルネックは、トークンを1つずつ生成する自己回帰デコードにある。1000トークン生成するには約1000回のターゲットモデル推論が必要で、KVキャッシュやアテンション計算が毎ステップ走る。特にKimi-K2.5やMiniMax-M2.5のようなMoE(Mixture of Experts)+アテンション重視の大規模モデルでは、このコストが顕著だ。
投機的デコード(Speculative Decoding)はこの問題に対する実用的な解法だ。軽量なドラフトモデルが複数の次トークン候補をまとめて提案し、ターゲットモデルが1回のフォワードパスでまとめて検証する。ドラフトが正解していれば複数トークンを一度に確定でき、デコードステップ数が減る。誤りがあっても最初の不一致トークンから再生成するため、出力品質はターゲットモデルのそれと一致することが保証される(ロスレス)。
Eagle3とは
投機的デコードの手法はいくつか存在するが(小型ドラフトモデル、Medusa、MTPなど)、Eagle3はその中でも受容率と実効スループット向上の両面で優れた手法とされている。
Eagle3の特徴は、ドラフトモジュールをターゲットモデルの内部表現に密接に結びつけている点だ。ターゲットモデルの低・中・高レベルの意味特徴を組み合わせて学習することで、ターゲットが承認しやすいトークン候補を提案できる。無関係な小型LLMをドラフトに使う手法より受容率が高くなる。
AMDチームが取り組んだ3つの課題
今回の記事はAMD Quarkチームによる実装報告だ。取り組みは以下の3点に整理される。
1. vLLMバックエンドの修正(最重要)
ROCm固有の高速アテンション実装であるAITER MLAバックエンドとEagle3を同時に使おうとすると、KVキャッシュのブロックサイズ非互換が発生し、どちらかを諦める必要があった。
vLLM PR #39616でこの問題を解決した。ブロックサイズ制約をカーネルの内部トークンインデックスから切り離し、マルチトークン検証時のメタデータ処理を簡素化することで、AITER MLAとEagle3の両立を実現した。このPRはupstreamにマージ済みだ。
2. AMD Quark FP8によるドラフトモデル量子化
ドラフトモデルは「軽くて精度が高い」という相反する要求を満たす必要がある。重すぎればドラフトのオーバーヘッドがスループット改善を相殺し、精度が低すぎれば受容率が下がる。
AMD QuarkでKimi-K2.5のEagle3ドラフトモデルをFP8量子化したモデル(amd/kimi-k2.5-eagle3-fp8)をHugging Face上で公開している。LMヘッドは量子化せずBF16のまま保持し、Quark FP8メタデータをモデルconfigに格納する実装だ。FP8ドラフトパスはtorch._scaled_mm(hipBLASLt rowwise scaled FP8 GEMM)を経由する。
3. InferenceXベンチマークへの組み込み
SemiAnalysisのInferenceXベンチマークにEagle3設定を追加した。対応PRは以下の通りだ。
| 対象 | PR |
|---|---|
| Kimi-K2.5 BF16 Eagle3 | #1116 |
| MiniMax-M2.5 BF16 Eagle3 | #1234 |
| Kimi-K2.5 FP8 Eagle3 | #1515 |
ベンチマーク設定は1K/1K、1K/8K、8K/1K(プロンプト/出力トークン長)のパターンと、TP/EP(テンソル並列/エキスパート並列)= 4/4 および 8/8 の組み合わせをカバーする。
実測スループット
測定条件は1K/1K(入力1024トークン、出力1024トークン)、AMD Instinct MI355X、TP=4。
Kimi-K2.5(ターゲット: amd/Kimi-K2.5-MXFP4)
| 同時リクエスト数 | 非投機的デコード (tok/s/GPU) | BF16 Eagle3 | FP8 Eagle3 |
|---|---|---|---|
| 4 | 82.7 | 157.0(1.90x) | 165.2(2.00x) |
| 8 | 142.2 | 269.1(1.89x) | 270.1(1.90x) |
| 16 | 220.5 | 399.6(1.81x) | 412.7(1.87x) |
| 32 | 342.2 | 627.6(1.83x) | 633.8(1.85x) |
| 64 | 533.3 | 901.6(1.69x) | 936.6(1.76x) |
FP8ドラフトが低並列時に最大2.00xを達成している。BF16ドラフトと比較してもFP8の方がわずかに上回るケースが多く、量子化によるドラフト効率改善の効果が見える。
MiniMax-M2.5(ターゲット: MiniMaxAI/MiniMax-M2.5)
| 同時リクエスト数 | 非投機的デコード (tok/s/GPU) | BF16 Eagle3 |
|---|---|---|
| 4 | 90.4 | 161.7(1.79x) |
| 8 | 162.6 | 277.9(1.71x) |
| 16 | 282.7 | 436.9(1.55x) |
| 32 | 484.8 | 759.0(1.57x) |
| 64 | 822.0 | 1137.3(1.38x) |
8K/1Kでスピードアップが落ちる理由
入力が長い(8K/1K)シナリオでは投機的デコードの効果が薄れる。記事では定量的な分析が示されている。
ドラフトの受容率や品質は変わらない。問題はターゲットモデルの検証コストだ。1K入力時の検証コストが約0.30ms/トークンだったのが、8K入力では約2.56ms/トークン(約8.5倍)に跳ね上がる。長いKVキャッシュへのアテンション計算とTP/MoE通信が原因だ。
その結果、ドラフト+検証サイクルに占める検証コストの割合が、1K時の約14%から8K時の約59%まで上昇する。検証コストが支配的になると、複数トークンをまとめて確定する利点が薄れる。
出力品質について
投機的デコードはターゲットモデルによるトークン検証を前提とするため、理論上は出力分布が変わらない。記事ではGSM8K(数学推論タスク)での評価を参照し、Eagle3適用前後で精度劣化が観測されなかったことを確認している。
詳細はAccelerating Large-Scale LLM Inference on AMD Instinct MI350X/MI355X with Eagle3 and AMD Quarkを参照していただきたい。