P-EAGLE: vLLM における並列投機的デコード

TL;DR

P-EAGLE は vLLM (v0.16.0 以降) に統合された並列投機的デコード手法で、自己回帰的ドラフトの逐次ボトルネックを排除します。すべての K 個のドラフトトークンを単一のフォワードパスで生成することで、NVIDIA B200 GPU を使用した実際のワークロードで、従来の EAGLE-3 に対して最大 1.69 倍の速度向上を達成します。

自己回帰的ドラフトのボトルネックの解決

EAGLE のような標準的な投機的デコード手法はトークンを自己回帰的にドラフトし、K 個のドラフトトークンを生成するにはドラフトモデルへの K 回の逐次フォワードパスが必要です。これにより、投機深度に比例してレイテンシが線形に増加し、ドラフトのオーバーヘッドが全体の性能向上を食いつぶすことなく、システムがどれだけ積極的に投機できるかが制限されます。

P-EAGLE はドラフトプロセスを自己回帰的から並列へと変換することでこの上限を取り除きます。逐次的なステップの代わりに、P-EAGLE はドラフトトークン全体を単一のフォワードパスで生成し、ドラフトトークン数と必要なフォワードパス数を切り離します。

P-EAGLE のアーキテクチャとメカニズム

P-EAGLE はドラフトトークンを生成するために主に 2 つのステップで動作します：

ステップ 1: プレフィリング

対象モデルはプロンプトを処理し、新しいトークンを生成します。この過程で、P-EAGLE は内部隠れ状態を取得します：各プロンプト位置の h_prompt と新しく生成されたトークンの h_context。これらの隠れ状態がドラフターの予測を導きます。

ステップ 2: 並列ドラフト

ドラフターはトークン埋め込みと隠れ状態の組み合わせを用いて、各位置の入力を並列に構築します：

• プロンプト位置: 各プロンプトトークン埋め込み emb(p) は、対象モデルからの対応する h_prompt と組み合わされ、1 位置シフトされて位置 $i$ のトークン予測を可能にします。
• 次トークン予測 (NTP): 最初の位置では、新しく生成されたトークン埋め込み emb(new) と h_context を組み合わせます。
• マルチトークン予測 (MTP): 位置 2 から K まで、トークン埋め込みと隠れ状態がまだ存在しない場合、P-EAGLE は 2 つの学習可能パラメータを使用します：共有マスクトークン埋め込み emb(mask) と共有隠れ状態 h_shared。これらは中立的なプレースホルダーとして機能します。

すべての位置は N 個のトランスフォーマ層と言語モデルヘッドを通過し、ドラフトトークン $t_1$ から $t_K$ を同時に予測します。

長シーケンスでのトレーニング

並列ドラフトはトレーニング時のメモリ要件を増大させます。長さ N のシーケンス上で K 個の並列グループをトレーニングすると、合計 $N \times K$ の位置が生成されるためです。例えば、$N=8,192$、$K=8$ の場合、単一のトレーニング例は 65,536 の位置を含み、4 億要素を超える巨大なアテンション行列が生じます。

これに対処するため、P-EAGLE は シーケンス分割アルゴリズム を導入し、シーケンス内部で分割を行います。このアルゴリズムは $N \times K$ の位置シーケンスを連続したチャンクに分割し、境界を越える正しいアテンション依存関係を維持しつつ、同一シーケンスのチャンク間で勾配を蓄積します。

vLLM の実装詳細

vLLM に並列ドラフトを統合するには、バッチメタデータとメモリ管理に関するいくつかの技術的課題を克服する必要がありました：

統合 Triton カーネル

並列ドラフトは、追加された MASK プレースホルダーによりドラフトと検証のバッチ形状の一貫性が崩れます。複数の GPU 操作でバッチメタデータを再構築するオーバーヘッドを回避するため、vLLM は統合 Triton カーネルを実装しています。このカーネルは、前のトークン ID と位置のコピー、ボーナストークンの挿入、並列ドラフトスロットへの MASK トークン ID の埋め込み、そして必要なメタデータ（拒否トークンマスク、マスクトークンマスク、隠れ状態マッピング）を単一のパスで生成します。

隠れ状態管理

隠れ状態はトークン ID よりもはるかに大きいため、vLLM は専用のコピーカーネルを使用して学習された隠れ状態プレースホルダー (parallel_drafting_hidden_state_tensor) をマスクトークンスロットにブロードキャストし、対象の隠れ状態は新しい位置にマッピングします。

KV キャッシュと CUDA グラフ

• KV キャッシュ: 拒否されたトークンは PADDING_SLOT_ID (-1) にマッピングされ、不要なキャッシュ書き込みを防止します。
• CUDA グラフ: キャプチャ範囲は $K \times \text{max_num_seqs}$ だけ拡張され、より大きなドラフトバッチサイズに対応します。

パフォーマンスベンチマーク

NVIDIA B200 GPU を使用した GPT-OSS-20B の評価では、P-EAGLE が複数のベンチマーク (MT-Bench、HumanEval、SPEED-Bench) で従来の EAGLE-3 を大幅に上回ることが示されています。

スループット向上

低同時実行数 (c=1) では、P-EAGLE は 55〜69% のスループット向上を提供します。高同時実行数 (c=64) でも 5〜25% の向上が維持されます。具体的な EAGLE-3 に対する速度向上比は次の通りです：

• SPEED-Bench: 最大 1.69 倍 (c=1 時)
• HumanEval: 最大 1.55 倍 (c=1 時)
• MT-Bench: 最大 1.55 倍 (c=1 時)

受容長さ (AL)

P-EAGLE は、ラウンドあたりの受容されたドラフトトークンの平均数 (AL) が EAGLE-3 より高くなります。投機深度 $K=7$ の場合、HumanEval での AL は 30% 高く (3.94 対 3.03)、SPEED-Bench では 31% 高く (3.38 対 2.59) です。

投機深度効率

自己回帰的ドラフターは $K=3$ でスループットが最大になるのに対し、P-EAGLE はすべての同時実行レベルで一貫して $K=7$ でピークスループットを達成します。これは、P-EAGLE が逐次ドラフトに伴う線形レイテンシペナルティなしに、より深い投機から利益を得られることを示しています。

デプロイ

vLLM で並列ドラフトを有効にするには、SpeculativeConfig クラスで "parallel_drafting": true を設定します。事前学習済みの P-EAGLE ヘッドは、以下のモデル向けに HuggingFace で利用可能です：

• GPT-OSS 120B
• GPT-OSS 20B
• Qwen3-Coder 30B