大模型推理優(yōu)化實踐：KV cache復(fù)用與投機(jī)采樣

發(fā)布于 2024-4-8 10:49

瀏覽

0收藏

1、背景

RTP-LLM 是阿里巴巴大模型預(yù)測團(tuán)隊開發(fā)的大模型推理加速引擎，作為一個高性能的大模型推理解決方案，它已被廣泛應(yīng)用于阿里內(nèi)部。該引擎與當(dāng)前廣泛使用的多種主流模型兼容，并通過采用高性能的 CUDA 算子來實現(xiàn)了如 PagedAttention 和 Continuous Batching 等多項優(yōu)化措施。RTP-LLM 還支持包括多模態(tài)、LoRA、P-Tuning、以及 WeightOnly 動態(tài)量化等先進(jìn)功能。

隨著大模型的廣泛應(yīng)用，如何降低推理延遲并優(yōu)化成本已成為業(yè)界關(guān)注的焦點。我們不斷地在這一領(lǐng)域內(nèi)探索和挖掘新方法。在本文中，我們將詳細(xì)介紹兩種在業(yè)務(wù)中實踐的優(yōu)化策略：多輪對話間的 KV cache 復(fù)用技術(shù)和投機(jī)采樣方法。我們會細(xì)致探討這些策略的應(yīng)用場景、框架實現(xiàn)，并分享一些實現(xiàn)時的關(guān)鍵技巧。

2、多輪對話復(fù)用KV cache

在淘寶問問中，有兩類多輪對話的場景：一是問答類的場景，它每次請求模型時會拼接之前的問答；二是LangChain應(yīng)用，它在模型生成結(jié)束后會調(diào)用外部插件，拼接插件返回的結(jié)果再次請求模型。這兩類場景共同的問題是：隨著對話輪數(shù)的增加，請求長度變長，導(dǎo)致模型的First Token Time（下稱FTT）不斷變長。

模型的FTT變長，本質(zhì)上是因為第一次進(jìn)入模型時，越來越多的token需要生成KV cache。考慮到這兩種多輪對話場景存在一個共同點：前一輪對話的輸出構(gòu)成后一輪對話輸入的一部分，或者存在較長的公共前綴。且大部分自回歸模型（除了chatglm-6b）的Attention Mask都是下三角矩陣：即某一位置token的注意力與后續(xù)token無關(guān)，因此兩輪對話公共前綴部分的KV cache是一致的。進(jìn)而能夠想到的解決辦法是：保存上一輪對話產(chǎn)生的KV cache，供下一輪對話時復(fù)用，就能減少下一輪需要生成KV cache的token數(shù)，從而減少FTT。根據(jù)這個思路改進(jìn)前后的模型如下：

大模型推理優(yōu)化實踐：KV cache復(fù)用與投機(jī)采樣-AI.x社區(qū)

2.1 框架設(shè)計

用戶請求對應(yīng)的KV cache存放在機(jī)器顯存中，因此不同輪次的對話需要請求同一臺機(jī)器，才能復(fù)用KV cache。但是在生產(chǎn)環(huán)境中，模型部署在由多臺機(jī)器組成的機(jī)器集群，用戶層的請求由統(tǒng)一域名服務(wù)轉(zhuǎn)發(fā)到機(jī)器集群中某一臺機(jī)器上，這樣的架構(gòu)設(shè)計導(dǎo)致不同輪對話命中同一臺機(jī)器的概率微乎其微。

最直觀的解決辦法是讓用戶去記錄首次請求的機(jī)器信息，并將后續(xù)請求同一臺機(jī)器。這個方法可行但是不合理，用戶不僅需要感知機(jī)器集群的具體信息，還需要對自己鏈路做大量改造；進(jìn)而能想到的辦法是增加一層轉(zhuǎn)發(fā)層，用戶將多輪請求攜帶同樣的標(biāo)識id并發(fā)送給轉(zhuǎn)發(fā)層，轉(zhuǎn)發(fā)層感知集群信息并匹配標(biāo)識id和下游機(jī)器。這樣不同輪對話就能打到同一臺存有KV cache的下游機(jī)器。至于如何在轉(zhuǎn)發(fā)機(jī)器間同步匹配信息，可以使用分布式數(shù)據(jù)庫記錄，我們采樣的方法是使用統(tǒng)一的哈希算法，將相同id哈希到固定的機(jī)器。只要選擇合適的哈希算法，就能在機(jī)器集群負(fù)載均衡的同時讓多輪對話命中同一臺機(jī)器。

