1、背景

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

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

2、多輪對話復用KV cache

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

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

2.1 框架設計

用戶請求對應的KV cache存放在機器顯存中，因此不同輪次的對話需要請求同一臺機器，才能復用KV cache。但是在生產環境中，模型部署在由多臺機器組成的機器集群，用戶層的請求由統一域名服務轉發到機器集群中某一臺機器上，這樣的架構設計導致不同輪對話命中同一臺機器的概率微乎其微。

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

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

2.2 總結與反思

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

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

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

03、投機采樣

3.1 介紹

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

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

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

3.2 設計思路

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

3.3 性能評估

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

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

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

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

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

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

改進后的流程不影響結果分布，且大幅度減少了計算量和kernel數量，極大程度減少了采樣需要的時間。我們測試優化后的采樣流程需要的時間是原來的1/10。

3.4 總結

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

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

4、存在的問題

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

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

5、總結與致謝

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

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

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

相關資料

[01] FasterTransformer

??https://github.com/NVIDIA/FasterTransformer??

[02] TensorRT-LLM

??https://github.com/NVIDIA/TensorRT-LLM??

[03] Flash-Attention2

??https://github.com/Dao-AILab/flash-attention??

[04] cutlass

??https://github.com/NVIDIA/cutlass??

[05] vllm

??https://github.com/vllm-project/vllm??

[06] hf transformers

??https://github.com/huggingface/transformers??

[07] Medusa

??https://github.com/FasterDecoding/Medusa??

[08] llava

??https://github.com/haotian-liu/LLaVA??

[09] qwen-vl

??https://github.com/QwenLM/Qwen-VL??