vLLM 在 AMD GPU 上的預測性解碼:效能與方法

簡要重點

vLLM 中的預測性解碼讓快速的草稿元件提出多個未來的 token,由目標模型在一次傳播中驗證,而在 AMD Instinct MI300X/MI355X GPU 上,此技術對某些模型-方法組合可帶來最高約 2.9× 的吞吐量提升,同時保持原始模型的輸出行為。


vLLM 中預測性解碼的工作原理

預測性解碼在原始(目標)大語言模型前新增了一個「草稿與驗證」階段。草稿元件產生一連串候選 token;目標模型隨即在單一前向傳播中評估整個候選區塊。被目標模型接受的 token 會被確認,而第一個被拒絕的 token 則由目標模型自身的輸出取代。此過程重複進行,直到生成完成。

關鍵特性

  • 目標模型仍是唯一真實來源;在驗證完成前不會輸出任何 token。
  • 一次目標模型傳播可確認多個草稿 token,減少高成本目標傳播次數。
  • 接受性評估從左到右進行;一旦拒絕,整個候選區塊的後續處理即中止。

評估的草稿方法

vLLM 支援五種具體的草稿方法,每種方法在草稿網路如何接收目標模型資訊,以及是否以序列或平行方式產生 token 上有所不同。

方法 草稿元件 使用的目標模型資訊 token 產生方式
原生 MTP 模型內建的輔助預測路徑 目標模型的隱藏表示(或前一 MTP 步驟) 序列式(重複 MTP 呼叫)
Gemma 4 MTP 與目標模型配對的獨立 MTP 檢查點 目標激活 + 共享 KV 快取 序列式
EAGLE‑3 專用的自回歸預測器 早期、中期與晚期目標層的融合隱藏狀態 序列式,每個草稿 token 依賴前一個
DFlash 專用的平行預測器 目標隱藏狀態投影至每個草稿層皆可存取的 KV 對 所有位置在一次前向傳播中同時預測
DSpark DFlash 主幹 + 輕量級馬可夫頭 與 DFlash 相同的目標上下文 平行主幹 + 序列式修正 token 選擇

原生 MTP

內建於目標模型中,MTP 使用輔助頭預測固定數量的未來 token。草稿 token 一個接一個產生,每一步都消耗前一次 MTP 輸出的隱藏狀態。當 num_speculative_tokens 超過原生深度時,vLLM 會透過額外的前向傳播重複使用 MTP 路徑。

Gemma 4 MTP

一個獨立的檢查點(即「助手」模型)與目標模型共享 KV 快取與激活,使其能重用已計算的上下文。草稿 token 仍以序列方式產生,但由於草稿元件重用目標張量,記憶體開銷相當低。

EAGLE‑3

EAGLE‑3 訓練一個預測器,接收三個目標層隱藏狀態的融合表示。該預測器以自回歸方式運作:第一個草稿 token 結合融合的目標特徵與採樣的 token 嵌入;後續 token 則依賴前一個草稿輸出。這產生與提案長度成正比的序列草稿工作。

DFlash

DFlash 以平行方式預測整個未來位置區塊。一個已由目標模型驗證的「錨點」token 開始區塊,其餘被遮蔽的位置同時填入。目標隱藏狀態被投影為額外的 KV 對,讓每個草稿層都能存取,從而維持整個平行傳播中的目標上下文。

DSpark

DSpark 基於 DFlash,加入輕量級馬可夫頭,於平行主幹後引入 token 間的依賴關係。主幹為所有位置產生基礎 logits;馬可夫頭則利用先前選定的草稿 token 調整每個位置,實現平行速度與序列一致性之間的混合。


在 vLLM 中啟用預測性解碼

預測性解碼透過 --speculative-config 標記進行設定。JSON 負載指定方法、可選的草稿檢查點,以及 num_speculative_tokens(提案長度)。

# 原生 MTP(無額外檢查點)
vllm serve <target-model> \
  --speculative-config '{"method": "mtp", "num_speculative_tokens": 4}'

# Gemma 4 MTP(需助手檢查點)
vllm serve <target-model> \
  --speculative-config '{"method": "mtp", "model": "google/gemma-4-26B-A4B-it-assistant", "num_speculative_tokens": 4}'