大模型推理優(yōu)化實踐：KV cache復(fù)用與投機(jī)采樣-AI.x社區(qū)

在底層實現(xiàn)上，復(fù)用KV cache的邏輯和P-Tuning v2在實現(xiàn)上非常相似，通過復(fù)用參數(shù)，我們使用PTuning的算子支持了KV cache復(fù)用。

2.2 總結(jié)與反思

我們在Qwen13B/int8量化/A10機(jī)器的條件下，對不同輸入和前綴長度的請求進(jìn)行了測試：

大模型推理優(yōu)化實踐：KV cache復(fù)用與投機(jī)采樣-AI.x社區(qū)

可以看到在復(fù)用KV cache功能極大程度的減少了FTT，并且歷史長度的變化對FTT的影響較小，F(xiàn)TT更多的取決于本次請求的輸入長度。并且除了多輪對話場景外，KV cache復(fù)用功能也擴(kuò)展到復(fù)用Ptuning前綴和長System Prompt的場景，降低FTT和顯存占用。

雖然復(fù)用KV cache的功能能夠顯著減少多輪對話場景下的FTT，但是在服務(wù)壓力過大時，存放歷史KV cache的顯存可能被新請求占用，導(dǎo)致后續(xù)請求出現(xiàn)cache miss請求時間變長，加劇服務(wù)壓力最后導(dǎo)致雪崩。目前我們已實現(xiàn)的解決方案是使用LRU算法優(yōu)先移除較舊請求的KV cache。未來進(jìn)一步的策略是參照vllm的思路，將過期的KV cache轉(zhuǎn)移到內(nèi)存，必要時重新加載至顯存。這種策略比重新計算快，有助于減輕極端情況下的請求延遲，防止服務(wù)雪崩。

03、投機(jī)采樣

3.1 介紹

投機(jī)采樣最早在2022年的Fast Inference from Transformers via Speculative Decoding提出，因為不久前的gpt4泄密而被更多人知道。投機(jī)采樣的設(shè)計基于兩點認(rèn)知：在模型推理中，token生成的難度有差別，有部分token生成難度低，用小參數(shù)草稿模型（下簡稱小模型）也能夠比較好的生成；在小批次情況下，原始模型（下簡稱大模型）在前向推理的主要時間在加載模型權(quán)重而非計算，因此批次數(shù)量對推理時間的影響非常小。

基于以上兩點認(rèn)知，投機(jī)推理的每一輪的推理變成如下步驟: 1. 使用小模型自回歸的生成N個token 2. 使用大模型并行驗證N個token出現(xiàn)的概率，接受一部分或者全部token。由于小模型推理時間遠(yuǎn)小于大模型，因此投機(jī)采樣在理想的情況下能夠?qū)崿F(xiàn)數(shù)倍的推理速度提升。同時，投機(jī)采樣使用了特殊的采樣方法，來保證投機(jī)采樣獲得的token分布符合原模型的分布，即使用投機(jī)采樣對效果是無損的。

大模型推理優(yōu)化實踐：KV cache復(fù)用與投機(jī)采樣-AI.x社區(qū)

上圖是投機(jī)采樣的運行過程，每一行的綠色token代表小模型生成并被大模型接受的部分，紅色token是小模型生成但被大模型拒絕的部分，藍(lán)色token是大模型根據(jù)最后接受token的logits重新采樣出來的部分。由上可以看到使用投機(jī)采樣，在合適的場景下能夠大幅提高每輪生成的token數(shù)，降低平均單個token生成時間。

