透過推測加速 Sonar

推測式解碼（Speculative decoding）透過使用快速且小型的高效率草稿模型來產生補全候選項目，並由較大的目標模型進行驗證，藉此加快大型語言模型（LLM）的產生速度。在此方案下，昂貴的目標模型不會只執行一次並產生單一 Token，而是在單一步驟中發出多個 Token。在此我們將介紹各種推測式解碼的實作細節，這些細節已應用於 Perplexity，以縮短 Sonar 模型的 Token 間延遲。

推測式解碼

推測式解碼利用自然語言的結構以及 Transformers 的自迴歸（auto-regressive）特性來加快 Token 產生速度。儘管較大的模型（例如 Llama-70B）比較小的模型（例如 Llama-1B）具備更豐富的知識，但在某些較簡單的任務中，它們的表現不相上下。這種重疊性確實顯示某些序列由成本較低的模型產生會更適合，而將複雜的問題留給較大的模型。挑戰在於如何判斷哪些補全項目較好，以及較小模型的產生品質是否與較大模型相當。

幸運的是，LLM 是自迴歸 Transformers：當給定一組 Token 序列時，它們會輸出下一個 Token 的機率分佈。此外，從與輸入序列中的 Token 相關聯的中介特徵所衍生出的 logits，也指出模型發出那些精確 Token 的可能性。此特性啟用了推測：如果 Token 序列是由較小的模型從輸入前綴開始產生，則可以透過較大的模型執行該序列，以判斷它與目標模型的吻合程度。候選項目的每個前綴都會被賦予一個機率評分，並挑選出超過接受閾值的最長前綴。作為附加好處，目標模型還能免費提供後續的 Token：如果草稿模型產生了 n 個 Token，則在單一步驟中最多可以發出 n + 1 個 Token。

放大

在推論時，推測式取樣過程大致可分為 4 個階段：

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

草稿產生（Draft generation）：草稿模型會進行迭代以產生固定數量的 Token。草稿序列可以是線性的，或者模型可以探索高達給定深度的樹狀結構（EAGLE、Medusa）。在這裡，我們專注於線性序列。

接受（Acceptance）：目標模型在草稿序列上執行，建構對應至每個草稿 Token 的 logits。接著會決定最長可接受序列的長度。

目標產生（Target generation）：由於目標模型產生了 logits，在不匹配的位置或序列尾端，logits 會對應至後續的 Token。可以對這些 logits 進行取樣，以從目標模型提供強健的 Token，從而為序列收尾。

實作推測式解碼的方法有很多種。在本文中，我們將專注於我們用來加速 Sonar 模型的方案（使用內部 1B 模型），以及我們正在建構以加速 DeepSeek 規模模型的預測機制。

目標－草稿（Target-Draft）

推測式解碼可以透過將現有的小型 LLM 作為草稿模型與目標模型結合來產生候選序列，從而實現。在生產環境中，我們使用在與目標相同資料集上進行微調的 Llama-1B 模型來加速 Sonar。雖然此方法不需要從頭開始訓練草稿，但小型模型仍會消耗大量的 KV 快取容量，並引入微小的預填負擔，從而增加 TTFT。

在此方案下，解碼器僅對純解碼批次進行推測，並在預填期間或混合預填－解碼批次上透過標準取樣來產生 Token。在預填階段，會立即對目標 logits 進行取樣，以同時預填草稿的 KV 快取中新產生的 Token。此時尚未對草稿進行取樣，但它所產生的 logits 會被帶入解碼階段。

放大

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

EAGLE

EAGLE 是一種推測式解碼方案，它會探索多個草稿序列，這些序列是透過對可能的草稿 Token 進行樹狀走訪所產生的。透過連續執行草稿 Token 來探索固定形狀（EAGLE）或動態形狀（EAGLE-2）的樹，並在每個節點考慮 Top-K 候選項目，而不是遵循線性序列中得分最高的 Token。然後對序列進行評分，並選出最長且合適的序列來繼續，同時附加來自目標模型的額外 Token。

放大

為了實現更準確的預測，EAGLE 草稿模型的預測不僅基於 Token，還使用了目標模型的目標特徵（最後一層隱藏狀態）。EAGLE 的缺點是需要訓練客製化的小型草稿模型，這些模型必須足夠準確，才能在低延遲預算內產生適當的候選項目。通常，草稿模型是與原始模型的解碼器層相同的單一 Transformer 層，它透過綁定其嵌入和 lm_head 投影與目標緊密耦合。由於這需要較少的 KV 快取容量，因此 EAGLE 具有較低的記憶體佔用。

