P-EAGLE：vLLM 中的平行推測解碼

TL;DR

P-EAGLE 是一種整合至 vLLM（自 v0.16.0 起）的平行推測解碼方法，消除自回歸草稿的順序瓶頸。透過在單次前向傳播中生成所有 K 個草稿 token，它在使用 NVIDIA B200 GPU 的真實工作負載上相較於原始 EAGLE-3 可達最高 1.69 倍的加速。

解決自回歸草稿瓶頸

標準的推測解碼方法（如 EAGLE）以自回歸方式產生草稿 token，亦即產生 K 個草稿 token 需要對草稿模型執行 K 次順序前向傳播。這會導致延遲隨著推測深度呈線性增長，限制了系統在不讓草稿開銷侵蝕整體效能提升的情況下，能夠多積極地進行推測。

P-EAGLE 透過將草稿過程從自回歸轉為平行，移除這一限制。P-EAGLE 不再採用順序步驟，而是在一次前向傳播中生成整套草稿 token，將草稿 token 的數量與所需的前向傳播次數解耦。

P-EAGLE 架構與機制

P-EAGLE 透過兩個主要步驟產生草稿 token：

步驟 1：預填充

目標模型處理提示詞並產生新 token。在此過程中，P-EAGLE 捕獲內部隱藏狀態：每個提示位置的 h_prompt 與新產生 token 的 h_context。這些隱藏狀態用於指導草稿模型的預測。

步驟 2：平行草稿

草稿模型使用 token 嵌入與隱藏狀態的組合，平行地為每個位置構建輸入：

  • 提示位置： 每個提示 token 的嵌入 emb(p) 與目標模型對應的 h_prompt 配對，並向前移動一個位置，以便預測位置 $i$ 的 token。
  • 下一 token 預測 (NTP)： 第一個位置將新產生的 token 嵌入 emb(new) 與 h_context 配對。
  • 多 token 預測 (MTP)： 對於位置 2 到 K，尚未有 token 嵌入與隱藏狀態時，P-EAGLE 使用兩個可學習參數：共享的遮罩 token 嵌入 emb(mask) 與共享的隱藏狀態 h_shared。它們作為中性佔位符。

所有位置隨後通過 N 個 transformer 層與語言模型頭，同時預測草稿 token $t_1$ 到 $t_K$。

在長序列上的訓練

平行草稿在訓練期間會增加記憶體需求，因為在長度為 N 的序列上訓練 K 個平行組會產生 $N \times K$ 個總位置。例如，當 $N=8,192$ 且 $K=8$ 時，單一訓練樣本包含 65,536 個位置，導致巨大的注意力矩陣（超過 40 億個元素）。

為了解決此問題，P-EAGLE 引入了 序列分割演算法 用於序列內部切分。此演算法將 $N \times K$ 位置序列劃分為連續的區塊，同時在邊界之間保持正確的注意力依賴，並在同一序列的各區塊之間累積梯度。

vLLM 實作細節

將平行草稿整合至 vLLM 需要克服多項與批次中繼資料及記憶體管理相關的技術挑戰：

融合 Triton 核心

平行草稿因加入 MASK 佔位符而打破了草稿與驗證批次形狀之間的一致性。為避免透過多次 GPU 操作重建批次中繼資料的開銷，vLLM 實作了一個融合的 Triton 核心。此核心在一次傳遞中完成複製先前 token ID 與位置、插入額外 token、以 MASK token ID 填充平行草稿槽位，並產生必要的中繼資料（被拒 token 掩碼、被遮罩 token 掩碼，以及隱藏狀態映射）。

隱藏狀態管理

由於隱藏狀態的大小遠大於 token ID，vLLM 使用專用的複製核心，將已學習的隱藏狀態佔位符（parallel_drafting_hidden_state_tensor）廣播至 MASK token 槽位，同時將目標隱藏狀態映射至其新位置。

KV 快取與 CUDA 圖形

  • KV 快取： 被拒 token 會映射至 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 在每輪接受的草稿 token 平均數量 (AL) 上高於 EAGLE-3。於推測深度 $K=7$ 時，P-EAGLE 在 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

Sources