PD分離：優化大語言模型推理效率

在大語言模型的推理過程中，Prefill 和 Decode 是兩個關鍵階段。隨著模型規模的不斷擴大，如何高效地處理這兩個階段的計算任務，成為了一個亟待解決的問題。

一、什么是 Prefill 和 Decode？

在大語言模型的推理過程中，Prefill 和 Decode 是兩個不同的階段，它們各自承擔著不同的任務。

• Prefill：Prefill 是推理過程的初始階段，它的主要任務是提前計算好輸入文本的 KV cache（鍵值緩存），并生成第一個 token。這個階段通常需要處理大量的計算任務，因此速度相對較慢。
• Decode：Decode 是在 Prefill 完成后進行的階段，它的任務是基于已經生成的 KV cache 不斷生成后續的 token。由于 Decode 階段可以利用已經計算好的 KV cache，因此速度相對較快。

初始邏輯問題

在傳統的推理流程中，Prefill 和 Decode 是順序執行的。然而，這種設計存在一個明顯的問題：當一些用戶正在處于 Decode 階段時，如果突然來了一個新的 Prefill 請求，系統會優先處理 Prefill 請求，這會導致正在 Decode 的用戶被迫中斷。這種中斷不僅會影響用戶體驗，還會降低系統的整體效率。

二、PD 分離的解決方案

為了解決上述問題，研究者們提出了 PD 分離的概念。PD 分離的核心思想是將 Prefill 和 Decode 分別放在兩個不同的 GPU 上運行，一個 GPU 專門負責 Prefill，另一個 GPU 專門負責 Decode。這樣可以避免 Prefill 和 Decode 之間的沖突，提高系統的整體效率。

PD 分離的核心問題

1. 如何傳輸 KV cache？

   • Pooling mode：在這種模式下，發送方將 KV cache 寫入到一個共享的 KV cache 池中，接收方從池中獲取所需的 KV cache。這種模式的優點是簡單易實現，但可能會因為池的共享訪問而導致延遲。
   • P2P mode：在這種模式下，發送方需要知道接收方是誰，并直接建立連接傳輸 KV cache。這種模式的效果通常比 Pooling mode 更好，但代碼設計會更加復雜。

2. 如何從 vLLM 中提取（和注入）KV cache？

   • Connector API：通過設計一個連接器 API，可以在模型前向傳播之前接收 KV cache（將 KV cache 注入到 vLLM 的分頁內存中），在模型前向傳播之后，從 vLLM 的分頁內存中提取 KV cache 并發送出去。

3. 何時發送請求到 P 和 D 節點？

   • First P then D：Prefill 先收到請求，處理完成后將結果發送給 Decode。
   • First D then P：Decode 先收到請求，如果發現沒有 KV cache，則請求 Prefill。

三、其他相關想法

除了 PD 分離，還有一些其他的想法可以進一步優化推理效率：

• Chunked Prefill：讓算子可以計算不同的 shape，從而提高 Prefill 階段的靈活性和效率。

四、相關工具和庫

隨著 PD 分離概念的提出，一些新的工具和庫也應運而生，例如：

• LM Cache：用于高效管理 KV cache 的工具。
• MoonCake：一種支持 PD 分離的框架。
• NIXL：用于優化 KV cache 傳輸和管理的庫。