為了驗證目標模型中的樹狀序列，必須使用自訂注意力遮罩。不幸的是，對於實際的輸入長度而言，為整個序列使用自訂注意力遮罩會顯著降低注意力的速度（高達 50%），從而抵消了透過推測所能實現的部分加速。基於此原因，我們尚未在生產環境中部署完整的樹狀探索，而是專注於 DeepSeek-V3 技術報告中所提出的類似 MTP 方案的單一 Token 預測的特殊情況。

MTP

此方案與 draft-target decoding 類似，但不同之處在於會將隱藏狀態與 token 一併用於預測。與一般的 draft-target speculation 相比，在 prefill 與 decode 階段都需要多做一些工作。草稿模型同時使用 token 和隱藏狀態：token t_{i+1} 是從與 token t_i 對應的 logits L_i 中取樣而得，而這些 logits 又是由隱藏狀態 H_i 推導出來。因此，輸入 token 緩衝區必須相對於 target 輸出的隱藏狀態向左移一個步長。下圖標示了訓練時使用的對應關係，以及推理期間的位移。

放大

解碼流程與草稿－目標解碼非常相似，不同之處在於隱藏狀態和 logits 都會被保留並帶入下一步。我們的實作共用了所有相關的取樣和 logits 處理，僅針對模型前向呼叫進行專門化。當預測多個 Token 時，草稿模型會使用草稿隱藏狀態進行預測，並根據其自身的特徵來填入 KV 快取項目。長遠來看，這可能會降低準確度。隨後，當執行草稿模型以填入目標預測的 KV 快取項目時，我們會在整個序列上執行它，並將更準確的目標隱藏狀態作為輸入。由於這些草稿模型很小，處理額外 Token 所增加的成本微不足道。

訓練 MTP 頭（MTP Heads）

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

MTP 訓練的目標是使從目標隱藏狀態推斷出的草稿隱藏狀態和 logits，與目標的下一個 Token logits 和隱藏狀態相匹配。由於隱藏狀態的推論成本高昂，我們使用我們經過推論優化後的目標模型實作來預先計算它們，以便在訓練期間使用。然而，為了驗證推論 MTP 實作並確保量化或優化所產生的數值差異不會妨礙結果，對於驗證損失和準確度評估，我們完全重複使用了目標模型和草稿模型的推論實作。

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

放大

雖然我們針對沒有推測的推論完全並行化了 CPU 端和 GPU 端的 工作，但我們發現推測式解碼的 CPU-GPU 平衡更加複雜。主要的挑戰源自於一個事實：被接受 Token 的數量決定了後續運行的序列序列長度，這引入了一個難以避免的 GPU-to-CPU 同步點。我們試驗了不同的排程方案，以便最好地隱藏 CPU 工作的延遲。

使用推測式解碼進行推論

在推論引擎中，為了替輸入序列產生 Token，需要先將它們組合成大小合理的批次，然後必須在 KV快取中為接下來的 Token 配置頁面。接著，輸入 Token 和 KV 頁面資訊會被打包成一個廣播緩衝區，發送給執行該模型的所有並行行程（parallel ranks）。最後，詮釋資料會被複製到 GPU 記憶體中，並執行模型以產生 logits，從中對下一個 Token 進行取樣。

不同於某些實作透過包裝器（wrapper）鬆散地耦合草稿和目標推論伺服器，並在它們之間協調請求，我們的草稿－目標配對是緊密耦合的，並且步調一致地進行產生。所有形式的推測式解碼都在模型之間共用批次排程和 KV 頁面配置：這統一了將模型連接到整體推論伺服器的邏輯，因為它們都公開了相同的介面。

Perplexity 的推論執行階段（runtime）圍繞著 FlashInfer 建構，該工具決定了為了設定和排程注意力核心（attention kernel）而需要建構的詮釋資料。給定形成批次的某些輸入序列（用於預填、解碼或驗證），必須在 CPU 端進行工作以配置中介緩衝區，並填入注意力中使用的某些常數緩衝區。這項工作除了批次排程和 KV 頁面配置的成本之外，還會產生必須隱藏的延遲，以便最大化 GPU 的利用率。

