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