1.   什么是大模型推理引擎

大模型推理引擎是生成式語言模型運轉(zhuǎn)的發(fā)動機，是接受客戶輸入 prompt 和生成返回 response 的樞紐，也是拉起異構(gòu)硬件，將物理電能轉(zhuǎn)換為人類知識的變形金剛。

大模型推理引擎的基本工作模式可以概括為，接收包括輸入 prompt 和采樣參數(shù)的并發(fā)請求，分詞并且組裝成 batch 輸入給引擎，調(diào)度 GPU 執(zhí)行前向推理，處理計算結(jié)果并轉(zhuǎn)為詞元返回給用戶。

• 和人類大腦處理語言的機制類似，大模型首先會把輸入的 prompt 進行統(tǒng)一理解，形成具有記憶能力的上下文。這個階段通常稱為 Prefill 階段。
• 在結(jié)束 Prefill 階段之后，大模型引擎會根據(jù)生成的上下文不停地推斷下一個可能出現(xiàn)的詞語，如此往復(fù)循環(huán)，直到遇到停止符或者滿足采樣參數(shù)中的停止條件。這是一個自回歸過程，通常稱為 Decoder 階段。

由于 Prefill 階段和 Decoder 階段所完成的任務(wù)不同，通常來講，會從用戶視角出發(fā)使用 SLO（Service Level Object）: TTFT（Time To First Token）和TPOT（Time Per Output Token）去評測引擎。

• TTFT 就是首 token 延遲，用于衡量 Prefill 階段的性能。也就是用戶發(fā)出請求之后，收到第一個詞元返回的間隔，也就是系統(tǒng)的反應(yīng)時間。對于客戶來說，這個指標(biāo)越低越好。
• TPOT 就是出字間隔，用于衡量 Decoder 階段的性能。也就是每生成兩個詞元之間的間隔。通常需要比人眼閱讀文字的速度要快，這個指標(biāo)同樣也是越低越好。

當(dāng)然，只用這些 SLO 并不能完全評測推理引擎對資源的使用狀態(tài)，所以，和其他使用異構(gòu)資源的系統(tǒng)一樣，會使用吞吐來評測引擎對資源的使用效率，常用的指標(biāo)就是極限出字率。

極限出字率 TPS（Tokens Per Second ）就是系統(tǒng)在滿載的情況下，使用所有可用的資源在 1s 內(nèi)可以生成的詞元的最大數(shù)量。這個指標(biāo)越高，代表硬件的效率越高，可以支持的用戶規(guī)模就越多。

目前市面上流行的推理引擎有很多，比如說 vLLM、SGLang、LMDeploy、TRT-LLM 等。其中 vLLM 是業(yè)界第一個完美解決了大模型不定長特性來各種問題的推理引擎，也是市面上使用最多，社區(qū)最活躍的推理引擎。

vLLM 的原創(chuàng)性高效顯存管理、高吞吐、極易用、易拓展、模式眾多、新特性支持快，社區(qū)活躍等特性是其受歡迎的原因。但是，vLLM 對復(fù)雜調(diào)度邏輯的處理沒有做到極致，引入了大量的 CPU 操作，拉長了 TPOT。TPOT 的拉長會降低用戶體驗，降低了出字率，造成了 GPU 資源浪費。

2.   影響 TPOT 的罪魁禍?zhǔn)?—— Token 間間隔

區(qū)別于小模型推理以 batch 為最小推理單位，大模型推理的最小單位是 step。這也是由大模型推理中自回歸的特點所決定的。

每一次 step 會給 batch 內(nèi)部的每個請求生成一個詞元，如果有請求生成了結(jié)束符，那么這個請求將會提前結(jié)束，并且從下個 step 的 batch 中剔除，空余出來的資源將會被引擎動態(tài)的分配給其余正在排隊的請求。用戶可以感知到的觀測指標(biāo) TPOT，就是每次 step 的執(zhí)行時間。

每個 step 的執(zhí)行邏輯可以簡單的概括為一下兩部分：前向推理和 Token 間間隔。

• 前向推理是調(diào)用 GPU 計算資源對 Transfomer 結(jié)構(gòu)進行運算的過程，是一個典型的 GPU 密集計算型任務(wù)。
• Token 間間隔，則負(fù)責(zé)做詞元拼接、結(jié)束檢測、用戶響應(yīng)、請求調(diào)度、輸入準(zhǔn)備等工作，是典型的 CPU 邏輯密集型任務(wù)。

優(yōu)化推理引擎的終極目標(biāo)其實就是，極限提升前向推理的吞吐，同時極限壓縮 Token 間間隔，最終提高極限出字率。

然而，vLLM 的實現(xiàn)中，這兩者天然存在著矛盾。極限提升前向推理的吞吐，（即充分發(fā)揮 GPU 算力）要求在適當(dāng)范圍內(nèi)盡可能增加 batch 內(nèi)的請求數(shù)。然而更多的請求數(shù)卻拉長了 Token 間間隔，這樣不僅會使 TPOT 拉長，還會導(dǎo)致 GPU 斷流，出現(xiàn)空閑。在最差的情況下（比如 batch 為 256），Token 間間隔和前向推理時間幾乎相同，GPU 的利用率只有 50%-60%。

為了提升極限出字率，同時確保高 GPU 利用率，優(yōu)化 Token 間間隔成為了提升推理速度的關(guān)鍵。

3.   百度百舸 AIAK 優(yōu)化 Token 間間隔的方案

百度百舸的 AI 加速套件 AIAK 基于 vLLM ，在優(yōu)化 TPOT 持續(xù)發(fā)力，并且始終保持著對社區(qū)在同周期的技術(shù)領(lǐng)先。

3.1.    解決方案 1：多進程架構(gòu)

