透過投機加速 Sonar

投機性解碼（Speculative decoding）利用一個快速且輕量化的草稿模型來產生候選完成內容，並由較大的目標模型進行驗證，藉此加速大型語言模型（LLMs）的生成速度。在此機制下，昂貴的目標模型不會在單一步驟中僅產生單一 Token，而是能一次輸出多個 Token。本文介紹了應用於 Perplexity 以降低 Sonar 模型 Token 間延遲的各種投機性解碼實作細節。

投機性解碼

投機性解碼利用自然語言的結構與 Transformer 的自動迴歸特性來加速 Token 生成。儘管像 Llama-70B 這樣較大的模型擁有的知識量遠超如 Llama-1B 等較小的模型，但在某些較簡單的任務上，它們的執行效果相當。這種重疊現象顯示，某些序列由成本較低的模型生成效果更好，從而將複雜問題留給較大的模型。其挑戰在於判定哪些完成內容較佳，以及較小模型的生成品質是否與較大模型相當。

幸運的是，LLMs 是自動迴歸 Transformer 模型：當給定一個 Token 序列時，它們會輸出下一個 Token 的機率分佈。此外，由輸入序列中與 Token 相關聯的中間特徵所導出的 Logit，也顯示了模型發出這些確切 Token 的可能性。此特性促成了投機：如果一個 Token 序列是由較小模型從輸入前綴開始生成的，則可將其輸入較大模型中，以判定其與目標模型的契合程度。每個候選前綴都會根據機率進行評分，並選取高於接受閾值的最長前綴。作為額外紅利，目標模型還會免費提供一個後續 Token：如果草稿模型生成 n 個 Token，則在一個步驟中最多可輸出 n + 1 個 Token。

放大

在推論階段，投機性採樣過程可粗略分為 4 個階段：

預填充（Prefill）：目標模型與草稿模型都必須在輸入序列上執行，以填充 KV 快取項目。雖然某些機制（例如 Medusa）使用較簡單的密集層進行預測，但本文專注於需要自身 KV 快取的基於 Transformer 的草稿模型。

草稿生成：草稿模型反覆運算以產生固定數量的 Token。草稿序列可以是線性的，或者模型也可以探索達到給定深度的樹狀結構（如 EAGLE、Medusa）。在此，我們專注於線性序列。

接受：目標模型在草稿序列上執行，構建對應於每個草稿 Token 的 Logit。藉此確定可接受序列的最長長度。

目標生成：由於目標模型產生了 Logit，在不匹配位置或序列末尾，Logit 對應於一個後續 Token。這些 Logit 可經由採樣以提供來自目標模型的穩健 Token，從而結束序列。

存在多種實作投機性解碼的方法。在本文中，我們將專注於我們用來加速 Sonar 模型的機制（使用內部 1B 模型），以及我們目前正在建構、旨在加速 DeepSeek 規模模型效能的預測機制。

目標-草稿

投機性解碼可透過將現有的小型 LLM 作為草稿模型與目標模型耦合，進而產生候選序列來達成。在生產環境中，我們使用在與目標相同數據集上進行微調的 Llama-1B 模型來加速 Sonar。雖然此方法不需要從頭開始訓練草稿模型，但小型模型仍會使用顯著的 KV 快取容量，並引入輕微的預填充開銷，進而增加 TTFT（首個 Token 時間）。

在此機制下，解碼器僅對僅解碼（decode-only）批次進行投機，並在預填充階段或混合預填充-解碼批次期間透過標準採樣產生 Token。在預填充階段，目標 Logit 會立即被採樣，以同時預填充草稿 KV 快取中新生成的 Token。此時草稿尚未進行採樣，但其產生的 Logit 會被攜帶至解碼階段。

放大

在解碼階段，草稿模型會推進，並在每個階段採樣頂部 Token。在達到所需的草稿長度後，這些 Token 會通過目標模型以產生 Logit，採樣器據此識別出被接受的序列長度。接受與否取決於比較草稿與目標的完整機率分佈。由於目標模型總是會在被接受的草稿序列後輸出單組 Logit，該 Logit 會被採樣以產生額外輸出。由於草稿模型尚未看過該被接受的 Token，因此會重新執行以填充對應的 KV 快取項目，為下一個解碼步驟做準備，並再次攜帶 Logit。

EAGLE

EAGLE 是一種投機性解碼機制，它透過對可能草稿 Token 進行樹狀遍歷來探索多個草稿序列。使用固定（EAGLE）或動態形狀（EAGLE-2）的樹，透過連續執行草稿 Token 來進行探索，並在每個節點考慮 Top-K 候選，而非在線性序列中僅遵循得分最高的 Token。接著對這些序列進行評分，選出最長且合適的序列以繼續，並附加一個來自目標模型的額外 Token。

放大

為了達到更準確的預測，EAGLE 草稿模型不僅基於 Token 進行預測，還使用了目標模型的目標特徵（最後一層隱藏狀態）。EAGLE 的缺點在於需要訓練客製化的小型草稿模型，以在低延遲預算內準確地生成合適的候選內容。通常，草稿模型是與原始模型解碼器層相同的單一 Transformer 層，透過連結其嵌入與 lm_head 投影與目標模型緊密耦合。由於這需要的 KV 快取容量較少，EAGLE 的記憶體佔用空間較低。

