大型語言模型 (LLMs) 正在推動當前研究和開發的重大進展。研究目標和方法發生了顯著變化，越來越多地朝著以 LLM 為中心的方法發展。然而，這些模型的使用成本很高，讓許多人無法大規模使用。因此，降低操作延遲成為一個重要挑戰，特別是在需要快速反應的動態應用中。
KV 快取用於 LLM 的自回歸解碼。它在推理的預填充階段儲存多頭注意力中的鍵值對。在解碼階段，新的 KV 對會被添加到記憶體中。KV 快取儲存了注意力機制中的中間鍵和值激活，從而將複雜度從二次降低到線性。KV 快取提高了效率，但隨著批次大小、序列長度和模型大小的增加而線性增長。KV 快取的記憶體大小增長超過了 GPU 的處理能力，將其轉移到 CPU 會引入多個瓶頸，增加延遲並降低吞吐量。
PCIe 接口成為一個限制因素，特別是在將快取從 CPU 轉移到 GPU 進行計算時。緩慢的 PCIe 接口可能導致延遲超過正常水平，造成 GPU 大量閒置時間。
之前的研究嘗試解決 PCIe 性能緩慢的問題，但這些方法往往因數據傳輸和 GPU 計算時間不匹配而失敗，尤其是在批次和上下文大小較大的情況下。其他方法依賴於 CPU 資源，這也成為了限制因素。本文討論了一種新穎的 PCIe 和 GPU 優化方法。
南加州大學的研究人員提出了一種高效的 CPU-GPU I/O 感知 LLM 推理方法，以優化 PCIe 的使用。他們利用部分 KV 快取的重新計算和異步重疊來解決加載大型 KV 快取的系統瓶頸。他們的過程涉及將較小的激活段轉移到 GPU，而不是轉移整個 KV 快取。然後，GPU 從這些較小的激活片段重建整個快取記憶體。關鍵在於計算注意力分數，以確保最小的信息損失。
作者提出了一種完全自動化的方法來確定重新計算和通信的分割。這項工作由三個模塊組成，以最小化 GPU 的延遲：
分析模塊
收集系統硬體信息，例如 PCIe 帶寬和 GPU 處理速度。
調度模塊
將問題表述為線性規劃任務，利用硬體信息和用戶配置來確定最佳的 KV 分割點。目標是最大化計算和通信過程之間的重疊。
運行時模塊
協調兩個設備之間的數據傳輸並管理記憶體分配。
調度模塊負責尋找最佳的 KV 分割，工作方式有兩種：
逐行調度：通過逐行執行計劃來減少延遲。在這裡，GPU 在剩餘激活異步加載的同時開始重建 KV 快取。逐列調度：通過在批次之間重用模型權重來最大化吞吐量並適應較大的批次推理。它將 KV 快取和激活的傳輸與多批次的多頭注意力計算重疊，而不是在每個批次中依次處理每一層。進一步使用六進程通信並行策略，運行時模塊使 GPU 計算和 CPU-GPU 通信能夠同時進行。
作者使用連接到 CPU 的 NVIDIA A100 GPU 進行了高效 LLM 推理的框架測試，通過 PCIe 4.0 x16 接口進行連接。實驗有兩個目標來評估框架的性能：
延遲導向工作負載
所提出的方法超越了基準，將延遲降低了 35.8%。
吞吐量導向工作負載
該方法相對於基準達到了高達 29% 的改進。
結論：
CPU-GPU I/O 感知的 LLM 推理方法有效地減少了延遲，同時提高了 LLM 推理的吞吐量。它利用部分 KV 快取的重新計算，並將其與數據傳輸重疊，以最小化 GPU 的閒置時間並提高效率。
查看論文。所有研究的功勞都歸於這個項目的研究人員。此外，別忘了在 Twitter 上關注我們，加入我們的 Telegram 頻道和 LinkedIn 群組。如果你喜歡我們的工作，你會喜歡我們的電子報。別忘了加入我們的 60k+ 機器學習 SubReddit。
🚨 [與我們合作]：‘下一期雜誌/報告 – 開源 AI 在生產中的應用’
新聞來源
本文由 AI 台灣 使用 AI 編撰，內容僅供參考，請自行進行事實查核。加入 AI TAIWAN Google News，隨時掌握最新 AI 資訊！