這個方案的目標(biāo)是盡可能縮短 Token 間間隔，將 detokenizer 所耗費的時間從 TPOT 中拿去。

我們發(fā)現(xiàn)在處理輸入請求和生成返回的過程中，tokenize/detokenize 過程（token id 和字符串的轉(zhuǎn)換）是完全可以獨立于 GPU 推理運算的邏輯操作。

所以，我們借助 NVIDIA Triton 框架，將 tokenize/detokenize 的過程從推理流程中抽象出來作為單獨的 Triton 模型部署，借助 Triton 的 ensemble 機制，把串行過程轉(zhuǎn)變?yōu)?3 階段（ 3 進程）流水，實現(xiàn)了 tokenize/detokenize 和 GPU 推理 overlap，有效縮短了 Token 間隔時間。盡管這個優(yōu)化只把 Token 間間隔中一部分 CPU 操作消除了，但是依然有將近 10% 的收益。

圖片

3.2.    解決方案 2：靜態(tài) Slot 方案

這個方案主要改造了 vLLM 的調(diào)度邏輯，全方位優(yōu)化了詞元拼接、結(jié)束檢測、用戶響應(yīng)、請求調(diào)度、輸入準(zhǔn)備，提高了各個模塊的并行效率，實現(xiàn)了對上一個方案中的「剩余部分」耗時的壓縮。

我們發(fā)現(xiàn) vLLM 的調(diào)度邏輯是面向全局視角的。也就是說每個 step 的調(diào)度都會從全局中進行重新篩選，相當(dāng)于當(dāng)前 step 結(jié)束之后，調(diào)度器會把當(dāng)前 batch 中的句子「放回」全局請求池子中，然后在下一個 step 開始前，從這個全局池子中「取回」適當(dāng)請求進行運算，這一放一取引入了額外的 overhead。

為了實現(xiàn)全局調(diào)度，vLLM 在詞元拼接等其他環(huán)節(jié)引入了大量的 for 循環(huán)去串行的處理每個請求，由于這些操作都在發(fā)生在 CPU 上，導(dǎo)致在輸入打包過程中，必須要引入耗時較長的 host to device 操作。

事實上，step 之間的很多信息是可以復(fù)用的（每次放回去的請求和取回來的請求很大一部分是重復(fù)的）。也正是基于這個洞見，百度百舸的 AIAK 把 GPU 每次可以迭代的 batch 當(dāng)成一批固定的 slot，一旦某個請求被調(diào)度到某個 slot 后，在完成請求所有推理迭代之前，都不會被喚出，也正是有了這些固定 slot 抽象，AIAK 實現(xiàn)了：

• 將全局調(diào)度改造為局部調(diào)度。也就是在下一個 step 調(diào)度時，最大程度復(fù)用上一個 step 的信息，避免全局搜索，只做增量調(diào)度。
• 串行轉(zhuǎn)并行。也正是有了 slot 的引入，詞元拼接、結(jié)束檢測等這些原本串行的操作可以用 CUDA Kernel 做并發(fā)處理，耗時從 ms 級別降低到 us 級別。
• 避開 host to device 操作。輸入打包的工作得以復(fù)用前序的顯存，有效避開了 host to device 操作。

圖片

3.3.    方案 3：異步化執(zhí)行

多進程架構(gòu)將邏輯上容易獨立的部分解耦到其他進程做流水并行，靜態(tài) Slot 方案則直面 token 間耗時問題，優(yōu)化調(diào)度模式壓榨各個環(huán)節(jié)的耗時。有了這兩個方案，Token 間間隔已經(jīng)從 35ms 降低到 14ms，GPU 的利用率已經(jīng)從 50% 提升到了 75%，但是距離 100% 的 GPU 利用率和零耗時 Token 間間隔的目標(biāo)還有不少距離。

百度百舸 AIAK 通過異步調(diào)度模式，將前一個方案中的「剩余部分」全部取出，最終實現(xiàn)了上述極限目標(biāo)。

簡單來講，就是將 CPU 操作密集的 Token 間間隔和 GPU 計算密集的前向推理完全分開到兩條流水線上做二級流水并行。

• 從邏輯上來講，核心調(diào)度邏輯擺脫了對前向推理的同步依賴，實現(xiàn)異步化調(diào)度。
• 從效果上來說，GPU 避免 token 間同步導(dǎo)致的斷流問題，處于一直繁忙狀態(tài)，實現(xiàn)了推理過程中 100% 利用率和 0 Token 間間隔。

為了簡化實現(xiàn)，我們將操作相對簡單的前向推理當(dāng)做一個任務(wù)放在后臺線程中進行運行，主線程則運行核心的復(fù)雜的調(diào)度邏輯。兩個線程通過一個隊列進行交互，分別沖當(dāng)生產(chǎn)者和消費者，通過線程信號量和 GPU 流上的事件進行信號同步，實現(xiàn)二級流互相 overlap。

圖片

和其他任何使用 GPU 類似的硬件作為加速器的系統(tǒng)一樣，追求 100% 的利用率一直是所有工程師的終極目標(biāo)。百度百舸的 AI 加速套件 AIAK 在優(yōu)化 TPOT，同時打滿 GPU 利用率這一目標(biāo)上經(jīng)歷漫長而又艱辛的探索，最終才徹底實現(xiàn)了 0 Token 間間隔和 100% 利用率這一目標(biāo)。

當(dāng)然，除去在這個過程中使用的諸多巧妙的優(yōu)化手段外，百度百舸的 AIAK 還在量化、投機式、服務(wù)化、分離式、多芯適配等領(lǐng)域做了大量工作，致力于實現(xiàn)一個適用于全場景、多芯片、高性能的推理引擎，助力用戶在「降低推理成本，優(yōu)化用戶體驗上」更上一層樓。