為了在目標模型中驗證樹狀序列，必須使用客製化的注意力遮罩（Attention Masks）。遺憾的是，針對整個序列使用客製化注意力遮罩會顯著減慢實際輸入長度的注意力計算速度（最高達 50%），抵銷了透過投機所能獲得的部分加速。基於此原因，我們尚未在生產環境中部署完整的樹狀探索，而是專注於 DeepSeek-V3 技術報告中提出的透過 MTP 類機制進行單 Token 預測的特殊情況。

MTP

此機制與草稿-目標解碼類似，差異在於使用隱藏狀態與 Token 一起進行預測。與常規草稿-目標投機相比，在預填充與解碼階段都必須執行稍多的工作。草稿模型同時使用 Token 與隱藏狀態：Token t_{i+1} 是從對應於 Token t_i 的 Logit L_i 採樣而來，而這些 Logit 又是從隱藏狀態 H_i 導出。因此，輸入 Token 緩衝區必須相對於目標輸出的隱藏狀態向量向左移動一步。下圖標記了訓練所使用的對應關係，以及推論期間的位移。

放大

解碼流程與草稿-目標解碼相當類似，差異在於隱藏狀態與 Logit 皆會被攜帶過來。我們的實作共享所有相關的採樣與 Logit 處理邏輯，僅針對模型前向呼叫進行了專門化處理。當預測多個 Token 時，草稿模型使用草稿隱藏狀態進行預測，同時根據其自身特徵填充 KV 快取項目。長遠來看，這可能會降低準確度。隨後，當執行草稿模型以填充目標預測的 KV 快取項目時，我們會在整個序列上執行它，並將更準確的目標隱藏狀態作為輸入。由於這些草稿模型很小，處理額外 Token 的成本可以忽略不計。

訓練 MTP 頭（MTP Heads）

為了從 MTP 獲益，我們建構了訓練 MTP 頭所需的基礎設施，這些頭附著在我們於 Perplexity 數據集上微調的模型上，並在一個具有 8xH100 裝置的節點上執行。我們可以在約一天內為從 Llama-1B 到 Llama-70B 以及 DeepSeek V2-Lite 等模型建構頭。對於較大的模型，我們則依賴於在微調過程中建構的 MTP 頭。

MTP 訓練的目標是將從目標隱藏狀態推導出的草稿隱藏狀態與 Logit，與目標的下一個 Token Logit 與隱藏狀態進行匹配。由於隱藏狀態的推論代價昂貴，我們使用我們針對推論最佳化的目標模型實作預先計算這些狀態，以便在訓練期間使用。然而，為了驗證推論 MTP 實作並確保量化或最佳化導致的數值差異不會阻礙結果，在驗證損失與準確度估計方面，我們完全重複使用目標與草稿模型的推論實作。

當從原始論文中使用的 ShareGPT 數據集擴展到更大的樣本時，我們注意到 EAGLE 論文中概述並實作的 MTP 頭架構無法為 70B 大小的模型進行訓練。與包含大量較短序列的 ShareGPT 不同，我們是在數量較少但提示詞長度明顯更長的數據上進行訓練。由於原始 EAGLE 頭的結構與典型 Transformer 略有不同，我們重新引入了一些被移除的 RMS 正則化層。我們發現這不僅使訓練得以收斂，還將頭的準確度提升了幾個百分點。

放大

層正則化不僅有助於訓練，重新引入這些正則化在數學上也符合直覺。MTP 頭重複使用目標模型的嵌入與 Logit 投影，因為它們的體積可能相當龐大（對於 Llama 70B 約為 2 GB）。在訓練期間，這些內容會被凍結，預期 MTP 層能學習將預測嵌入到與原始模型投影層在訓練期間所學習到的相同的向量空間中。透過移除正則化，單一 MLP 預期能學習到與「MLP 後跟隨正則化」相同的函數，這會阻礙草稿與目標模型隱藏狀態之間的匹配。

使用投機性解碼進行推論

在推論引擎中，為了產生輸入序列的 Token，需要先將其分組為大小合理的批次，然後在 KV 快取中為下一個 Token 分配頁面。接著，輸入 Token 與 KV 頁面資訊會被封裝到廣播至所有執行模型之平行處理單元（ranks）的緩衝區中。最後，元數據會被複製到 GPU 記憶體中，並執行模型以產生 Logit，從而採樣出下一個 Token。

與某些透過封裝器在請求之間進行協調以鬆散耦合草稿與目標推論伺服器的實作不同，我們的草稿-目標對是緊密耦合的，並步調一致地進行生成。對於所有形式的投機性解碼，批次排程與 KV 頁面分配都在模型間共享：這統一了連接模型與總體推論伺服器的邏輯，因為它們都暴露了相同的介面。

