在SageMaker HyperPod上實現LLM推理的分離式預填充和解碼

在大型語言模型推理過程中，預填充和解碼是兩種截然不同的階段。預填充是計算密集型，它並行處理整個輸入提示以生成初始鍵值快取；而解碼是記憶體密集型，一次生成一個令牌，需要大量記憶體頻寬來訪問模型權重和不斷增長的鍵值快取。當預填充和解碼共享同一GPU時，長提示會阻塞每個併發請求的令牌生成。分離式預填充和解碼（DPD）透過將這兩個階段分配到透過彈性結構介面卡（EFA）連線的不同GPU池上，消除了這種干擾，從而可以獨立調整首令牌延遲和令牌間延遲，更可靠地控制尾部延遲，並防止長上下文預填充阻塞正在進行的解碼請求。

vLLM透過連續批處理和PagedAttention提高了單節點效率，但大規模部署時仍面臨多節點編排和路由最佳化方面的挑戰。本文展示瞭如何使用HyperPod推理運算子在Amazon SageMaker HyperPod上實現基於vLLM的DPD。

何時使用分離式推理

DPD對於長上下文、高併發的流式工作負載效果最為顯著，例如聊天助手、代理管道、文件分析端點和具有大量檢索上下文的檢索增強生成。在這些情況下，單個長提示會阻塞所有其他請求的解碼，導致令牌延遲峰值，而DPD透過構造消除了這一現象。當您的輸入提示經常超過4096個令牌、有多個併發使用者或請求、需要流式響應且對令牌交付一致性有要求，或者混合了長提示和短提示流量時，可以考慮使用DPD。對於GPU爭用不成問題的批次或離線工作負載、低併發部署或僅包含短提示的流量，共置部署更簡單。低於路由閾值的請求將直接傳送到解碼器，無需手動路由。

DPD至少需要一個預填充節點和一個支援RDMA的EFA網路的解碼節點。支援的例項型別包括P5和P6系列。

架構

HyperPod DPD實現基於vLLM生產棧路由器，並使用LMCache提供KV快取傳輸層。部署包含三個元件：智慧路由器、預填充器和解碼器。路由器負責令牌化並決定請求是否走分離路徑；預填充器為長提示計算KV快取並推送到解碼器；解碼器僅執行解碼，保證延遲穩定。KV快取傳輸透過LMCache、NIXL、libfabric和EFA四層堆疊實現，延遲極低。

部署概述

部署DPD需要HyperPod推理運算子3.2或更高版本，並選擇合適的支援RDMA的例項型別。本文詳細介紹了部署Llama 70B模型的步驟，包括安裝運算子、準備模型檢查點、配置部署清單等。部署清單中的關鍵欄位包括spec.pdSpec宣告預填充/解碼拓撲，routingThreshold設定路由閾值，以及各個角色的引數和資源。