雖然我們在沒有推測的情況下，完全平行化了推論的 CPU 端與 GPU 端工作，但我們發現推測式解碼的 CPU-GPU 平衡更加複雜。主要的挑戰在於，接受的 Token 數量決定了後續執行的序列長度，這引進了一個難以避免的 GPU-to-CPU 同步點。我們實驗了不同的排程方案，以期最大程度地隱藏 CPU 工作的延遲。

目標－草稿排程（Draft-Target Schedule）

儘管整個 LLM 作為草稿時比目標模型小，但它仍然會在 GPU 上引入相當大的延遲，從而提供了一些餘裕來隱藏昂貴的 CPU 操作。由於小型模型無法從張量並行（tensor parallelism）中受益，目標模型和草稿模型所分片（sharded）跨越的行程（ranks）數量之間存在不匹配。在我們的實作中，草稿模型僅在 TP 群組的領導行程（leader rank）上執行。

放大

如前所述，解碼步驟會將 logits 帶入下一次運行。這允許我們將草稿模型的一次執行與 CPU 端的批次排程工作重疊。在組合好批次後，對取樣器和草稿的重複呼叫會產生草稿 Token。同時，為目標模型組合驗證批次，並與並行工作程序（workers）進行同步。對目標 logits 進行驗證和取樣，以決定被接受的序列長度。此時，為了決定後續的序列長度，GPU-to-CPU 同步是必要的。由於草稿模型僅在領導節點上執行，因此其批次是以循序方式設定的，並啟動其執行，以便用目標產生的額外 Token 來填入其 KV 快取項目。在此次運行中，由該草稿運行所產生的 logits 將用於對後續運行的第一個草稿 Token 進行取樣。最重要的是，在草稿運行時，可以排程下一批次。

單一 Token 的 MTP 排程（MTP Schedule for a Single Token）

雖然執行階段尚未提供 Eagle 風格的草稿樹狀探索，但我們實作了此方案的特殊情況，即考慮由單一 Transformer 解碼器層大小的模型所產生的草稿 Token 線性序列。此方案可用於使用 DeepSeek R1 開源權重的草稿預測。預測單一 Token 的子情況很有趣，因為大型 MTP 層達到了足夠高的接受率，足以證明其額外負擔是合理的。

MTP 排程稍微複雜一些，因為草稿模型快得多，隱藏的 CPU 端延遲較少。此外，草稿與目標模型一起進行分片，需要透過共用記憶體傳輸來取得批次資訊。運行開始時，會傳輸批次資訊並從繼承的 logits 中取樣第一個 Token，這與先前的方案類似。接著，執行目標模型來驗證 Token，處理 2 * D 個 Token（其中 D 為解碼批次大小）。這非常適合在透過 InfiniBand 等較慢互連的混合專家（MoE）模型中進行微批次處理（micro-batching），因為批次可以均勻地分為兩半。目標模型的隱藏狀態會帶入下一個草稿運行，而 logits 則傳遞給取樣器進行驗證。

放大

透過在 GPU 上執行少量的額外工作，我們避免了在草稿序列接受後的 CPU-to-GPU 同步。在目標的輸入 Token 位移後，核心（kernel）會將下一個目標 Token 插入其對應的位置。然後，草稿模型會使用與目標相同的批次資訊重新運行，填入 KV 快取項目並建構下一個運行的 logits 和隱藏狀態，同時對未被接受的 Token 執行一些冗餘工作。在這些情況下，由於草稿模型的規模很小，未使用工作的延遲幾乎是微不足道的。與草稿運行並行，會在 CPU 上決定序列長度，並啟動下一批次的排程，而不必等待 GPU 工作終止。

草稿層中額外工作的負擔在注意力中並不明顯，然而 MLP 層則較具問題。由於矩陣乘法指令會沿著 Token 數量的維度對齊到 64 的邊界，如果倍增不需要顯著更多的區塊，則負擔會被隱藏起來。對於較長的草稿序列，負擔更為昂貴，因此正規草稿－目標模型所使用的方案效果更好。

參考文獻

透過推測式解碼實現 Transformers 的快速推論

EAGLE：推測式取樣需要重新思考特徵不確定性

EAGLE-2：透過動態草稿樹實現更快速的語言模型推論

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

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

層歸一化（Layer Normalization）

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}
}