Prefix Cache特性介紹
Prefix Cache 即前綴緩存,是一種用于優化大語言模型(LLM)推理性能的技術,主要應用于多輪對話、系統提示等具有大量共同前綴的場景。
原理
LLM 推理計算主要包括 Prefill 階段(Prompt 計算)和 Decode 階段。Prefill 階段為 Decode 階段準備 KV Cache,通常這些 KV Cache 只為單條推理請求服務,請求結束后會被清除。但在一些場景下,多次請求的 Prompt 可能共享同一個前綴,其 KV Cache 計算結果相同,Prefix Cache 就是將這些相同前綴的 KV Cache 存儲起來,供后續請求復用,避免重復計算,從而降低TTFT,提升推理計算性能。
實現方式
- 基于基數樹(Radix Tree):如 vLLM 中的 Prefix Caching 使用 RadixAttention 算法,該算法基于基數樹實現。基數樹的節點可以是一個變長序列,能動態分裂以滿足動態共享前綴的需求。vLLM 中通過 BlockSpaceManagerV1 類管理 block 分配,以 hash 碼作為物理 KV Block 的唯一標識,通過 prompt 中的 token_ids 來獲取 hash 值,確保不同 prompt 的 cache block 能獲取唯一 hash 碼,進而實現前綴緩存。
- 基于前綴樹(Trie):KVCache 的存取本質類似鍵 - 值對數據庫,可將 Trie 的每個節點代表一個 token,從根節點到某個節點的路徑表示一個 token 序列,以此處理前綴匹配問題。vLLM 會先將 tokens 根據 block_size 進行分塊,然后對每塊進行哈希,每個 chunk 的哈希值都包含了該 chunk 前綴的信息,以此模擬前綴樹的實現。
應用場景
Prefix Cache 適用于 common prefix 較多、prefill 計算占比較大的場景,例如多輪對話、system prompt、代碼補全等等。在多輪對話中,每一輪對話依賴歷史輪次對話的上下文,若使用 Prefix Cache,可避免歷史輪次中生成對話的重復計算,降低首 Token 的耗時。
緩存管理
由于顯卡顯存有限,當不同的 prefix 請求較多時,之前的 prefix cache 可能會被驅逐。此時可采用 LRU、LFU 等驅逐策略,通常會從后往前驅逐,盡可能復用 prefix cache。部分框架還提供了如 prefix cache offload 功能,在顯存 prefix cache 被驅逐時將其轉移到 cpu 內存上,命中時再加載回 gpu 中,以加速后續相同 prefix 的請求。
昇騰
當前大語言模型推理系統普遍采用KV Cache緩存機制,但該機制存在以下兩個問題:
隨著LLM支持的序列長度不斷增長,KV Cache所需要的顯存資源也急劇增加。
KV Cache只對當前session有效,如果跨session存在重復token序列的情況下無法實現復用。
Prefix Cache通過RadixTree保留session結束后的KV Cache,新的session請求在RadixTree中查找是否存在相同的Token序列,即可復用之前計算好的KV Cache,從而實現跨session的KV Cache復用。
其優勢主要包括:
更短的prefill時間:由于跨session的重復token序列對應的KV Cache可以復用,那么就可以減少一部分前綴token的KV Cache計算時間,從而減少prefill的時間。
更高效的顯存使用:當正在處理的sessions相互之間存在公共前綴時,公共前綴部分的KV Cache可以共用,不必重復占用多份顯存。
限制與約束
-
Atlas 800I A2 推理服務器和Atlas 300I Duo 推理卡硬件支持此特性。
-
僅Qwen2系列模型支持對接此特性。
-
當跨session公共前綴token數大于等于Page Attention中的block size,才會進行公共前綴token的KV Cache復用。
-
Prefix Cache支持的量化特性:W8A8量化與稀疏量化,其他量化特性暫不支持。
-
該特性不能和多機PD分離、Multi-LoRA、Function Call、長序列以及多機推理特性同時使用。
-
該特性可以和單機PD分離、并行解碼和SplitFuse特性同時使用。
-
開啟Prefix Cache特性需要配置補充參數。
效果
Prefix Cache 技術本身主要優化的是首 token 響應時間(TTFT, Time to First Token),但在特定條件下也能間接提升系統的吞吐量(Throughput)。下面從多個角度分析其影響及優化方向:
- Prefix Cache 對吞吐量的直接影響
正面影響
減少重復計算:當多個請求共享相同前綴時,Prefix Cache 避免了這些前綴的 KV Cache 重復計算,釋放了計算資源(如 GPU),使系統能同時處理更多請求。
降低排隊延遲:對于首 token 響應時間敏感的場景(如實時對話),更快的 TTFT 能減少請求在隊列中的等待時間,從而提高整體吞吐量。
負面影響
緩存占用顯存:Prefix Cache 需要額外存儲 KV Cache,可能導致系統可同時處理的請求數量減少(因顯存被占用)。
緩存管理開銷:維護緩存(如查找、驅逐策略)可能引入額外的計算開銷。 - 增大吞吐量的其他優化方法
(1)模型層面優化
量化(Quantization):如 INT8、INT4 量化減少模型參數量,降低顯存占用,提高并發請求數。
模型并行:將模型分布到多個 GPU 上,提升計算能力。
FlashAttention:優化注意力計算的內存訪問模式,減少顯存占用。
(2)系統層面優化
批處理(Batching):將多個請求合并為一個批次處理,提高 GPU 利用率。
連續批處理(Continuous Batching):如 vLLM 的 PagedAttention 技術,動態分配 KV Cache 內存,支持請求動態加入 / 退出批次。
預取(Prefetching):提前加載熱門模型或數據到顯存,減少 IO 等待。
(3)調度策略優化
優先級調度:對短請求優先處理,減少長尾延遲。
資源隔離:為不同類型的請求分配專用資源(如 GPU 核心)。 - 結合 Prefix Cache 與其他技術的案例
假設使用 vLLM 框架,可同時啟用 Prefix Cache 和 PagedAttention:
from vllm import LLM, SamplingParams# 初始化LLM時啟用PagedAttention
llm = LLM(model="deepseek-ai/deepseek-llm-7b-chat", gpu_memory_utilization=0.9)# 對于多輪對話,系統提示會被緩存
system_prompt = "You are a helpful, creative, accurate, and harmless AI assistant."
user_prompt1 = "What is machine learning?"
user_prompt2 = "Explain deep learning in simple terms."# 首次請求(包含system_prompt)會計算并緩存前綴
sampling_params = SamplingParams(temperature=0.7)
response1 = llm.generate(f"{system_prompt}\nUser: {user_prompt1}", sampling_params)# 第二次請求共享相同前綴,直接復用緩存
response2 = llm.generate(f"{system_prompt}\nUser: {user_prompt2}", sampling_params)
通過上述代碼,系統提示部分的 KV Cache 會被復用,減少了重復計算。同時,PagedAttention 技術允許高效管理 KV Cache 內存,進一步提升吞吐量。
4. 性能測試與調優建議
- 測試工具:使用框架提供的性能分析工具(如 vLLM 的 benchmark 腳本)或自定義負載測試。
- 關鍵指標:監控 TTFT、吞吐量(requests/second)、GPU 利用率、顯存占用。
- 調優方向:
增大緩存命中率:通過調整緩存大小、驅逐策略(如 LRU)提高緩存復用率。
平衡緩存與并發:根據模型大小和顯存容量,調整可緩存的最大前綴長度。
異步處理:使用異步 API(如 vLLM 的async_generate)實現非阻塞請求處理。
總結
Prefix Cache 能通過減少重復計算間接提升吞吐量,但需結合其他優化技術(如量化、批處理、高效內存管理)才能實現最大性能提升。在實際部署中,建議根據具體場景(如請求模式、模型大小)進行性能測試和參數調優。