隨著 ChatGPT 的出現，大規模語言模型（LLM）研究及其應用得到學術界和工業界的廣泛關注。一方面，開源的 LLM 模型不斷涌現，比如 OPT、BLOOM、LLaMA 等，這些預訓練模型的推出極大地促進了 LLM 的相關研究，使得 LLM 可以被應用于解決愈發復雜的實際問題。利用這些開源模型，快速構建一套基于 LLM 的應用服務已經變得愈發容易，但 LLM 面臨著高昂的計算和存儲需求，其成本也令人望而卻步。

另一方面，以羊駝家族（如 Alpaca、Vicuna、Guanaco）為代表的，經過微調或蒸餾的小型化 LLM 也成為了當下的研究焦點之一，在多項測評中都展現出了優異的表現；此外，以 Quantization、LoRA、Offloading 為代表的多項系統優化技術使得以更低的資源需求部署這些 LLM 成為可能。但天下沒有免費的午餐，有關證據表明 [1]，這些小型化的 LLM 以及面向低資源場景的系統優化技術往往都會帶來模型質量的下降，影響最終應用的效果。

因此，如何在保證模型輸出質量的前提下，讓 LLM 推理變得高效和廉價，已經成為了 MLSys 領域非常重要的研究問題。近日，來自卡耐基梅隆大學（CMU）的 Catalyst Group 團隊發布了一款「投機式推理」引擎 SpecInfer，可以借助輕量化的小模型來幫助大模型，在完全不影響生成內容準確度的情況下，實現兩到三倍的推理加速。

• 論文鏈接：https://arxiv.org/abs/2305.09781
• 項目地址：https://github.com/flexflow/FlexFlow/tree/inference

論文作者之一、CMU 助理教授 Zhihao Jia 表示：「生成式大規模語言模型不僅推理效率低下而且部署成本很高；它們小型化的版本具有速度和價格上的優勢，但是也會影響生成內容的質量；而 SpecInfer 可以實現這兩方面的雙贏?！?/span>

同樣來自 CMU Catalyst Group 的助理教授 Tianqi Chen 也表示：「SpecInfer 可以適用于云上的 LLM 部署等場景，讓 LLM 推理更加可擴展。」

研究現狀

目前 LLM 推理主要依賴于自回歸式（auto-regressive）的解碼（decoding）方式，每步解碼只能夠產生一個輸出 token，并且需要將歷史輸出內容拼接后重新作為 LLM 的輸入，才能進行下一步的解碼?？紤]到這種數據依賴，現有 LLM 推理系統如 FasterTransformer 會采用一種增量式解碼（incremental decoding）技術，將已經解碼的 token 對應的 key/value 進行緩存，避免重新計算。但是，這類系統仍然面臨兩個關鍵的缺陷：1）由于逐 token 計算的解碼范式，算子并行度有限，GPU 硬件資源難以被充分利用；2）當序列過長時，KV-cache 空間消耗過大，有限的 GPU 顯存無法承載。因此，當面對超大規模的 LLM 推理時（如 GPT-4 32K tokens），現有系統往往面臨資源利用低效，推理延遲過高的問題。

Incremental Decoding 示意圖

為了解決上述問題，研究者提出了一種「投機式」推理引擎 SpecInfer，其核心思想是通過計算代價遠低于 LLM 的 “小模型” SSM（Small Speculative Model）替代 LLM 進行投機式地推理（Speculative Inference），每次會試探性地推理多步，將多個 SSM 的推理結果匯聚成一個 Speculated Token Tree，交由 LLM 進行驗證，通過高效的樹形解碼算子實現并行化推理，驗證通過的路徑將會作為模型的推理結果序列，進行輸出。

總體上來說，SpecInfer 利用了 SSM 的內在知識幫助 LLM 以更低廉的計算成本完成了主要的推理過程，而 LLM 則在一定程度上破除了逐 token 解碼的計算依賴，通過并行計算確保最終輸出的結果完全符合原始的推理語義。

SpecInfer 工作流程

系統設計

SpecInfer 系統架構

可學習推測器（Learning-based Speculator）

Speculator 的主要作用是利用 SSM 快速產生對 LLM 未來輸出的推測結果，SSM 可以是（微調后）小版本的 LLM（如 LLaMA 7B），也可以是量化或蒸餾的小規模 LLM，還可以是可供檢索的知識庫（如參考文本）亦或是用戶的自定義函數。總之，SSM 的輸出結果越接近 LLM，驗證時才會更容易通過，整體的推理效率才會更高。

為此，SpecInfer 引入集成學習的思想，將多個 SSM 的結果融合，提高輸出的差異化程度。為了盡可能提高匹配率，Speculator 提出了 Collective Boost-Tuning 方法，即在一個公開的通用數據集（如 OpenWebText）上，從一個較弱的 SSM 開始進行微調，將匹配程度較低的序列不斷從數據中過濾，交由新的 SSM 來學習，持續多次，提高整體的推測質量；此外，Speculator 還引入了一個可學習的調度器（scheduler）來決定選用哪些 SSM 以獲得更長的匹配序列長度。

Token 樹驗證器（Token Tree Verifier）

SSM 的推理速度優勢是 SpecInfer 能夠加速推理的前提，但另一個不可或缺的因素就是 LLM 對并行化推理的支持。在 SpecInfer 中，LLM 并不直接作為推理引擎產生輸出 token，但是它需要對 Speculator 中 SSM 產生的 token 進行驗證，確保輸出內容符合 LLM 的推理語義。