# EAGLE‑3
vllm serve <target-model> \
  --speculative-config '{"method": "eagle3", "model": "RedHatAI/<target>-speculator.eagle3", "num_speculative_tokens": 3}'

# DFlash(平行區塊)
vllm serve <target-model> \
  --speculative-config '{"method": "dflash", "model": "z-lab/<target>-DFlash", "num_speculative_tokens": 15}'

# DSpark(平行 + 馬可夫修正)
vllm serve <target-model> \
  --speculative-config '{"method": "dspark", "model": "RedHatAI/<target>-speculator.dspark", "num_speculative_tokens": 7}'

記憶體考量 – 原生 MTP 與目標模型共享權重,因此額外 GPU 記憶體開銷極小。其他所有方法都會載入獨立的草稿檢查點;所需空間取決於檢查點大小、精度(FP16/INT4)、張量平行度與執行時緩衝區。


基準測試方法

作者在 AMD Instinct MI300X 與 MI355X GPU 上,使用 ROCm™ 堆疊測量 輸出 token 吞吐量(token/秒)。基準測試使用任務導向資料集(GSM8K、MATH500、HumanEval、MBPP)以反映真實的接受模式。針對每個目標-方法組合,他們掃描 num_speculative_tokens(N),並記錄:

  • 吞吐量比率(預測性 / 基準)
  • 平均接受長度(MAL) – 每次驗證回合中平均接受的草稿 token 數量
  • 接受率(AR) – 通過驗證的提案 token 比例

所有測量值皆相對於相同硬體與軟體堆疊上執行的非預測性自回歸基準。


主要效能發現

觀測到的最高加速比

目標模型 方法 資料集 最佳 N 吞吐量比率
gemma‑4‑26B‑A4B‑it DFlash MATH500 7 2.87×
gemma‑4‑26B‑A4B‑it Gemma 4 MTP MATH500 5 2.74×
gemma‑4‑26B‑A4B‑it EAGLE‑3 GSM8K 5 2.27×
gemma‑4‑31B‑it DFlash MATH500 7 2.41×
Qwen3‑8B DSpark GSM8K 7 1.63×
Qwen3‑6‑35B‑A3B DFlash MATH500 7 2.06×
Kimi‑K2.5 DFlash GSM8K 7 2.37×

各方法的趨勢

  • 序列草稿(原生 MTP、Gemma 4 MTP、EAGLE‑3) 通常在提案長度適中時(N ≈ 3–5)提升吞吐量。超過此範圍後,額外的序列草稿工作量超過節省的目標傳播次數,導致吞吐量停滯或下降。
  • 平行草稿(DFlash、DSpark) 通常在較大 N(N ≈ 7–11)時達到最高比率。後續位置的接受率下降,但平行成本仍低,因此整體吞吐量仍持續上升。
  • DSpark 加入輕量級馬可夫頭;由於額外的序列修正步驟,其吞吐量通常比純 DFlash 在相同 N 下略低數百分點。
  • 模型家族至關重要 – 原生 MTP 建立於 Qwen 3.5/3.6 模型中,表現優於相同模型大小的 DFlash;然而,對於 Gemma 4,平行方法則占優勢。
  • 工作負載依賴性 – 以程式碼為主的資料集(HumanEval、MBPP)通常偏好較短提案,因為後續 token 預測性較低;而以數學為主的資料集(GSM8K、MATH500)則能容忍較長提案。

接受行為

在所有實驗中,第一個草稿 token 的接受率超過 90%。接受率隨位置遞減;在 DFlash 且 N = 15 時,第 15 個 token 的接受率可能低於 5%。平均接受長度(MAL)與吞吐量強烈相關:MAL 越高 → 目標傳播次數越少 → 加速比越高。