3.2 設(shè)計思路

我們在RTP-LLM中基于論文的思路，使用大小模型進(jìn)行了投機(jī)采樣的實踐。在代碼設(shè)計上我們一方面考慮系統(tǒng)的可維護(hù)性，希望這部分能夠和原始流程解耦；其次投機(jī)采樣優(yōu)化需要與其他優(yōu)化正交，使投機(jī)采樣時兩個模型都能夠用上FT的其他優(yōu)化。最后我們的設(shè)計是為投機(jī)采樣封裝了一層編排層，對外提供統(tǒng)一的API，在內(nèi)部組織參數(shù)順序調(diào)用正常流程。

大模型推理優(yōu)化實踐：KV cache復(fù)用與投機(jī)采樣-AI.x社區(qū)

3.3 性能評估

在實現(xiàn)過程中，我們著重關(guān)注投機(jī)采樣引入的額外負(fù)擔(dān)。我們希望做到在系統(tǒng)每輪接受token數(shù)較少的情況下，也能有與原始模型相近的表現(xiàn)。在實踐中，我們測得額外時間消耗主要有兩塊：小模型順序生成token引入時間和采樣。

首先最直觀的額外消耗，就是小模型推理所占用的時間。在小模型順序生成N個token時，會有N*T_{small\_model}的時間，在N較大且接受token數(shù)少的情況下，這部分開銷會非常大。值得一提的時最初我們假設(shè)模型消耗時間和參數(shù)規(guī)模成正比，而實際上這個猜測是錯誤的。我們測試得到在參數(shù)規(guī)模減少的情況下, lm_head在模型調(diào)用的占比會顯著增加。以下是Qwen1.8B和Qwen13B在A10/half條件下，單個token在Transformer網(wǎng)絡(luò)(transformer_layer)和輸出層(lm_head)的時間對比:

大模型推理優(yōu)化實踐：KV cache復(fù)用與投機(jī)采樣-AI.x社區(qū)

造成上述比例不一致原因在于：模型參數(shù)規(guī)模從13B變成1.8B時, Transformer網(wǎng)絡(luò)在層數(shù)(40 -> 24)和權(quán)重大小(5120 -> 2048)兩個維度減少，而輸出層的參數(shù)僅從[5120, 152064]變成[2048, 152064] 。同時因為模型詞表通常很大，因此輸出層的時間通常也比較長。除了輸出層的影響以外，在一些情況下小模型矩陣乘對硬件的利用率并沒有大模型這么高，因此在選擇小模型時需要對這部分開銷進(jìn)行更謹(jǐn)慎的估計。

其次重復(fù)多次的采樣也引入了巨大的開銷。從上面流程圖可以看到，在一輪投機(jī)采樣流程中，需要進(jìn)行N次小batch采樣和1次大batch采樣。我們以a10/half/vocab_size=152064/top_k=0.5/top_p=0.95的情況下用huggingface采樣邏輯進(jìn)行了測試：

大模型推理優(yōu)化實踐：KV cache復(fù)用與投機(jī)采樣-AI.x社區(qū)

假設(shè)在原始請求batch為2, 投機(jī)采樣每次出5個token的條件下，需要1.15*5+1.47=7.221.15?5+1.47=7.22毫秒的時間，這接近上述1.8B小模型一次推理的總時長。

好在FT的采樣流程針對存在top_k參數(shù)的情況，通過融合算子對原版（Huggingface流程）進(jìn)行了優(yōu)化，改進(jìn)后的流程分成兩步：對維度是[batch, vocab]的輸入進(jìn)行TopK采樣后，使用輸出維度是[batch, k]的tensor進(jìn)行后續(xù)流程；省略TopP步驟，直接在采樣過程中對TopP進(jìn)行判斷。改進(jìn)前后的流程對比如下：

大模型推理優(yōu)化實踐：KV cache復(fù)用與投機(jī)采樣-AI.x社區(qū)

改進(jìn)后的流程不影響結(jié)果分布，且大幅度減少了計算量和kernel數(shù)量，極大程度減少了采樣需要的時間。我們測試優(yōu)化后的采樣流程需要的時間是原來的1/10。

