眾所周知，對於自迴歸模型而言，記憶體頻寬一直是制約推理效能的關鍵瓶頸。當模型生成文字時，它需要反覆執行前向傳播來預測每個 token，這個過程不僅需要大量的計算資源，更重要的是會頻繁訪問記憶體來獲取模型權重和中間狀態。這種記憶體密集型的特性導致了較高的延遲，影響使用者體驗。
去年，學界興起了一種稱為推測解碼（Speculative Decoding）的技術用以緩解這一問題。這種方法使用較小的“草稿”模型來預測可能的未來 token 序列，再由主模型驗證，從而實現並行化生成，提高推理效率。
圖丨自迴歸解碼和推測解碼（來源：arXiv）
今年初，提出了一種新的推測解碼方法 Recurrent Drafter （ReDrafter），對原有方法進行了改進。透過將 RNN 草稿模型與動態樹注意力機制相結合，ReDrafter 在每個生成步驟中最多可接受 3.5 個 token，超越了此前方法的水平。
最近，蘋果又宣佈與展開深度技術合作，將其創新的 ReDrafter 推測解碼技術整合到英偉達 TensorRT-LLM 推理加速框架中。這項合作將幫助開發者在英偉達 GPU 上實現更高效的 LLM 推理效能。
ReDrafter 的核心是一個基於（RNN）的草稿模型。與先前的方法（如 Medusa）類似，它使用 Transformer 最後一層的輸出作為輸入。創新之處在於，它還將歷史 token 的嵌入作為迴圈輸入提供給草稿頭。
當 LLM 生成一個 token（比如“quiet”）時，草稿模型會獲取這個 token 的嵌入向量 et 和 LLM 最後一層的輸出 h，然後透過 RNN 更新隱藏狀態：st = f(Ust-1 + W et + b)。更新後的狀態 gt = [st, h] 被用來預測下一個 token。這種遞迴設計讓模型能夠維護動態的上下文表示，大大提升了預測準確性。
圖丨草稿模型以 LLM 的最後隱藏狀態 h 作為輸入來預測接下來的幾個標記。（來源：arXiv）
另一方面，ReDrafter 使用束搜尋來探索多個可能的序列延續。為了提高效率，它引入了動態樹注意力演算法來處理候選序列中的共同字首。
來看一個具體例子：當生成多個候選序列時，比如“morning sipping coffee and watching”和“morning sipping coffee on her”，這些序列往往共享相同的字首。傳統方法會對每個序列獨立計算，而 ReDrafter 透過動態構建字首樹來合併重複計算。系統使用張量運算實現了高效的字首檢測和合並處理，在實踐中可以減少 30%-60% 的計算量。
此外，為了提高草稿模型的預測質量，ReDrafter 採用了一種新的知識蒸餾訓練方法。系統會直接學習 LLM 的 token 生成分佈，而不是簡單地擬合真實序列。訓練目標被定義為最小化 LLM 和草稿模型輸出分佈之間的 KL 散度：min(pdraft) KL(pllm(y1:T)|pdraft(y1:T))。
在每個訓練位置 t，系統從 LLM 取樣未來 T 個 token，並最佳化經驗損失：min(pdraft) Σt −log pdraft(ŷt+1:t+T|y1:t)。這種方法確保了草稿模型能夠更好地模仿 LLM 的行為。
研究結果顯示，在 MT-Bench 基準測試中，ReDrafter 在 Vicuna 7B 和 13B 模型上實現了 2.8 倍的加速。且 ReDrafter 在每個生成步驟平均可接受 4.2 個 token，這意味著相比傳統的每步生成一個 token 的方式，ReDrafter 顯著提升了並行度。在大規模批處理場景下，系統的最高吞吐量可達到每秒 1636 個 token。
並且，在不同的硬體平臺上，ReDrafter 都展現出良好的適應性。在搭載 M2 Ultra 的系統上，13B 引數量的模型達到了 1.94 倍的加速比。即使在計算資源相對有限的 M1 Max 上，系統仍然實現了 1.37 倍的效能提升。
為了讓 ReDrafter 的效能優勢更好地在 TensorRT-LLM 上發揮，英偉達的工程團隊進行了一系列創新。其最重要的改進是將驗證邏輯直接整合到引擎內部，而不是依賴執行時處理。這種設計不僅降低了系統開銷，還為核心選擇和排程最佳化提供了更大的自由度。
在批處理方面，研究團隊採用了一種創新的 inflight-batching（IFB）策略。當處理混合的請求流時，系統會智慧地將上下文階段和生成階段的請求分開處理，每類請求使用專門最佳化的計算流水線。這種設計在保持低延遲的同時，顯著提升了 GPU 利用率。
圖丨 ReDrafter 的 TensorRT-LLM 引擎的線上批處理相容計算工作流程（來源：英偉達）
尤其值得一提的是系統對字首樹計算的最佳化。英偉達開發了專門的運算元來處理動態樹結構，這些運算元經過深度最佳化，能充分利用 Tensor Core 等硬體特性。同時，透過精心設計的記憶體管理策略，系統能夠高效處理不同長度的序列，減少記憶體碎片。
在實際部署測試中，研究團隊使用一個擁有數百億引數的生成模型進行了全面評估。在貪婪解碼模式下，系統實現了 2.7 倍的生成速度提升。這種效能改進不僅降低了使用者感知延遲，還減少了所需的 GPU 數量和能源消耗。
圖丨使用英偉達 TensorRT-LLM 和 ReDrafter 與自迴歸相比，每秒令牌速度加快。（來源：蘋果）
這種效能改進將直接轉化為實際價值：它能最大限度地減少使用者延遲並減少所需的硬體數量。英偉達在其部落格中強調，這次合作顯著增強了 TensorRT-LLM 框架的能力。透過提供標準化的介面和最佳化的實現，使得開發者能夠更容易地部署複雜的語言模型。
ReDrafter 此前已在 GitHub 開源，專案地址：https://github.com/apple/ml-recurrent-drafter。
參考資料：
1.https://machinelearning.apple.com/research/redrafter-nvidia-tensorrt-llm
2.https://arxiv.org//2403.09919v5
3.https://developer.nvidia.com/blog/nvidia-tensorrt-llm-now-supports-recurrent-drafting-for-optimizing-llm-inference/
排版：劉雅坤