實務調校建議

  1. 從小開始 – 對任何方法,從 num_speculative_tokens = 1(無額外成本)開始,並驗證正確性。
  2. 掃描 N – 在代表性工作負載上,逐步增加 N(例如 1–3–5–7–11–15),並測量吞吐量、MAL 與 AR。
  3. 觀察各位置接受率 – 若在位置 k 後接受率急劇下降,將 N 降低至 k,以避免浪費草稿工作。
  4. 記憶體預算 – 確保有足夠的 GPU 記憶體容納草稿檢查點;平行方法(DFlash/DSpark)通常需求最高。
  5. 硬體特定調校 – AMD Instinct GPU 可從 ROCm 優化注意力後端(triton_attn)獲益。平行草稿能充分利用 MI300X/MI355X 的高記憶體頻寬。
  6. 批次大小與 token 限制 – 較大的 max-num-batched-tokens 與更長的上下文視窗可改善草稿成本的分攤,特別是對平行方法而言。

如何取得草稿檢查點

發布者 支援方法 例子檢查點
Google Gemma 4 MTP google/gemma-4-26B-A4B-it-assistant, google/gemma-4-31B-it-assistant
LightSeek EAGLE‑3, EAGLE‑3.1 lightseekorg/kimi-k2.5-eagle3-mla
Red Hat AI EAGLE‑3, DFlash, DSpark RedHatAI/gemma-4-26B-A4B-it-speculator.eagle3, RedHatAI/gemma-4-31B-it-speculator.dspark
Z‑Lab DFlash z-lab/gemma-4-26B-A4B-it-DFlash, z-lab/Qwen3.8B-DFlash-b16
DeepSeek AI EAGLE‑3, DFlash, DSpark deepseek-ai/eagle3_qwen3_8b_ttt7, deepseek-ai/dflash_qwen3_8b_block7
Inferact EAGLE‑3, DSpark Inferact/MiniMax-M3-EAGLE3, Inferact/Kimi-K3-DSpark

訓練新預測器(高階工作流程)

  1. 收集代表性提示 用於預期工作負載(對話、程式碼、數學等)。
  2. 使用相同的分詞器、對話範本與採樣設定 產生目標模型的回應,與推理時一致。
  3. 選擇隱藏狀態提取模式 – 即時(即時)、離線(預先儲存)或混合(先快取第一個 epoch)。
  4. 提取所需的目標層(例如 EAGLE‑3 需早期、中期、晚期;DFlash/DSpark 需所有選定層)。
  5. 以相同詞彙、隱藏大小與 token 嵌入 訓練預測器。包含方法特定的頭(平行區塊、馬可夫修正等)。
  6. 驗證 接受率、MAL 與端到端吞吐量;若接受率低,可迭代調整提示混合或訓練超參數。
  7. 打包 檢查點,並使用 --speculative-config 與目標模型一同提供服務。

未來方向

  • 非學習型預測(例如 n-gram 後綴預測)可補足學習型預測器,適用於高度重複的程式碼編輯工作負載。
  • 更廣泛的併發研究 – 在多使用者批次負載、不同批次大小與不同採樣溫度下測量預測性解碼。
  • 預測器資料影響分析 – 系統性分析提示多樣性與領域特定訓練資料如何影響程式碼、數學、對話與多語言任務中的接受率。
  • 深度剖析 – 在 AMD GPU 上隔離草稿產生、目標驗證、KV 快取重用與圖執行開銷,以指導核心層優化。

致謝

感謝 Hongxia Yang 與 Peng Sun(AMD)以及 Pin Siang Tan、Jun Kang Chow 與 Ye Hur Cheong(Embedded LLM)提供硬體存取與合作。


免責聲明

測量在兩種 AMD Instinct 配置上執行:

  • 硬體 1:8 × MI300X GPU(gfx942),搭配雙 EPYC 9654 96 核 CPU。
  • 硬體 2:8 × MI355X GPU(gfx950),搭配雙 EPYC 9575F 64 核 CPU(用於 MiniMax‑M3‑MXFP8)。 軟體堆疊:Ubuntu 22.04.5 LTS,ROCm 7.2.53211,vLLM 0.23.1rc1.dev1120+g0f0f28b53,PyTorch 2.11.0,Transformers 5.13.1,Python 3.12.13。結果可能因不同驅動程式、核心版本或硬體配置而異。

Sources

相關

  • Dispatch
  • Dispatch
  • Dispatch
  • Dispatch
  • Dispatch