分離式預填與解碼
跨裝置分離預填與解碼可加速 LLM 推理並改善效能
為了從輸入提示詞生成輸出詞元,LLM 推理會拆分為兩個階段:預填 (prefill) 和解碼 (decode)。預填會在輸入詞元上運行,並填入 KV 快取,然後進入逐個生成詞元的解碼階段。單個解碼步驟通常只需數十毫秒,但預填所需的時間則長得多。如果在相同的裝置上運行,將預填與解碼混合會降低解碼效能。在本文中,我們探討了一種既有的解決方案,即分離式預填與解碼,將它們在不同的裝置上運行,以將預填傳輸量和解碼延遲都發揮到極致。
預填與解碼效能
在典型的 LLM 服務引擎中,批次排程器會在模型的每個執行步驟中選取要處理的請求。當在單個裝置或節點上運行時,預填和解碼請求會被批次處理在一起。沿著序列長度進行聚合的注意力成本,無論是對於預填還是解碼,都會隨著 KV 快取中項目的長度 (kv_len) 成比例增長。解碼請求通常會轉發單個詞元 (qo_len=1),透過獨立運作於序列詞元上的其他層,其成本極低。預填請求則會透過密集層轉發數千或數萬個詞元 (qo_len 很大),成本相當高。
前向傳播的延遲更受通過稠密層的獨立 token 數量(qo_len)影響,而非注意力期間從 KV cache 取回的 token 數量(kv_len)。注意力可同時在請求數量與與序列長度成比例的 kv_len 上並行化,從而達到良好的資源利用率。Prefill 屬於計算受限:當 qo_len 很高時,GEMM kernel 可以沿 M 維度分配足夠的區塊,充分利用現代 GPU 的計算能力。Decode 屬於記憶體受限:由於批次大小通常較低,沿 M 的輸入數量通常很少,只足以對應一個區塊。雖然 Split-K GEMM kernel 可在較低 token 批次大小下改善 SM 利用率,但快取與矩陣乘法單元通常仍未被充分利用。
當混合在一起時,包含預填請求的批次會在前向傳遞中產生較高的延遲,從而對整個實例的解碼傳輸量產生負面影響。雖然將預填請求與解碼請求混合或採用分段預填可以稍微改善解碼效能,但很難維持足夠的預填傳輸量來在實例上處理足夠的請求,以將解碼傳輸量最大化。對於具有典型輸出長度的大型模型而言,為了維持解碼的大批次大小,必須經常執行預填,這會顯著降低平均延遲並導致輸出出現停頓。
這些問題可以透過使用一組獨立的節點來執行預填和解碼來解決。透過將一個預填節點與多個解碼節點相關聯,可以排程足夠的請求進行預填以最大化傳輸量,並在解碼節點上維持足夠數量的並行請求,以同樣將解碼傳輸量最大化。預填節點會填入 KV 快取,然後將其傳輸至解碼節點。由於解碼器不再需要為了預填而中斷,延遲會變得更具決定性,因為活躍請求不斷增長的 kv_len 的整體影響會顯著減小。這需要付出首次詞元時間 (TTFT) 增加的代價,因為透過網路傳輸 KV 快取可能需要數十到數百毫秒。
KV Messenger
在 Perplexity,我們針對分離式預填與解碼的實作是圍繞著 KV messenger 建立的,它與 LLM 引擎互動,透過網路協調從預填節點到解碼節點的 KV 快取傳輸。在預填端,messenger 接受來自解碼節點的請求,將其交給批次排程器,並追蹤前向傳遞的執行,以盡可能以最低的延遲分派 KV 快取。在解碼端,配置好不可逐出的頁面後,messenger 會阻擋請求被排程進行解碼,直到收到 KV 快取和解碼器上下文傳輸完成的通知為止。
分離預填需要高傳輸量、低延遲的連線,因此我們的實作是針對 RDMA 量身打造的,同時支援 EFA 和 ConnectX 網路介面卡 (NIC)。KV Messenger 建立在 libfabric 之上,使用我們的 fabric-lib 包裝函式在遠端直接記憶體存取 (RDMA) 原語之上提供更高階的低延遲抽象,實作高效的頁面和後設資料傳輸,以及低延遲信號傳遞。在背景中,fabric-lib 會協調 GPU 及其直接連接的 NIC,以將資料從預填節點複製到解碼節點。
收到後,預填節點會配置一組對應的來源 KV 頁面,並使用其本機引擎將請求排程進行預填。為了將延遲降到最低,傳輸不會等待前向傳遞:相反地,一旦模型完成將 KV 快取項目附加到個別層的 KV 快取後,就會立即啟動 KV 頁面複製。由於預填請求可以分段,因此批次排程器會在執行前將目前排程的分段通知 KV messenger。為了在能夠追蹤層的同時支援 CUDA 圖形,messenger 會保留一個專用執行緒來輪詢在注意力的輸出投影之後遞增的計數器。該計數器僅在分片環境中的主導節點上維護:儘管 KV 快取項目在附加後和注意力之前是有效的,但輸出投影會在各個排名之間進行縮減,從而隱式同步它們。一旦觀測到計數器的變更,就會通知 messenger,並呼叫 fabric-lib 來啟動層的傳輸。
最後一個分段的傳輸完成後,任何額外的後設資料也會被複製過去:推測式解碼或 MTP 需要將 logits 和隱藏狀態移動到解碼器。這些複製也透過 RDMA 執行,往返於預先配置的緩衝區。
最後一個分段的所有暫停傳輸完成後,預填節點會取消配置 KV 頁面並完成請求。解碼器節點不會收到明確的通知;相反地,它會使用即時計數器來追蹤已完成的操作數量。預填端的 RDMA 操作數量與傳輸的頁面數量成比例。在完成已知數量的頁面和上下文複製後,fabric-lib 會呼叫 KV messenger 以指示請求已準備好進行解碼。messenger 會取消配置任何上下文,並將請求交給 LLM 引擎。
分片式 KV 快取傳輸
如果預填器和解碼器依賴張量平行 (TP) 並以相同方式對 KV 快取進行分片或複製,則單個傳輸引擎會協調多個裝置來傳送和接收所有複本的頁面。為了能夠在模型執行器跨多個裝置和處理程序複製的情況下使用單個 messenger 和傳輸引擎,會使用 cuMem 和 cuMemImportFromShareableHandle 來配置支援 KV 快取的裝置記憶體,並將其對應到主處理程序中。傳輸引擎會檢查節點的拓樸,以尋找最接近 NUMA 節點中的 NIC 和 CPU,以用於每個 KV 快取片段的傳輸。
如果來源和目的地的分片方式相同,則傳輸非常簡單,因為來源和目的地的裝置和頁面之間存在一對一的對應關係。在這種情況下,分片會隱式協助傳輸延遲:透過使用更多 GPU,可以部署更多相關的 NIC,從而更接近全頻寬利用率。但是,如果存在不匹配的情況,傳輸引擎必須根據來源和目的地片段之間的比例來分割或重建頁面。
如果預填器將 KV 快取分割到更多裝置上,則會透過從預填器裝置發送對應的半部,在解碼器上重建完整的頁面。如果解碼器有更多分片,它會從多個來源接收頁面。解碼器需要了解預填器的分片方案,以便能夠計算預期接收的 RDMA 寫入次數。如果涉及複製,預填器會將裝置分組為複本集,這些複本集會在自身內部複製完整的 KV快取。目的地複本集會被隨機指派其中一個來源集,以便使用所有可用的裝置來啟動 RDMA 寫入。
分片傳輸需要對 KV cache 作些微調整。預設情況下,FlashInfer 依賴 NHD 佈局,該佈局按 heads 內的 page 來排列 tokens。由於 cache 最有可能沿 attention heads 的數量進行分片,這會在 head 內造成不連續。RDMA 傳輸並不隱式支援跨距寫入,因此每個 head 都需要一個操作才能完成傳輸。相對地,為了減少與 libfabric 的互動次數,我們使用 HND 佈局來組織 KV cache,讓 head 維度位於 token 數量之前。這可確保連續性,讓一個 page 能以單次寫入複製過去。
推測式解碼
推測式解碼需要對分離式預填-解碼進行輕微調整。在我們的實作中,不允許預填節點對詞元進行取樣。由於 Perplexity 的 Sonar 模型支援結構化輸出,我們不希望承擔跨預填器和解碼器同步綱要處理程序實作的複雜性。在 MTP 和推測式解碼機制中,將草稿模型預填到最後一個詞元涉及從目標模型取樣詞元。
為了繞過這些問題,預填不包含輸入序列的最後一個詞元。相反地,會傳輸最後一個詞元之前的預填隱藏狀態或 logits,並在解碼器的下一步中將其視為解碼詞元。雖然這會稍微增加延遲(因為必須在預填後執行完整的解碼步驟來發出第一個詞元),但實作的複雜性大幅降低。
分離式部署
我們已經部署或試驗了具有不同模型的多種分離式組態,以支援生產流量或內部評估工作負載。根據模型的大小和注意力機制,我們為預填器和解碼器節點選擇了合適的分片方案,以充分利用 GPU。
DeepSeek-R1
對於 DeepSeek,我們同時考慮了張量平行 (TP) 和資料平行 (DP) 部署。如先前的部落格文章所述,TP 部署可提供更好的延遲,代價是傳輸量較低,需要更多 GPU 來服務龐大的流量。DP 部署的負載擴展性好得多,但由於裝置間或節點間通訊的成本,其峰值傳輸量較低。
DeepSeek 依賴多頭潛在注意力 (Multi-Head Latent Attention) 來壓縮 KV 快取。由於所有 KV 頭都被壓縮成一個潛在向量,因此 TP 無法對 KV 快取進行分片,因為它必須在所有排名上複製潛在向量。分片發生在解密之後,因為每個排名可以從相同的潛在表示中提取不同的頭。因此,所有 KV 快取分片在預填器和解碼器分片中都是相同的。
透過節點內 TP 設定,預填器和解碼器的分片方式相同。傳輸是從所有排名分派的,以充分利用所有可用的 NIC。然而,透過 DP 部署(其中 TP 排名大小較小,或每個 DP 排名被指派到單個 GPU),持有 KV 快取複本的任何預填器裝置都可以分派它。為了平衡所有可用 NIC 上的請求,我們隨機選擇一個 GPU 和一個 NIC,將 KV 快取從預填器發送到解碼器。
使用混合預填-解碼時,由於頻繁發生數百毫秒量級的預填中斷,我們的 R1 部署很難持續超過 50 TPS。相比之下,透過分離預填,每個請求的 TTFT 懲罰約為 100 毫秒,但單個預填節點可以在 3 個解碼器節點上維持一致的批次大小,在處理每個解碼器節點約 1 QPS 的負載時,提供超過 90 TPS 的傳輸量。在資料平行部署中,TPS 稍低,約為 50,但實例可以處理每個排名 1 QPS 的負載,單個節點有 8 個排名。
Qwen3-Coder
這個 480B 模型使用群組查詢注意力 (GQA),因此注意力可以輕易分片,並且可以受益於張量平行,而不會犧牲 KV 快取的記憶體。因此,我們可以將模型跨 8 個 GPU 進行分片以進行預填和解碼,將大約 3 個解碼器節點與單個預填節點配對。由於注意力是分片的,我們依賴 HND KV 快取配置來對預填器和解碼器 KV 快取進行分片,將預填器排名與解碼器排名配對,並充分利用所有 NIC 來平行傳輸片段。
引用
@online{perplexity2025DisaggregatedPrefillAndDecode,
author = {{Perplexity} and Nandor Licker and Lequn Chen and Kevin Hu and Erin McKnight and Phi Hoang},
title = {{Disaggregated prefill and decode}},
date = {2025-08-01},
url = {https://perplexity.coolkeypoint.com/hub/blog/disaggregated-prefill-and-decode}
}