AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com
本文作者袁鐿博士是騰訊公司專家工程師，負責無量系統和一念LLM等機器學習訓練和推理框架研發。
以 OpenAI 的 GPT 系列模型為代表的大語言模型（LLM）掀起了新一輪 AI 應用浪潮，但是 LLM 推理的高昂成本一直困擾著業務團隊。
騰訊 PCG 機器學習平臺中心自研了高效能 LLM 推理引擎：一念 LLM。在傳統的運算元融合，ContinousBatching 等推理加速技術的基礎上，透過視訊記憶體最佳化，非同步排程和計算複用等技術，在相同精度的推理中，一念 LLM 相比 vLLM，TensorRT-LLM 等著名開源框架的推理單價低 20%+。
另外，為了應對國外高階 GPU 卡供應不足的問題，一念 LLM 在高效能 LLM 推理框架領域第一次同時支援 Nvidia GPU 卡和華為 NPU 卡。目前一念 LLM 已在 QQ 智慧體等 PCG 主要的 LLM 業務場景上線。
本文以一個簡單的公式，逐步分析 LLM 推理的效能瓶頸，介紹當前 LLM 推理的相關技術，以及一念 LLM 的設計決策邏輯。
“夫一心具十法界，一法界又具十法界、百法界；一界具三十種世間，百法界即具三千種世間。此三千在一念心，若無心而已，介爾有心即具三千。”
-- 出自：佛教天台宗摩訶止觀卷五上（大四六・五四上）
“一念亦稱一心，指心念活動之最短時刻；三千表示世間與出世間一切善惡、性相等人、物差別之總和。一念三千即謂，於凡夫當下一念之中，具足三千世間之諸法性相。”
-- 出自：佛光大辭典 (慈怡法師主編) 詞條 “一念三千”
一念 LLM，取 “一念三千” 之意，寓意 “一念之間，用大模型生成世間永珍”。
簡介
以 OpenAI 的 ChatGPT 為代表的大語言模型（LLM）掀起了新一輪 AI 應用浪潮，業務團隊都在探索基於 LLM 重構現有應用或者構建新的 APP。大語言模型的大引數量導致了巨大的計算和視訊記憶體需求，使得 LLM 的請求推理成本高昂。LLM 推理框架成為 2023 年以來的業界研究熱點。當前有多個著名的開源專案，比如：UCBerkeley 的 vLLM 和 Nvidia 的 TensorRT-LLM。
這些框架實現了諸多業界先進的推理加速技術，比如：ContinousBatching、PagedAttention 等，但是也存在兩方面的問題：
1. 為了便於演算法人員使用和嘗試新技術，vLLM 採用了 Python 作為主要排程管理功能的實現語言，導致視訊記憶體管理和排程效率較低。
2. 主要支援 Nvidia 的 GPU 等國外主流廠商的硬體，對國產硬體沒有支援。國產硬體配套的推理框架，缺乏對業界最新的推理加速技術的支援。
一念 LLM 透過對異構硬體的底層抽象，構建統一的高效能排程管理，實現了：
1. 應用業界最新的推理加速技術，推理單價相比業界開源框架低 20%+。結合業務場景進行針對性最佳化，單價可以降低 60%+。
2. 將最新的技術同時應用到國外主流的 Nvidia GPU 和國產的華為 NPU 上，避免軟體技術被硬體供應能力影響。
一念 LLM 已開源，歡迎共建。程式碼地址：https://github.com/pcg-mlp/KsanaLLM
問題分析
為了構造一個高效能的 LLM 推理框架，我們需要從源頭上分析推理效能的瓶頸。我們從兩個方面來分析：1）排程和視訊記憶體管理；2）計算效能最佳化，類似運算元融合，量化等計算最佳化等。
排程與視訊記憶體管理
在一個 LLM 推理系統中，我們希望 token 的生成速度能夠最大化。由於 GPU 硬體的特性，將多個請求組合成一個大 batch 平行計算是 LLM 推理主要的計算速度提高手段。從 A100 的推理壓測結果圖可以給我們不少啟示：
圖 1 平行計算 token 數與 GPU 實際計算效率關係。圖片來源：https://github.com/microsoft/DeepSpeed/blob/master/blogs/deepspeed-fastgen/README.md
當前向推理的 token 數增大時，GPU 有效的計算量吞吐逐步增大，當到達 200Tflops 之後，趨於穩定。這個穩定的上限與硬體的最大 Tflops 引數（A100 的 float16 的標稱 flops 為 312Tflops）以及 LLM 推理運算元的實現效率有關。
圖 2 GPT-2 模型推理過程示例。圖片來源：https://medium.com/@joaolages/kv-caching-explained-276520203249
在 LLM 模型的推理過程大致分為 prefill 和 decoding 兩個階段。在 prefill 階段（圖 2 中 '$' 之前部分輸入，生成 'A' 的過程），prompt 的多個 token 被一起輸入給模型，輸入的 token 數量可能幾百或者幾千。在 decoding 階段（圖 2 中紅色輸入部分），模型每次輸入上一次生成的 token，生成下一個 token，迴圈這個邏輯直到回答結束。
從圖 1 中，我們可以標出兩個階段所處的工作區間。在輸入的 token 數超過 300 的情況下，prefill 階段可以處於 GPU 全力工作的區間，而由於 decoding 階段一個請求每次輸入的 token 只有一個，則處在 GPU 極大浪費的狀態。decoding 階段如果要達到 GPU 計算資源的充分利用，batch size 應該增大到 300 左右。然而，實際情況下由於 GPU 視訊記憶體的限制，batch size 遠遠小於這個數。
我們需要進一步分析視訊記憶體是如何影響 batch size 的。我們可以列一個簡化的公式來幫助分析：
其中 M 是模型引數佔用的視訊記憶體，α 是每個請求推理過程中的視訊記憶體佔用，BS 是 batch size，β 是每個 token 對應的 kv cache 所需的視訊記憶體，TN 是快取 kv cache 的 token 數量，Mem 是 GPU 的視訊記憶體大小。在不使用量化等手段的情況下，選定模型和 GPU 硬體後，β 和 Mem 是固定。如果要增大 batch size，就需要降低 M，α 和 TN。
M 主要由放在視訊記憶體上的引數量決定的。α 主要是由模型計算邏輯中的中間變數佔用的視訊記憶體空間決定的，而 TN 與 BS 相關，一個簡單的關係是如果 batch 內請求的 token 平均數量為 TA，那麼。γ 表示 batch 中不同請求之間 token 不能複用 kv-cache 的比例。所以，從視訊記憶體節省的角度最佳化系統的吞吐，就有下面兩條路徑：
最佳化 M：在對 latency 影響較小的前提下，將引數 offload 到記憶體或者其他儲存上。
最佳化 α：最佳化推理計算邏輯中的中間變數視訊記憶體佔用。
最佳化 γ：最佳化 batch 中不同請求之間複用的 kv-cache 比例。
計算效能最佳化
LLM 模型由於模型結構非常固定，尤其 ChatGPT 的成功讓比傳統 transformer 結構更簡單的 decoder-only transformer 結構模型成為當前的主流。這種穩定而簡單的結構讓手擼大運算元成為了 LLM 推理加速框架的首選，類似 FlashAttention 等針對單個大結構深度最佳化的運算元庫深受大家追捧。各個開源框架都紛紛推出自己的定製運算元，Nvidia 等硬體廠商也都提供了與自身硬體高度適配的運算元庫，甚至不惜為同一結構的不同引數大小模型單獨開發運算元。
對開源運算元的支援能力，決定了框架是否能有一個持平業界的基線效能。
方案設計
下面我們先簡單介紹一下一念的主要模組，稍後按照從前面提到的多個效能最佳化角度介紹一念 LLM 的設計。
一念 LLM 的基本結構
一念 LLM 主要由以下模組組成：
圖 3 一念 LLM 模組圖
記憶體 / 視訊記憶體統一管理模組負責分配和管理記憶體和視訊記憶體，其中最重要的功能是分配和管理 PagedAttention 機制所需的 Block 和推理過程中所需的臨時視訊記憶體。在後期，與排程配合，實現更細化的排程能力。
請求排程器模組負責排程多個請求執行，協調記憶體 / 視訊記憶體統一管理，實現推理過程的流水線化。
kv-cache 快取排程負責 kv-cache 在請求之間的快取複用。
加速運算元包括不同硬體的模型並行，計算量化，運算元融合等功能相關的運算元。包括了自研的高效能運算元，經過框架適配的開源框架優秀運算元。相關運算元隨著業界發展迭代。
統一抽象介面負責將不同硬體的運算元以相同的操作方式對接到執行流水線。採用是類似 Nvidia Cuda API 的介面。
LLM 模型是基於統一抽象介面和加速運算元庫來支援的 LLM 模型。
模型請求排程模組用於排程不同的請求到後端推理節點。與傳統機器學習推理不同，LLM 模型具有推理時間長，狀態資料大的特點，請求排程模組更加請求的特點和後端服務節點的狀態進行排程，最佳化系統效能。
ContinousBatching&&PagedAttention（最佳化有效 batch size）
在大語言模型推理過程中，一般會使用到 GPU 進行加速。在一個請求的依次生成 token 的過程中會有大量使用 kv-cache 來降低計算量，但是 kv-cache 本身會佔用 GPU 視訊記憶體資源。目前 LLM 推理的效能瓶頸主要是因為 LLM 引數量大導致的視訊記憶體頻寬瓶頸，為了提高服務吞吐，需要儘量使用多個請求組成一個大 batch 來推理，以充分利用 GPU 的計算能力。
在通常情況下，由於大語言模型計算過程中用到了一個自增長的 kv-cache，為了加速 GPU 的計算過程，通用方案（圖 4 (a)，典型代表為 FasterTransformer）都是按 batch 為單位排程執行。由於 batch 中不同的請求 token 數量差異大，batch 粒度的排程方式會導致 GPU 計算能力的浪費，後續的請求不能得到及時處理，影響推理服務的吞吐能力。在 batch 執行的後期，可以理解為有效輸出 token 的 batch size 在逐漸變小。
以圖 4 (a) 為例，到第 5 步時，只有兩個請求還在推理，到第 6 步，有效 batch size 就只有 1 個了。
圖 4 不同排程方案示意圖。
為了充分利用 GPU 的計算能力， 需要細化請求排程的粒度。於是有了按請求粒度排程的 ContinousBatching 技術，如圖 4 (b) 所示，第二個 batch 的第一條和第三條請求在第一個 batch 最後一條請求執行完之前就開始了執行，GPU 計算資源的利用效率得到了提升。Batch 越大，請求長度的差異也越大，ContinousBatching 對系統吞吐的提升就越大。
ContinousBatching 的技術被提出後，並沒有引起推理加速框架的爆發增長。其中最大障礙是原有的 GPU 計算中對 kv-cache 連續空間訪問方式，導致 ContinousBatching 在 token 生成後排程請求有很大的視訊記憶體操作開銷。為了解決這個問題，PagedAttention 技術提出了類似作業系統虛擬頁的視訊記憶體管理機制，將 kv-cache 的整個連續空間切分為多個連續塊（Block），使得按請求粒度的排程變得高效，讓 ContinousBatching 技術被廣泛應用。
為了實現 GPU 計算資源的充分利用，大語言模型推理框架必須要實現 ContinousBatching 功能，一念 LLM 有了請求管理器。在前面問題分析的時候提到過，在不同的場景下，排程器的最佳化邏輯不同，甚至需要設計比 ContinousBatching 更小粒度的排程策略。我們抽象出了排程策略介面，用於實現不同的排程策略。純 C++ 的實現讓排程的非同步邏輯更高效。
為了實現 PagedAttention 的功能，一念 LLM 設計了視訊記憶體 / 記憶體統一管理模組，同時為了便於後期實現多模型，Multi-LoRA，狀態快取等功能，視訊記憶體 / 記憶體統一管理模組收攏了一念 LLM 主要的記憶體和視訊記憶體操作。
多硬體運算元抽象（硬體和運算元決定系統的 TFlops 上限）
在國外高階卡進口受限的局面下，形成的新問題。目前業界最新的推理框架（比如：最早提出 PagedAttention 的 vLLM 和 Nvidia 的 TensorRT-LLM）主要支援 Nvidia 的 GPU 等國外主流廠商的硬體，對國產硬體缺乏支援。國產硬體配套的推理框架，缺乏對業界最新的推理加速技術的支援。目前相關的新技術主要集中在排程或者更高的演算法層面，與硬體關係不大，所以最合理的方式是使用統一的運算元抽象來遮蔽下層硬體差異，從而實現一次最佳化，所有硬體上可用。
在 Nvidia GPU 生態下，一念 LLM 的運算元庫包含了來自 FasterTransformer，vLLM，TensorRT-LLM，pytorch 的開源專案運算元，以及部分自研運算元。
在華為生態，推薦的使用方式是用華為生態軟體，使用圖最佳化等方式來加速，但是圖計算存在最佳化控制粒度的問題。在 LLM 這種相對穩定的模型結構上，也不能發揮計算圖最佳化的優勢。一念選擇了相對底層的 AscendC 介面來實現自定義運算元的方案。這套介面與 Nvidia Cuda 的介面類似，有 device，stream 等常用的物件介面。AscendC 介面當前在成熟度和效能方面與 Nvidia Cuda 還有不少差距。透過與華為共建和華為卡的廣泛使用，我們相信 AscendC 這層介面實現的 LLM 運算元效能會越來越好。
在運算元使用上，透過效能和效果兩個維度來選擇運算元。從效能方面，根據不同運算元在不同硬體上的效能特點，擇優選擇。與效能相對，有的業務場景會希望推理的結果與訓練的結果嚴格對齊，從而降低評估和效果調優成本。比如：要求生成的長文內容對齊。導致推理服務和訓練框架在長文字生成內容上不對齊的主要原因是推理過程普遍使用的 float16 的表示精度有限，不同運算元實現在數學上等價，但是實際精度誤差不同，而且誤差會隨著推理長度增長而累積，於是出現了不同運算元的推理結果在前幾十個 token 相同，然後結果差異越來越大的情況。
當出現這種情況時，框架需要在效能與效果之間進行 tradeoff，有時就會為了對齊效果，將對效果影響最大的運算元替換為效能更低的運算元。
Prefix Caching，基於 prefix-token 的 kv-cache 快取排程（最佳化 γ）
ContinousBatching 方案的排程仍然是請求粒度，並未對請求輸入內容進行針對性最佳化。如圖 1 (a,b) 所示，所有請求的前三個 token 都是 (1，2，3)，我們稱請求的相同輸入部分為 prefix-tokens。在當前的排程邏輯下，prefix-tokens 的計算在每個請求的計算中都會被執行，而在當前主流的 decoder-only 的模型結構下，prefix-tokens 的計算結果以及對後續 token 計算結果的影響是相同的，也就是說對 prefix-tokens 的計算只需要進行一次。
當前請求粒度的排程導致了重複計算，從而導致了 GPU 計算資源的浪費。而且這個浪費的計算比例會隨著 prefix-tokens 的佔比和 batch size 增大而增大。在類似角色扮演等應用場景中，prefix-tokens 的佔比可能達到 50% 以上，而 batch size 會超過 30，那麼將會有超過 50% 的計算被浪費掉。
要實現高效的 prefix-token 的視訊記憶體和計算複用，面臨以下問題：
常規的計算過程是以矩陣方式計算的，相關運算元的最佳化也都是基於矩陣的規則大小計算。如果要實現不規則的計算，需要重新開發和最佳化運算元，技術通用性和開發成本都非常高，可能新實現的運算元最終的效能與現有運算元差異巨大，得不償失。
排程上 batch 內的不同請求的相同輸入長度短，導致計算節省的收益小。
所以要實現收益的最大化，我們需要實現下面的最佳化：
排程請求到不同的 batch，實現 batch 內請求的相同 prefix-token 長度最大化。
在 batch 的輸入處理邏輯中，需要基於開源運算元，以最小的代價去掉相同輸入的計算。
排程上需要匹配視訊記憶體與計算的複用邏輯，讓計算的排程與視訊記憶體的管理協調一致。
基於這樣的需求，我們設計了以下的架構框架：
圖 5 Prefix Caching 功能模組關係圖。
總體上，為了提升 batch 內請求的相同 prefix-token 長度，增加了基於 prefix-token 分析的請求路由器模組。在排程器上改造為 prefix-token 與剩餘部分的兩階段排程，同時排程策略上針對 prefix-token 和 kv-cache 快取情況進行了最佳化。為了配合排程策略的執行，增加了 kv-cache 快取管理器模組，後期可以實現 kv-cache 快取在視訊記憶體 / 記憶體 / 外部儲存的三級排程管理能力。
在傳統的分散式系統中，請求路由器主要考慮後端服務節點的請求響應和負載情況進行請求分發。為了提高後端服務節點 prefix-tokens 快取的命中率，在路由模組上增加根據 prefix-tokens 路由的策略，構造了下圖所示的路由表。其中 PT 代表 prefix-tokens，用 PTi 表示第 i 個 prefix-tokens。同時我們用 S 表示請求的服務節點 Server，用 Sj 表示第 j 個節點。用 SS 表示多個服務節點的集合 ServerSet，用 SSk 表示第 k 個 ServerSet。
圖 6 Prefix token 路由表示例。
透過將不同 PT 對映到 SS 上，實現不同 PT 請求的服務擴縮容。同時透過控制單個服務節點所歸屬的 SS 數量來控制需要處理的 PT 種類，從而提高快取命中率。
在特徵批次處理的場景中，prefix-token 在輸入中的佔比 80%+。一念開啟 KV-cache 快取功能後的吞吐率提升 60%+，等效於單價下降 40%+。
CPU/GPU 混合推理（最佳化 M）
在 transformer 模型的執行過程中，業界常規的做法是將所有的運算元放到 GPU 上執行，相應的模型引數也被放到了視訊記憶體中。由於 LLM 模型引數量大，導致用於並行推理的視訊記憶體空間被擠壓。在業務常用的 7B，13B 模型中，模型的詞表有變大的趨勢，導致 token embedding 的引數量佔比較大。以 llama-13B 為例，原始詞表大小為 3.2 萬，token embedding 引數佔比為 1.2%。如果詞表大小擴充套件到 30 萬，embedding 引數佔總引數量的 11.8%。但是我們發現 token embedding 的操作並不是計算密集型的，而是一個典型的 sparse 查表操作。於是一念將 token embedding 引數放到記憶體中，用 CPU 執行 token embedding 操作，實現 CPU/GPU 混合推理，如下圖所示。在詞表大小為 30 萬的 llama-13B 模型上，提升吞吐率 10%+。
圖 7 Cpu/GPU 混合推理示意圖。
臨時視訊記憶體最佳化（最佳化 α）
在深度學習網路執行過程中，會使用到很多臨時變數來儲存中間結果。在不同的框架中，會有不同的臨時變數回收策略，一般是基於計算圖來最佳化的，在 LLM 模型的推理過程中，很多變數大小都是動態增長的，會導致計算圖的視訊記憶體最佳化失效。一念沒有采用計算圖的方式來進行推理，而是採用運算元拼接的方式直接描述模型，從而實現臨時變數的自管理。透過預先分配視訊記憶體然後重複使用的方式，最小化臨時變數的視訊記憶體消耗。
未來計劃
現在一念算是有了第一階段的起點，解決了排程最佳化和多硬體支援的基礎問題。
在國產硬體支援方面，目前只支援華為 NPU，後期還要支援騰訊自研的紫霄以及其他國產晶片。
在排程 / 視訊記憶體 / 運算元層面還需要根據業務場景和硬體的特點持續最佳化。同時在演算法技術層面，還不斷有一些新的方向出現。下面簡單說一下兩個有趣的方向：Speculative Decoding 和稀疏化。
Speculaitve Decoding：在前面的分析過程中，一直是以加大 batch size 的方式來提升服務整體的 throughput，其中的主要瓶頸點是視訊記憶體。可能存在下面的場景：a）視訊記憶體有剩餘，但是有不足以增加一個請求到 batch 中；b）請求很少，batch size 不能放大。SpeculativeDecoding 可以透過猜測多個可能的 token 輸出，然後並行驗證的方式，降低 latency，讓當前請求儘快結束，從而釋放出視訊記憶體空間來響應新的請求。這個過程中猜測準確率是關鍵。於是有多種預測方式，比如：UCBerkeley 的 Big Little Decoder 利用小模型來快速猜測，螞蟻金服的 Lookahead 框架基於 Trie-based retrieval 來進行猜測等。
稀疏化：大語言模型的大引數量和 kv-cache 都會帶來大的計算量和顯示卡儲存消耗，讓人不禁會問，這些儲存和計算是否都是必須的。圍繞各種問題，產生了很多稀疏化的嘗試。大致可以分為模型引數稀疏化和 kv-cache 稀疏化兩個方向。模型引數稀疏化在深度學習模型推理加速領域一直有比較多的研究，本文不再贅述。kv-cache 作為 transformer 結構引入的新變數，也有很多有意思的研究，比如：基於 kv-cache 內容壓縮的 GEAR，基於 token 重要性壓縮的 KeyFormer。
如果要讓這些技術成功落地，對視訊記憶體和排程管理都提出了更嚴苛的要求。
結語
大語言模型的能力越來越強，但大語言模型在應用場景中 ROI 正向仍然是一個非常挑戰的問題。LLM 推理在 LLM 應用成本中佔比大，任何小小的進步都能獲得不錯的成本收益，切實幫助業務實現更好 ROI。
在國外高階硬體供應不足的當下，統一框架以及國產硬體支援的可控亦是實現業務安全的必要路徑。在相關軟體生態不成熟的背景之下，會有很多困難。相信隨著國產硬體的成長，會越來越好。
一念 LLM，篳路襤褸，以啟山林。