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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

圖片

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

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

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

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

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

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

圖片

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

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

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

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

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

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

圖片

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

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