在上一篇 博客文章 中，我們展示了通過將鍵值 (KV) 緩存卸載到 CPU 內存來重復使用該緩存如何在基于 x86 的 NVIDIA H100 Tensor Core GPU 上將首次生成令牌 (Time To First Token, TTFT) 的速度提升高達 14 倍，在 NVIDIA GH200 Superchip 上可將其提升高達 28 倍。在本文中，我們介紹了可以進一步推動 TTFT 加速的 KV 緩存重復使用技術和最佳實踐。
KV 緩存簡介?
許多任務（包括問答和代碼生成）都迅速采用了 LLM 模型。為了生成響應，這些模型首先將用戶的提示符轉換為 tokens，然后將其轉換為 dense vectors。隨后會進行大量的 dot-product operations，以數學方式對 tokens 之間的關系建模，并構建對用戶輸入的上下文理解。生成這種上下文理解的計算成本會隨著輸入序列的長度而二次增加。
此資源密集型流程會生成鍵和值，這些鍵和值會被緩存，以避免在生成后續 token 時進行重新計算。重復使用 KV 緩存可減少生成額外 token 所需的計算負載和時間，實現更快、更高效的用戶體驗。
重復使用 KV 緩存時，必須仔細注意其在內存中的剩余時間、在內存已滿時首先移除哪些組件，以及何時可重復用于新的傳入提示。優化這些因素可能會帶來 KV 緩存重復使用的性能增益。NVIDIA TensorRT-LLM 提供了三個關鍵功能，專門解決了這些問題。
早期 KV 緩存復用?
傳統的復用算法需要先完成整個 KV 緩存的計算，然后才能根據新用戶的提示重復使用其任何部分。在企業聊天機器人等場景中，系統提示（添加到用戶查詢中的預定義指令）對引導 LLM 根據企業指南做出響應至關重要，這種方法可能效率低下。
當大量用戶同時與聊天機器人交互時，每個用戶都需要單獨計算系統提示 KV 緩存。借助 TensorRT-LLM，我們可以重新使用實時生成的系統提示，使其在突發期間在所有用戶之間共享，而不是為每個用戶重新計算。這可以顯著加快推理速度，最高可達 5 倍。
靈活調整 KV 緩存塊大小?
在重用實現中，只能分配整個緩存內存塊以供重用。例如，如果緩存內存塊大小為 64 個令牌，KV 緩存為 80 個令牌，則僅存儲 64 個令牌以供重用，而其余 16 個令牌則需要重新計算。但是，如果將內存塊大小減少到 16 個令牌，則可在五個內存塊中存儲全部 64 個令牌，從而消除重新計算的需求。
當輸入序列較短時，這種效果最為明顯。對于長輸入序列，更大的塊會更有好處。顯然，您對 KV 緩存的控制越細致，就越能針對特定用例進行優化。
TensorRT-LLM 可對 KV 緩存內存塊進行細粒度控制，使開發者能夠將這些塊分割成 64 到 2 個令牌之間的較小塊。這可優化已分配內存的使用，提高復用率并改善 TTFT。在 NVIDIA H100 Tensor Core GPU 上運行 LLAMA70B 時，我們可以通過將 KV 緩存塊的大小從 64 個令牌減少到 8 個令牌，在多用戶環境中將 TTFT 速度提升高達 7%。
高效的 KV 緩存回收協議?
將 KV 緩存分割成較小的塊并移除未使用的塊對于內存優化可能有效，但這會引入依賴項復雜性。當使用特定塊生成響應并將結果存儲為新塊時，它會形成樹狀的依賴結構。
隨著時間的推移，跟蹤源塊（分支）使用情況的計數器可能會隨著依賴節點（葉子）的重復使用而過時。然后，移除源塊需要移除所有依賴塊，這將需要重新計算 KV 緩存以用于新用戶提示，從而增加 TTFT。
為應對這一挑戰，TensorRT-LLM 包含智能驅逐算法，該算法可以追蹤源節點中的依賴節點，并首先驅逐依賴節點，即使這些節點具有較新的重復使用計數器。這可以確保更高效的內存管理，同時防止不必要的依賴塊驅逐。
TensorRT-LLM KV 緩存重用入門?
在推理過程中生成 KV 緩存需要大量的計算和內存資源。高效使用 KV 緩存對于改善模型響應、加速推理和提高系統吞吐量至關重要。TensorRT-LLM 為希望進一步優化 TTFT 響應時間以實現峰值性能的開發者提供了高級復用功能。
要開始使用 TensorRT-LLM KV 緩存，請參閱 我們的 GitHub 文檔 。
?
?
