本文作者袁鐿博士是騰訊公司專家工程師，負責無量系統和一念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 模塊圖

1. 內存 / 顯存統一管理模塊負責分配和管理內存和顯存，其中最重要的功能是分配和管理 PagedAttention 機制所需的 Block 和推理過程中所需的臨時顯存。在后期，與調度配合，實現更細化的調度能力。
2. 請求調度器模塊負責調度多個請求執行，協調內存 / 顯存統一管理，實現推理過程的流水線化。
3. kv-cache 緩存調度負責 kv-cache 在請求之間的緩存復用。
4. 加速算子包括不同硬件的模型并行，計算量化，算子融合等功能相關的算子。包括了自研的高性能算子，經過框架適配的開源框架優秀算子。相關算子隨著業界發展迭代。
5. 統一抽象接口負責將不同硬件的算子以相同的操作方式對接到執行流水線。采用是類似 Nvidia Cuda API 的接口。
6. LLM 模型是基于統一抽象接口和加速算子庫來支持的 LLM 模型。
7. 模型請求調度模塊用于調度不同的請求到后端推理節點。與傳統機器學習推理不同，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 的顯存和計算復用，面臨以下問題：

1. 常規的計算過程是以矩陣方式計算的，相關算子的優化也都是基于矩陣的規則大小計算。如果要實現不規則的計算，需要重新開發和優化算子，技術通用性和開發成本都非常高，可能新實現的算子最終的性能與現有算子差異巨大，得不償失。
2. 調度上 batch 內的不同請求的相同輸入長度短，導致計算節省的收益小。

所以要實現收益的最大化，我們需要實現下面的優化：

1. 調度請求到不同的 batch，實現 batch 內請求的相同 prefix-token 長度最大化。
2. 在 batch 的輸入處理邏輯中，需要基于開源算子，以最小的代價去掉相同輸入的計算。
3. 調度上需要匹配顯存與計算的復用邏輯，讓計算的調度與顯存的管理協調一致。

基于這樣的需求，我們設計了以下的架構框架：

圖 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，篳路襤褸，以啟山林

本文轉自 機器之心 ，作者：機器之心

原文鏈接:??https://mp.weixin.qq.com/s/rlyJwaOfDfNYMZEH7kfKGA??