在 SpecInfer 中，SSM 產生的輸出序列會被組織成 token tree 的樹形結構，避免冗余的存儲開銷。為了能夠在 token tree 上進行并行化的驗證，SpecInfer 提出了一種樹形注意力（Tree Attention）計算方法，通過構造的 mask 矩陣和基于深度優先的 KV-cache 更新機制，Verifier 可以在不增加額外存儲的同時，盡可能并行化樹中每一條路徑的解碼過程。相比于樸素的逐序列或逐 Token 的解碼方式，樹形解碼可以同時在內存開銷和計算效率上達到最優。

Tree-based Decoding 示意圖

大規模 LLM 和小規模 SSM 協同工作

Speculative Inference 執行 Timeline 對比

大規模的 LLM 在參數量上通?？梢赃_到小規模 SSM 的幾十倍甚至上百倍，而 SSM 相比于 LLM，在推理速度上，基于通常的系統實現，也有數倍到數十倍的性能優勢，SpecInfer 結合了 SSM 極低的推理延遲以及 LLM 的并行驗證能力，大幅降低了較為耗時的 LLM 推理次數，最終可以在保證推理結果質量的情況下顯著提升模型推理速度。

系統實現

SpecInfer 基于 FlexFlow 系統實現，支持用戶自定義模型結構，導入模型參數，兼容主流深度學習框架的 operator 或 layer 抽象，現已支持常規的 GPT、LLaMA 等多種主流基礎模型。值得注意的是，FlexFlow 是一款面向分布式場景的深度學習系統，由來自 CMU、Stanford、MIT、NVIDIA 等機構的研究人員共同維護，是機器學習系統領域最早提出 “自動并行” 的工作之一 (MLSys’19, ICML’18) [2,3]，也是最早將計算圖優化以及自動并行優化集成進行聯合優化的工作 (Unity, OSDI’22) [4]。

借助于 FlexFlow 的自動并行能力，SpecInfer 可以自動完成大規模 LLM 的最優分布式部署。與此同時，SpecInfer 還可以支持 Offloading 操作，以較低的成本擴展模型的規模。SpecInfer 通過獨特的「投機式推理」機制，可以大幅降低 LLM 所需的推理步數，從而減小分布式場景的網絡通信開銷，緩解 Offloading 場景下的 PCIe 傳輸帶寬瓶頸。

實驗結果

端到端推理延遲

端到端實驗：使用 LLaMA-7B 作為 LLM，LLaMA-160M 作為 SSM，在五個對話數據集上進行了測試，相比于依賴于增量式解碼的 LLM，SpecInfer 可以使推理延遲降低 1.9-2.8 倍。

單次推理平均步長（LLM：OPT-13B + SSMs：OPT-125M）

單次推理平均步長（LLM：LLaMA-7B + SSMs：LLaMA-160M）

匹配長度測試：分別使用 OPT 和 LLaMA 系列模型，測試 SpecInfer 中 LLM 的平均驗證通過序列長度，可以看出，隨著 SSM 數量的提升，在各個對話數據集上，LLM 的驗證通過長度均會得到提升，以 5 個 SSM 為例，OPT 和 LLaMA 在 5 個數據集上平均可達 3.68 和 2.67，相比于僅使用單一 SSM，分別提升 26.4% 和 24.8%。

更多更詳細的實驗結果可以參考論文原文：https://arxiv.org/abs/2305.09781

總結

SpecInfer 是首個基于「推測式解碼」的分布式 LLM 推理引擎，通過集成多個小模型，以及基于 token tree 的原創系統實現優化，可以幫助現有的主流 LLM 減少內存訪問需求，實現兩到三倍的無損推理加速，大幅降低推理成本。

作者介紹

SpecInfer 項目的指導老師是 Zhihao Jia，他目前在卡耐基梅隆大學計算機學院擔任助理教授。他的研究興趣主要包括面向機器學習、量子計算以及大規模數據分析的系統研究。此前他曾畢業于清華大學的姚班，博士畢業于 Stanford 大學，師從 Alex Aiken 和 Matei Zaharia，曾獲 Stanford Arthur Samuel Best Doctoral Thesis Award，NSF CAREER Asward 以及來自 Amazon, Google, Meta, Oracle, 以及 Qualcomm 的多項研究獎項，個人主頁：https://www.cs.cmu.edu/~zhihaoj2/。

孵化 SpecInfer 項目的主要是 CMU 的 Catalyst Group 實驗室，該實驗室由 Zhihao Jia 與 Tianqi Chen（陳天奇）在 CMU 共同主持，致力于集成來自于機器學習算法、系統、硬件等多方面的優化技術，構造自動化的機器學習系統。此前，該實驗室還推出了 MLC-LLM [5] 等開源項目，推進 LLM 大模型相關系統的研究和應用。實驗室主頁：https://catalyst.cs.cmu.edu。

論文的共同一作分別是 Xupeng Miao（博士后研究員），Gabriele Oliaro（博一）以及 Zhihao Zhang（博一），均來自于 CMU Catalyst Group 團隊。其中，Xupeng Miao 博士畢業于北京大學，主要研究方向包括機器學習系統、數據管理和分布式計算，曾獲 VLDB2022 最佳可擴展數據科學論文獎、2022 年 ACM 中國優博獎、2022 年世界人工智能大會（WAIC）云帆獎等榮譽，個人主頁：https://hsword.github.io。