Perplexity 的推論執行階段是圍繞 FlashInfer 建構的，該引擎確定了配置與排程注意力核心所需的元數據。給定組成批次的某些輸入序列，對於預填充、解碼或驗證，必須在 CPU 端執行工作以分配中間緩衝區並填充注意力中使用的某些常數緩衝區。這項工作是批次排程與 KV 頁面分配成本之外的額外開銷，這些開銷同樣會產生延遲，必須隱藏起來以最大化 GPU 利用率。

雖然我們完全平行化了無投機推論的 CPU 端與 GPU 端工作，但我們發現投機性解碼的 CPU-GPU 平衡更為複雜。主要挑戰源於「被接受的 Token 數量決定了後續執行的序列長度」，這引入了一個難以避免的 GPU 到 CPU 同步點。我們嘗試了不同的排程機制，以求最佳地隱藏 CPU 工作的延遲。

草稿-目標排程

儘管整個 LLM 用作草稿時比目標模型小，但它仍在 GPU 上引入了相當大的延遲，為隱藏昂貴的 CPU 操作提供了一些空間。由於較小的模型無法從張量平行化（tensor parallelism）中獲益，目標與草稿分片所在的 rank 數量之間存在不匹配。在我們的實作中，草稿模型僅在 TP 組的領導 rank 上執行。

放大

如前所述，解碼步驟會將 Logit 攜帶到下一次執行中。這使我們能夠將草稿模型的一次執行與 CPU 端的批次排程工作重疊。在批次組合完成後，重複呼叫採樣器與草稿模型會產生草稿 Token。同時，目標模型的驗證批次被組合並與平行工作單元同步。目標 Logit 被驗證並採樣以確定被接受的序列長度。此時，為了確定後續序列長度，GPU 到 CPU 的同步是必要的。由於草稿模型僅在領導節點上執行，其批次是依序設定的，並啟動執行以使用目標模型產生的額外 Token 填充其 KV 快取項目。此草稿執行產生的 Logit 將用於採樣後續執行中的第一個草稿 Token。最重要的是，在執行草稿的同時，下一個批次即可進行排程。

單 Token 的 MTP 排程

雖然目前的執行階段尚未提供 Eagle 風格的草稿樹探索，但我們實作了該機制的特殊情況，考慮由單一 Transformer 解碼器層大小的模型所產生的草稿 Token 線性序列。此機制可用於使用 DeepSeek R1 開源權重進行草稿預測。預測單個 Token 的子情況非常有趣，因為大型 MTP 層達到了足夠高的接受率，足以抵銷其開銷。

MTP 排程稍微複雜一些，因為草稿模型快得多，隱藏的 CPU 端延遲較少。此外，草稿與目標模型一起進行分片，需要進行批次資訊的共用記憶體傳輸。執行開始時，類似於先前的機制，傳輸批次資訊並從攜帶的 Logit 中採樣第一個 Token。接著，目標模型執行驗證 Token，處理 2 * D 個 Token，其中 D 是解碼批次大小。對於在諸如 InfiniBand 等較慢互連上的專家混合（MoE）模型中的微批次處理而言，這非常理想，因為批次可均勻拆分為兩半。目標的隱藏狀態會攜帶到下一次草稿執行中，而 Logit 則被傳遞至採樣器進行驗證。

放大

透過在 GPU 上執行有限的額外工作，我們避免了在草稿序列接受後進行 CPU 到 GPU 的同步。在目標的輸入 Token 位移後，核心會將下一個目標 Token 插入其對應位置。接著，草稿模型會使用與目標模型相同的批次資訊重新執行，填充 KV 快取項目並建構下一輪所需的 Logit 與隱藏狀態，對未被接受的 Token 做了一些冗餘工作。在這種情況下，由於草稿模型規模較小，未使用工作的延遲幾乎無法測量。在草稿執行的同時，CPU 上會確定序列長度並啟動下一個批次的排程，無需等待 GPU 工作結束。

草稿層中額外工作的開銷在注意力計算中並不顯著，但 MLP 層較為棘手。由於矩陣乘法指令會沿著 Token 數量維度填充至 64 的邊界，如果加倍不需要顯著更多的區塊，開銷則會被隱藏。對於較長的草稿序列，開銷更昂貴，常規草稿-目標模型使用的機制效果更好。

參考資料

透過投機性解碼實現 Transformer 的快速推論

EAGLE：投機性採樣需要重新思考特徵不確定性

EAGLE-2：利用動態草稿樹加速語言模型推論

EAGLE-3：透過訓練時測試擴展大型語言模型的推論加速

Medusa：具備多重解碼頭的簡單 LLM 推論加速框架

層正則化

FlashInfer：用於 LLM 推論服務的高效且可自訂的注意力引擎

DeepSeek-V3 技術報告

引用

@online{perplexity2025AcceleratingSonarThroughSpeculation,
  author = {{Perplexity} and Nandor Licker and Erin McKnight},
  title = {{Accelerating Sonar through speculation}},
  date = {2025-06-10},
  url = {https://perplexity.coolkeypoint.com/hub/blog/accelerating-sonar-through-speculation}
}