3.4 總結(jié)

我們在店鋪起名和文案生成兩類任務(wù)，對原模型和投機(jī)采樣模型進(jìn)行了性能對比。其中原模型是int8量化的Qwen13B模型，投機(jī)采樣使用量化后的Qwen13B和Qwen1.8B模型，在A10機(jī)器測試結(jié)果如下：

大模型推理優(yōu)化實踐：KV cache復(fù)用與投機(jī)采樣-AI.x社區(qū)

在兩類任務(wù)下投機(jī)采樣對模型均有加速，效果隨接受token數(shù)和輸入token長度變化。我們測得在使用上述條件每輪生成5個token的情況下，短序列跑一輪需要60ms，長序列跑一輪需要70ms。而原模型跑一輪需要30ms，因此長/短序列只有在拒絕全部token的情況下會劣于原模型，其他情況則是與原模型平均token時間相近或者優(yōu)于原模型。由于測試條件限制，相比原論文的大小模型比例（70B:6B），我們大模型和小模型的規(guī)模（7B:1.8B）更接近，并且Qwen中文模型的詞表大小是152064，相比其他模型（如Llama詞表大小是32000）大了數(shù)倍，這也額外開銷時間變大，因此在其他測試場景下應(yīng)該還能有更好的表現(xiàn)。

4、存在的問題

上文我們著重介紹了兩個優(yōu)化對推理速度的影響，但除了推理速度外，并行度也是影響大模型吞吐的重要因素。影響并行度的主要因素是顯存，大模型的顯存占用分三塊：模型權(quán)重占用顯存、運行時顯存和KV cache顯存，KV cache顯存越多，模型能夠同時承載的請求數(shù)越多，并行度越大。

序列長度較長時，對運行時顯存運行最大的是Softmax Buffer，它的大小和序列的平方成正比，知名優(yōu)化FlashAttention除了能降低模型第一次運行時間外，更重要的作用是消除了Softmax Buffer對顯存的占用。但比較遺憾的是，F(xiàn)lashAttention優(yōu)化的開源實現(xiàn)要求Attention計算的QKV維度一致，而KV csache復(fù)用和投機(jī)采樣都未滿足這個條件，導(dǎo)致對顯存有額外的占用。除此之外，投機(jī)采樣因為要額外加載小模型的權(quán)重，且運行時需要多保存一份小模型的KV cache，還需要額外的顯存。

5、總結(jié)與致謝

以上是我們在大模型推理上做的一些優(yōu)化嘗試，有根據(jù)業(yè)務(wù)場景和實際問題的，也有參考論文實現(xiàn)的，并且都取得了一定的加速效果。但是從極致性能的角度，我們做的還遠(yuǎn)遠(yuǎn)不算完美，這些功能在算子層和框架層都還有優(yōu)化空間，這些是我們后續(xù)需要改進(jìn)的。

除了上述介紹的功能外，RTP-LLM還支持了非常多的功能，和上文相關(guān)的對System Prompt進(jìn)行緩存的Multi Task Prompt復(fù)用Medusa投機(jī)采樣，以及動態(tài)LoRA和不規(guī)則剪枝模型支持。未來我們也會持續(xù)的添加新功能，優(yōu)化底層算子性能，打造更好的大模型推理框架。

我們的項目主要基于FasterTransformer，并在此基礎(chǔ)上集成了TensorRT-LLM的部分kernel實現(xiàn)。FasterTransformer和TensorRT-LLM為我們提供了可靠的性能保障。Flash-Attention2和cutlass也在我們持續(xù)的性能優(yōu)化過程中提供了大量幫助。我們的continuous batching和increment decoding參考了vllm的實現(xiàn)；采樣參考了hf transformers，投機(jī)采樣部分集成了Medusa的實現(xiàn)，多模態(tài)部分集成了llava和qwen-vl的實現(xiàn)。感謝這些項目對我們的啟發(fā)和幫助。