在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設置路由閾值，以及各個角色的參數和資源。