分而治之:全面解析分布式分離 Inference 系統
一、背景
大模型,如大語言模型(LLM)和大型多模態模型(LMM),正在改變自然語言處理和多模態任務的格局。然而,這些模型的 Inference 過程面臨大計算、大內存、高時延等諸多挑戰。為了應對這些問題,分布式分離 Inference 系統應運而生,旨在通過將模型的不同部分分開處理來優化性能。
大體來說,大模型 Inference 經歷了從單體到分布式,再到分離式的演進,并在繼續發展中:
1.單體 Inference 階段(2020 年前):
- 模型完整加載至單個設備(CPU 或 GPU),Inference 過程在同一設備完成
- 適用于小型模型(如 BERT、ViT、早期 GPT-2)
- 局限性:受單設備內存和計算能力限制
- 分布式并行階段(2020 - 2023):
- 采用模型并行(Tensor Parallel)和流水線并行(Pipeline Parallel)
- 模型仍作為整體處理,但計算分布在多設備
- 分離式階段(2024 至今):
- 按特性將模型拆分為獨立組件
- 組件可獨立部署和優化
- 引入專門的調度和 Cache 管理
當前的大模型正在經歷探索突破智能上限的階段,受到硬件的制約,架構由簡到繁是迫不得已的選擇。隨著硬件的快速迭代以及模型和算法的不斷演進,也必將經歷由繁到簡的過程。本文中我們基于之前關注的一系列工作,從幾個方面分別介紹當前大模型 Inference 系統的幾個關鍵技術:
- 分離式 Inference:
模態分離:Vision/Language 單獨處理
PD 分離:Prefill 和 Decoding 分布式擴展
Attention/MoE 分離:提升資源利用
- 分布式并行:
TP 及優化(MultiShot,Flux 等)
PP/EP/SPP
- 分布式 KV Cache 存儲和管理:
阿里 DistKV,Transfer Engine,NIXL,EIC 等。
- 調度策略
二、分離式 Inference 系統
2.1 概述
- 模態分離:涉及將不同模態分開處理,例如在 Vision-Language 模型中,Vision Encoder 可以獨立優化。
- PD 分離:將 LLM 的 Prefill 和 Decoding 階段分開,允許在不同硬件上運行,從而提高吞吐量,降低成本。
- Attention & MoE 分離:在 Decoding 階段,通過將 Attention 計算和 FFN(MoE)模塊分開,優化 MoE 模型的資源利用率。
2.2 模態分離
2.2.1 多模態模型范式
- 左側(Decoder-Only):Image Token 和 Text Token 一起作為 Input Token 輸入 LLM,在模型結構上與常見的 LLM 沒有本質區別。這種方式也是當前最常見的。
優勢:可以充分利用預訓練好的 LLM 模型,適配的成本低,效果好。
劣勢:常見場景的 Image Token 比較多,并且圖像越多,分辨率越高也往往意味著越多的 Image Token,導致較大的 Inference 成本。
- 右側(Cross-Attention):在 Transformer Block 中添加 Cross-Attention 模塊,Image Token 只在 Cross-Attention 模塊實現與 Text Token 的信息交互。Meta 去年的 LLaMA 3 ([2407.21783] The Llama 3 Herd of Models [2])中采用了這種 Cross Attention 的方案,但是在最近的 LLaMA 4(The Llama 4 herd: The beginning of a new era of natively multimodal AI innovation [3])中放棄了這種方案,轉向上述的早期模態融合方案。
優勢:即使 Image Token 比較多,也不會明顯增加 Inference 的成本。
劣勢:訓練和適配的成本較高,并且相應的效果可能不如上述早期融合方案。
2.2.2 Vision 模態分離
- 洞察一:各階段性能異質,Vision Encoder 成為 TTFT 瓶頸。
- 洞察二:Vision 與文本處理可以并行,解耦有巨大加速潛力。
- 洞察三:架構不同(Decoder-Only 和 Cross-Attention),LLM 后端的 Prefill 延遲差異巨大。
- 洞察四:Batching 策略、并行度對每一階段的影響大不相同。
- 洞察五:單體架構限制了分階段獨立伸縮,導致大幅資源浪費。
- 洞察六:路由與排隊策略必須模態感知,否則極易形成 HoL 阻塞和尾部延遲。
- 洞察七:動態負載需按 Token 級吞吐速率而非簡單 QPS 衡量與擴縮容。
總體來說,Inference 系統面臨如下挑戰:
- 多階段 Inference 流水線:典型的 LMM 推理包含 Vision 預處理、Vision 編碼和 LLM 后端生成等多個環節,計算特征與資源消耗各不相同。
預處理:包括 Image/Video 下載(網絡密集型)和 Image/Video 解碼和預處理(比如 Crop、Resize 等,通常使用 CPU,是 CPU 密集型)。
Vision 編碼:將 Image 和 Video 幀編碼為 Vision Token,通常使用 Vision Encoder 模型,GPU 密集型。此外,不需要緩存中間狀態(比如 KV Cache),因此對 GPU 顯存的需求不大。
LLM 后端:分為 Prefill 和 Decoding 階段,計算特性稍有不同,Prefill 通常是 GPU 計算密集,而 Decoding 通常是 GPU 內存密集,并且往往需要分布式。
- 多模態輸入的異質性:不同請求間輸入內容(圖片數、文本長度)和流量模式變動極大,會帶來負載預測和資源管理的挑戰。
- 現有 Inference 服務系統的局限:主流 LMM 服務往往采用“單體”架構,即所有子模塊(Vision、Text 等)在同一進程(設備)內共同調度、伸縮,難以針對具體環節和突發流量實現精細資源分配,導致高延遲和資源浪費。
為此,ModServe 論文中作者提出了對應的 ModServe 架構,如下圖 Figure 13 所示:
- 將 LMM Inference 分模塊化解耦,Vision 和 Text 相關計算在不同實例資源上獨立調度和優化。
- 支持自適應自動伸縮、模型切分與 Batching 策略,動態匹配真實時變流量,降低成本并精準滿足延遲 SLO。
- 設計模態感知調度與路由,有效應對圖像流量突發、請求異構性,降低尾部延遲。
PS:盡管 ModServe 不是最早探討模態分離研究的,但其在系統層面的深入分析非常具有代表性。
2.3 PD 分離
2.3.1 早期研究
Google 早在 22 年 11 月([2211.05102] Efficiently Scaling Transformer Inference [7])就對 Prefill 和 Decoding 進行了深入和系統的分析:
- Prefill 階段:輸入序列(通常為用戶輸入上下文/提示)全部一次性送入模型,可以對所有輸入 Token 同時并行計算,也就是高度并行化,通常是 Compute Bound。
- Decoding 階段:模型一步步生成新 Token,每生成一個 Token,都需要依賴前面(包括Prefill 階段)已生成/輸入的所有 Token,因而必須嚴格序列化,也就是并行化很低,通常是 Memory Bound。
如下圖 Figure 1 和 Table 3 所示,Prefill 的 Cost 相比 Decoding 更低,并且 Prefill 在比較小的 Batch Size 就能達到比較高的 MFU,而 Decoding 階段往往需要非常大的 Batch Size 才能達到比較大的 MFU。如紅框所示,Prefill 在 Batch Size 16-32 左右的 MFU 和 Decoding 階段 Batch Size 512-1024 的 MFU 相當。
不過上述工作中并沒有真的將 Prfill 和 Dcoding 進行分離部署。直到 23 年 11 月,微軟和華盛頓大學提出 Splitwise([2311.18677] Splitwise: Efficient generative LLM inference using phase splitting [8]),其為 LLM Inference 不同階段選擇不同的 GPU 類型。Prefill 階段為 Compute 密集型,選擇高算力 GPU,而 Decoding 階段為 Memory 密集型,使用算力不是特別強而訪存帶寬比較大的 GPU。同時為了兩個階段 KV Cache 的共享,需要在 GPU 間有高速的 IB 網絡互聯。
如下圖 Figure 10 所示為 Splitwise 的架構概覽。Splitwise 會維護兩個獨立節點池:Prompt Pool 和 Token Pool,用于 Prefill 和 Decoding 的處理。第三個混合池(Mixed Pool)根據工作負載的情況來擴展 Prompt Pool 和 Token Pool。
- 當新的推理請求到達時,調度進程會將其分配給一對節點(Prompt 和 Token)。Prompt 節點還會將 KV Cache 發送給 Token 節點。在 Token 節點上使用 Continuous Batching,以最大限度地提高其利用率。
- 為了滿足 SLO 并避免在較高負載下由于碎片而導致性能突降,Splitwise 維護了一個特殊的 Mixed Pool,該 Pool 根據輸入和輸出 Token 的請求速率和分布動態增長和收縮。Mixed Pool 中的節點仍然保留其作為 Prompt 節點或 Token 節點的標記,并在其掛起隊列中沒有相反類型的任務時返回其原始 Pool。
Splitwise 使用分層的兩級調度,Cluster 級調度(CLS)進程 1 負責將輸入請求路由到特定節點并重新調整節點的用戶。Machine 級調度(MLS)進程 2 維護掛起的隊列并管理每個節點上的批量請求。在較低的請求速率下,目標是在 Splitwise 中實現更好的延遲,而在更高的請求速率下,目標是避免由于 Prompt 節點和 Token 節點池之間的碎片而導致的任何性能或吞吐量降低。
2.3.2 MoonCake
24 年中期,Moonshot AI 團隊與清華大學共同提出 MoonCake([2407.00079] Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving [9]),它是 Kimi(Moonshot 旗下 AI應用)使用的核心 LLM Inference 系統。該系統以 KV Cache 為中心,采用解耦(Disaggregated)架構,專為支撐多樣化、海量、具有高時延需求的 Inference 流量而設計。由于其已經在生成環境大規模部署落地,在國內引起了大規模討論。
MoonCake 的核心創新點包括:
- KVCache 為中心的調度與存儲(KV Cache-centric):將模型 Inference 過程中的 KV Cache 單獨管理、分片、調度和分布式存儲,實現不同請求之間的最大化重復利用,減少重復計算。
- Prefill 與 Decoding 階段解耦:Inference 流程解耦為 Prefill 與 Decoding,分別由單獨的異構節點池負責,實現各自資源的最優調度。
- 分布式 KV Cache 池:充分利用 GPU 集群中未被完全利用的 CPU/DRAM/SSD 資源,通過 RDMA 實現節點間 KV Cache 高速轉移與副本復制,極大提升 Cache 命中率與帶寬利用。
- Cache 感知與超載導向調度算法:結合 KV Cache 分布和機器負載,實現“最大重用+最優負載”調度,針對高并發、高負載場景,還引入預測驅動的請求提前拒絕機制,保障 SLO 并減少無效資源消耗。
- 流水線分塊 Prefill 與層級 KV Cache 轉存:對長上下文輸入,采用分塊+流水線 Prefill(Chunked Pipeline Parallelism),降低 TTFT;同時 Prefill 過程中 KV Cache 分層異步加載/存儲與轉移,最大程度釋放寶貴 VRAM 空間。
如下圖 Figure 1 所示為其主要架構和流程:
- Conductor(全局調度器):根據 KV Cache 分布與預估請求負載,智能選擇 Prefill 節點與 Decoding 節點。
- KV Cache 復用:輸入經過 Token 分塊,優先復用已有 KV Cache 塊,需遠程拉取時平衡吞吐與時延。
- Prefill 分塊流水處理:對長輸入流按 Chunk 分發多節點串行處理,計算與 KV Cache 傳輸可并行化,顯著降低延遲。
- KV Cache 傳輸與存儲:KV Cache Pool 子模塊通過 RDMA 在 CPU 內存/GPU 顯存/SSD 間實現 KV Cache 高效傳遞與副本冗余,支持冷熱數據分級與 Cache 淘汰。
- Decoding 與批處理:Prefill 結束后,KV Cache 轉到預選 Decoding 節點,加入 Continuous Batching,高效生成后續 Token。
2.3.3 DeepSeek
24 年底,DeepSeek V3([2412.19437] DeepSeek-V3 Technical Report [10])發布,其中的 PD 分離部署及專家并行(Expert Parallelism)優化進一步引起大家對分離式部署方案的討論。其在 H800 集群上部署 DeepSeek-V3,為了同時確保在線服務的 SLO 和高吞吐量,采用 Prefill 階段和 Decoding 階段分離部署。
Prefill 階段:最小部署單元由 4 個節點組成,共 32 個 H800 GPU。
- Attention 部分采用 4 TP 結合 SP(Sequence Parallelism),并與 8 DP(Data Parallelism)相結合。其較小的 TP 規模(4)限制了 TP(Tensor Parallelism)通信的開銷。
- MoE 部分采用 32 EP,確保每個專家處理足夠大的 Batch Size,從而提升計算效率。
為了實現 MoE 部分不同 Expert 間的負載均衡,需確保每個 GPU 處理大致相同數量的 Token。引入了冗余 Expert 部署策略,即對高負載 Expert 進行復制并冗余部署。高負載 Expert 基于在線部署期間收集的統計數據檢測并定期調整(如每 10 分鐘一次)。確定冗余 Expert 集合后,根據觀察到的負載情況,在節點內的 GPU 間精心重新安排 Expert,力求在不增加跨節點 All2All 通信開銷的前提下,盡可能平衡各 GPU 的負載。對于 DeepSeek-V3 的部署,作者為 Prefill 階段設置了 32 個冗余 Expert。每個 GPU 除了原本承載的 8 個 Expert 外(256/32),還將額外承載一個冗余 Expert。
此外,在 Prefill 階段,為了提高吞吐量并隱藏 All2All 和 TP 通信的開銷,同時處理兩個計算工作量相近的 Micro Batch,將一個 Micro Batch 的 Attention 和 MoE 計算與另一個 Micro Batch的 Dispatching 和 Combining Overlap 執行。
Decoding 階段:將共享 Expert 視為路由 Expert 之一。由此視角出發,每個 Token 在路由時會選擇 9 個 Expert,其中共享 Expert 被視為高負載 Expert,始終被選中。Decoding 階段的最小部署單元由 40 個節點組成,共 320 H800 GPU。
- Attention 部分采用 4 TP 結合 SP,并與 80 DP 協同工作,而 MoE 部分則采用 320 EP。
- MoE 部分,每個 GPU 僅承載一位 Expert,其中 64 個 GPU 負責承載冗余 Expert 及共享 Expert。Dispatching 與 Combining 部分的 All2All 通信通過 IB Direct P2P 傳輸實現,以實現低時延。此外,還利用 IBGDA 技術進一步降低時延,提升通信效率(PS:DeepSeek 開源了對應的代碼庫 DeepEP: an efficient expert-parallel communication library [11])。
與 Prefill 階段類似,基于在線服務的 Expert 負載統計數據,在特定間隔內定期確定冗余專家集合。然而,由于每個 GPU 僅承載一位 Expert,因此無需重新安排 Expert。
2.3.4 NVIDIA Dynamo
Dynamo(Dynamo Inference Framework | NVIDIA Developer [12]) 是 NVIDIA 新開源的 Inference 服務框架(GitHub - ai-dynamo/dynamo: A Datacenter Scale Distributed Inference Serving Framework [13]),專為生成式 AI 和大模型的大規模部署而設計。Dynamo 同樣支持 Prefill 和 Decoding 的分離式部署,可以無縫兼容常見的 Inference Engine,比如 TensorRT-LLM、vLLM 和 SGLang 等。其核心模塊采用 Rust 編寫以獲得高性能,同時提供 Python 接口以便靈活控制。
如下圖 Figure 1 所示為 Dynamo 的架構示意圖,其中包含多個組件,每個組件都可以獨立伸縮:
- 包含 4 層:
API Server:接收用戶請求,兼容 OpenAI API 等主流接口,也會提供相應的 Metric 信息。
Smart Router:根據實際指標(比如每個 Worker 的 KV Cache 命中率和負載)對請求進行調度,盡可能平衡 KV Cache 復用和負載均衡。
Disaggregated Serving:真正的執行引擎,分為 Prefill Worker 和 Decoding Worker(PD 分離,也許后續還會有 Vision Worker)。
NVIDIA Inference Transfer Engine:也就是 NIXL(NVIDIA Inference Xfer Library),這是一個專門為 Inference 場景設計的高性能異步數據傳輸引擎。Prefill Worker 和 Decoding Worker 之間的 KV Cache 傳輸需要通過 NIXL。
- 其他組件:
- Planner 和 Event Plane 組件負責動態資源調度和監控。系統通過一個事件總線收集實時指標(例如請求隊列長度、Worker 利用率等),供 Planner 參考。當檢測到負載變化時,Planner 可實時增加或減少 Prefill/Decoding Worker 的數量,實現彈性擴展。
- KV Cache Manager:Dynamo 還提供了一個全局 KV Cache Manager,負責跨節點維護和調度所有 KV Cache。KV Cache Manager 記錄全局的 Cache 命中信息,并支持多層次內存卸載:當 GPU 顯存不足時,將冷數據卸載到 CPU 內存、SSD 或遠程存儲等更便宜的存儲介質上。這種分級 Cache 策略使得系統可以支持比單機顯存容量更長的上下文長度,同時保持較低的時延。
2.3.5 其他
PD 分離的推理系統還很多,這里就不一一介紹,其思路基本類似,可以參考:
- [2406.17565] MemServe: Context Caching for Disaggregated LLM Serving with Elastic Memory Pool [14]
- [2406.01566] Helix: Serving Large Language Models over Heterogeneous GPUs and Network via Max-Flow [15]
2.4 Attention & MoE 分離
MoE 模型在擴展 LLM 能力方面展現出巨大潛力,既能提升性能又可降低計算復雜度。然而,其稀疏激活架構使得 Inference 過程中 FFN 從 Compute Bound 轉變為 Memory Bound,導致 GPU 利用率顯著下降并增加運營成本。
字節跳動提出 MegaScale-Infer([2504.02263] MegaScale-Infer: Serving Mixture-of-Experts at Scale with Disaggregated Expert Parallelism [16]):一個高效且經濟的 MoE 模型服務系統。該系統在 LLM 的 Decoding 階段,對各模型層中的 Attention 模塊與 Expert 模塊解耦,實現獨立擴展、定制化并行策略及異構部署。
除此之外,MegaScale-Infer 創新性地采用 Ping-Pong 流水線并行技術,將請求 Batch 劃分為 Micro-Batch 并在 Attention 與 Expert 模塊間動態調度以完成 Inference。結合各模塊專屬的模型并行策略,該系統可以有效隱藏通信開銷并最大化 GPU 利用率。針對解耦后的 Attention 與 Expert 模塊特性,為最小化數據傳輸開銷(如 Token Dispatch),MegaScale-Infer 實現了高性能 M2N 通信庫,消除了不必要的 GPU-CPU 數據拷貝、Group 初始化開銷及 GPU 同步等待開銷。
如下圖 Figure 3 所示為 MegaScale-Infer 的架構,作者依然會采用常見的 PD 分離方案,也就是 Prefill 和 Decoding 有單獨的集群。本文的重點是 Decoding 階段,進一步對 Decoding 集群進行切分,包含 Attention 節點和 Expert 節點:
- Attention 節點:
每個節點包含全部 Attention 參數,采用 TP(Tensor Parallelism)。
- Expert 節點:
每個節點多個 GPU 存儲一個 Expert,采用 TP。
所有專家節點一起組成一個 EP(Exert Parallelism)組。
- TP:節點內 TP 可以充分利用高帶寬的 NVLink 通信。
為了解決資源浪費問題,作者引入了 Ping-Pong 流水線并行方案,與分布式訓練中的流水線并行思路完全類似。如下圖 Figure 4 所示,類似 Interleaved 1F1B,不過只有兩個設備,共 2*L 個 Stage(L 表示層數,每層都是 2 個 Stage);此外,Stage 間的通信方式也稍有不同。
采用分離式架構后,Attention 和 Expert 是相互獨立的節點,并且數量可能不同。假設 M 個 Attention 節點,N 個 Exert 節點,則 Dispatch(A2E) 變成 M2N 操作,而 Combine(E2A)變成 N2M 操作。為此,作者也開發了高效的 M2N 通信庫。
三、分布式并行 Inference 系統
3.1 概述
在單體 Inference 系統中,由于模型較大,使用單一設備進行 Inference 往往會導致 Latency 過高、性能不加,甚至出現無法部署的情況。而分布式并行技術可以有效的緩解上述問題,通過合理的并行策略,可以充分利用分布式計算資源,提高系統的吞吐量和響應速度。
雖然分離式 Inference 系統可以在一定程度上緩解了上述問題,但依然無法完全避免,這也是為什么在上述的分離式系統中還往往穿插著各種分布式并行技術。
3.2 Tensor Parallelism
3.2.1 原理
TP 是大模型 Inference 中最常使用的分布式并行策略。其將模型權重矩陣沿特定維度切分,這樣 Tensor 計算可以分散到多個設備上并行執行,這也是其能降低時延的最主要原因。如下圖所示為 MLP 和 Self-Attention 中 TP 的常見切分方案:
使用 TP 也會帶來一定的通信開銷,在 TP 的特定階段通常需要引入額外的 AllReduce 通信。對于標準的 Transformer Layer 而言,如果只是采用 TP 分布式 Inference,則每一層會引入 2 次 AllReduce 通信,并且常見的實現中這些 AllReduce 通信并不能被很好的 Overlap。也有一些列工作在優化這個問題。
3.2.2 NVIDIA MultiShot
NVIDIA 在 TensorRT-LLM 中利用 NVSwitch 的 MultiCast 能力對 AllReduce 進行優化(3x Faster AllReduce with NVSwitch and TensorRT-LLM MultiShot | NVIDIA Technical Blog [17]),可以有效降低 AllReduce 操作的時延(降低時延和增加吞吐是非常不一樣的場景,NCCL 中的 AllReduce 更關注吞吐,而 LLM Inference 中的 AllReduce 更希望降低時延)。
TensorRT-LLM 中的 MultiShot 實際上是真正的將 AllReduce 分成 ReduceScatter + AllGather,并且都進行了相應的優化。(PS:NVIDIA 的 Blog 中沒有詳細介紹這個優化,下面是我們的推測)
如下圖所示為 ReduceScatter 的優化:
如下圖所示為 AllGather 的優化:
3.2.3 美團 Flash Communication
美團在 [2412.04964] Flash Communication: Reducing Tensor Parallelization Bottleneck for Fast Large Language Model Inference [18] 中提出一種兩步量化策略,以替代傳統的 Ring AllReduce 方法,稱為 Flash AllReduce。
如下圖 Figure 6 展示了 Flash Communication 的通信原理:
- 首先,將每個 GPU 上的激活值按 Rank 的數量進行劃分。
- 在激活值上進行細粒度量化后,執行 All2All 通信(與我們猜測的 TRT-LLM 的 MultiShot 實現類似),使得每個設備接收其規約所需的計算負載。當然,接收后也需要反量化操作。
- 在設備內完成 Reduce,不涉及通信操作。
- 對得到的結果再次進行量化以加速傳輸。然后進行 AllGather 以匯總所有結果,并在每個設備上進行反量化以恢復浮點數值。
當然,量化可能影響精度,這也是該方案的最大問題。
3.2.4 字節跳動 Flux
字節跳動在 [2406.06858] FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion [19] 中提出 Flux,其將通信和計算操作分解為更細粒度的操作,并進一步融合成更大的 Kernel,從而在不損害 Kernel 效率的前提下有效隱藏通信。在融合 Kernel 的情況下,Flux 有望重疊高達 96% 的通信時間。
在 Flux 中,同樣是將 ReduceScatter 與 GEMM 進行 Overlap 和 Kernel 融合。ReduceScatter 操作可以分解為一個 AlltoAll 操作和一個 local 的 Reduce 操作。這里只有 AlltoAll 需要設備間通信,因此,將 All2All 融合到 GEMM 的尾部通常足以 Overlap 通信。與 ReduceScatter 不同,AllGather 的實現采用首部融合方式,直接嵌入 GEMM Kernel 中。具體而言,AllGather 的信號檢查功能被融合至 GEMM 內核的前序階段。
如下圖 Figure 5 展示了 ReduceScatter 中 Overlap 之間的主要差異:
- 現有 Overlap 方案 Tm 理論上可能比原始方法 Tc 執行得更快,但通常情況下,Tm 仍慢于原始 GEMM 操作時間 Tg。主要原因在于,將一個 GEMM Kernel 拆分為一系列較小的 GEMM Kernel 會降低 GPU GEMM 的執行效率。GEMM 通常需要合理大小的矩陣才能充分利用 GPU 的計算能力。這些具有數據依賴性的小型 GEMM 操作序列進一步阻礙了 GEMM Kernel 通過 GPU 多路復用技術并行運行,因此,Tensor 并行度越高,GPU 上的 GEMM 效率越低。
- 相比之下,作者提出的技術不存在上述限制。其 Tf 能夠在極小開銷下實現與原始 GEMM 操作 Tg 相當的性能。其細粒度分解策略完美契合現代 GPU 設計特性,即通過上下文切換的 Warp 和數百個在 SM 間并發活躍的 Warp 來隱藏延遲,如下圖 Figure 5 底部所示。最終,作者的方法在不影響 GEMM 計算效率的前提下,僅在執行末尾引入少量通信開銷。
3.2.5 字節跳動 TileLink
在上述的 Flux 中需要手寫 Kernel,成本比較高。為此,字節跳動進一步提出了 TileLink([2503.20313] TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitives [20]),旨在高效編譯并生成計算-通信 Overlap 執行的 Kernel。TileLink 由前端(Frontend)和后端(Backend)構成:
- 在前端,系統通過以 Tile 為中心的原語將通信與計算的設計空間解耦并建立關聯。
- 在后端,將這些原語轉換為底層指令,整合通信與計算組件以實現 Overlap 執行。
我們在之前的文章中已經詳細介紹過,這里不再贅述。最近作者也開源了相應的代碼,具體可以參考:GitHub - ByteDance-Seed/Triton-distributed [21]。
3.3 Pipeline Parallelism
3.3.1 原理
PP 通過將模型按層切分到不同設備上,形成 Inference 流水線,是處理超大規模模型的另一種分布式并行策略。
然而,對于每個 Query/Batch 而言,PP 中的不同層仍是串行執行,在系統沒有明顯瓶頸時并不會降低時延,因而在 Online 場景用用的并不是特別多,反而由于切分方式簡單、通信開銷少等優勢在 Offline 場景廣泛使用。
3.3.2 TP + PP 混合——微軟 DeepSpeed Inference
微軟在 DeepSpeed Inference([2207.00032] DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale [22])中探討了各種分布式策略混合的部署方案。也明確探討了 TP 和 PP 的優劣:
- TP
優點:高帶寬利用率,相鄰設備間通信延遲低(例如,單機多卡通過 NVLink 通信)
局限:跨節點大規模擴展時,受限于帶寬瓶頸和 AllReduce 通信;無法單獨用來突破單節點內存限制。
- PP
優點:適合多節點大規模擴展,通信只發生在相鄰 Pipeline Stage 間,通訊量小,易于橫向擴展;可以突破單節點內存上限。
局限:推理時存在 Pipeline Bubble(氣泡),尤其在自回歸生成場景,每步都要同步。自回歸依賴導致不可完全并行,Pipeline 難以100% 飽和;需要對 Batch size/Micro-batch 數量做精細調度才可最大化利用率,否則延遲較高。
3.3.3 TP + PP 混合——北大 FastServe
北大在 FastServe([2305.05920] Fast Distributed Inference Serving for Large Language Models [23])中也探索了 TP + PP 進行混合并行 Inference 的方案。其沒有通過實驗對 TP 和 PP 的性能進行詳細的對比,不過也確實提到了各自的優劣,也重點說明當不得不進行多機多卡分布式 Inference 時,考慮到跨節點通信成本很高,可以采用 TP + PP 的混合方案。
3.4 Expert Parallelism
3.4.1 原理
MoE 模型顯著提升了參數規模和表達能力,其關鍵做法是將 FFN 層由多個 “Expert(FFN)” 的組合,每次只激活其中少數 K 個 Expert,“稀疏” 地利用更大的參數空間。然而,MoE 模型參數極大,單一 DP 或 TP 已無法高效擴展和利用資源。
MoE 結構進一步進一步放大了 Decoding階段的 Memory Bound 問題。這主要是每次只激活個別專家,假設激活專家與總專家的比例是 1/16:
- 對于 Dense 模型,可能在 Batch Size 為 156(A100:312T FLOPS / 2TB/s = 156)即可以達到 Compute Bound。
- 對于 MoE 模型,需要 Batch Size 為 156*16=2496 才能達到 Compute Bound,并且由于負載不均,可能會更加嚴重。?
EP 就是為了解決上述問題,其把每個 Expert 分別放在不同 GPU 上。每個 GPU 只負責一部分 Expert。這樣一個 Batch 的 Token 只被路由(分配)到需要激活的 Expert,從而減少顯存需求,同時每個設備都有更大 Batch Size 的可能,幫助提升 GPU 利用率。
3.4.2 DeepSpeed Inference
微軟在 DeepSpeed Inference([2207.00032] DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale [24])中創新性地將 EP 與 TP 和 DP 結合,通過新的通信和 Kernel 優化,實現稀疏大模型的極致擴展。
3.5 Sequence Pipeline Parallelism
3.5.1 背景
隨著 LLM 應用拓展,如論文、書籍總結、長電影分析等, Inference 時對上下文長度的需求已提升到百萬級甚至千萬級 Token。然而,現有的大部分 Inference 系統在面對這種場景時存在顯著瓶頸:
- 自注意力計算復雜度為 O(n^2),計算和內存需求都飆升。
- 現有 Inference 系統(如 LoongServe)容易出現頭阻塞(Head-of-Line, HOL):短請求被長請求嚴重阻塞。
- 訓練時用到的一些長上下文并行技術(如 Ring Attention/Context Parallelism)在 Inference 階段并不適用,特別是面對變長請求和對時延 SLO 嚴格要求時。
為了解決上述問題,一系列工作將序列并行(Sequence Parallelism)的思路引入到了 Inference 中,并將其與 Pipeline Parallelism 的思路結合,提出了 Sequence Pipeline Parallelism(也稱作 Chunked Pipeline Parallelism)的方案。
3.5.2 MoonCake Chunked Pipeline Parallelism
我們在之前的 MoonCake PD 分離中提到,MoonCake 實現了 Chunked Pipeline Parallelism,也就是對長上下文輸入,采用分塊+流水線 Prefill,降低 TTFT。具體過程如下圖所示(C 表示 Chunk,L 表示 Layer,這里表示模型包含 6 個 Layer,Request 分為 4 個 Chunk):
- Request 切分:將長上下文 Request 的輸入 Token 按照預設的 Block 大小(如 1024)切分成多個 Chunk。
- 節點分組、流水線分配:將 Prefill 節點劃分為一組,采用流水線方式處理各 Chunk。假如分為 4 個節點,第一個節點處理第一個 chunk,處理結束傳至第二節點,以此類推,每個節點之間只有 Chunk 交接點存在通信。
- 并行執行、流水線推進:各節點可并行處理屬于自身的 Chunk,有效利用集群的并行能力并最大化吞吐。。
其中每個 Chunk 的前向計算需要依賴前一個 Chunk 同層的 KV Cache,因為自回歸模型生成某 Token 的輸出時,Attention 的 “key/value” 都必須包含前序所有 Token 的上下文。所以第 i 個 Chunk 進入第 L 層時,必須等第 i?1 個 Chunk 在同一層 L 的 KV Cache 計算完成。換言之,流水線是以 Chunk 維度并行,但跨 Chunk 內有 “層級同步” 依賴,如上圖中的箭頭所示。雖然需要 “層級同步”,但是其成本要比標準序列并行中的全局同步成本要低得多,并且可以和相鄰層的計算充分 Overlap。
當然 MoonCake 不是首先提出這種方案的,早在 21 年 2 月的 [2102.07988] TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models [25] 中就提到過類似思路。
3.5.3 MoonCake Sequence Pipeline Parallelism
Microsoft 在 [2409.17264] Medha: Efficiently Serving Multi-Million Context Length LLM Inference Requests Without Approximations [26] 中也提出了類似的方案。具體來說,Medha 系統提出三大核心創新,打造了 3D 并行 Inference 架構,實現了千萬級 Token Inference 的高效率和低時延。
- 自適應分塊(Adaptive Chunking)+ 松散調度(Slack-aware Scheduling)
核心思想:把大的 Prefill 請求拆成 Chunk,使用 Slack-aware 優先級調度算法排隊處理,能實現在長請求處理中靈活 “穿插” 短請求,不再出現長請求堵死管道。
Slack-aware(緊迫度優先):優先處理 “剩余時間比最少” 的請求,保障 SLO。
- 序列流水并行(Sequence Pipeline Parallelism, SPP)
與傳統 PP 相比:Medha 發現長請求 Prefill 各 Chunk 間實際上是“無依賴”的(不像 Decoding 有自回歸依賴),因此可以對 Prefill 各 Chunk 做 Pipeline “密集排布”,讓每個 Chunk 在流程中一階段結束后,馬上用新 Chunk 補位,因而 Prefill 總時延幾乎能隨 Pipeline 深度線下縮短,顯著降低首 Token 時延 (TTFT)。
效能顯著:在 128 張 H100 上 Prefill 百萬 Token,TTFT 降至 15s 以內。
- KV Cache 并行(KV Cache Parallelism, KVP)
核心思想:Decode 階段主要瓶頸在于大量 KV Cache 讀取帶寬。Medha 提出 KVP,把 KV Cache 沿序列切分,各 Server 持有部分 KV,每次 Decode 將 Q 片段廣播后,聚合 Partial Attention 結果,實現并行 Decoding,加速輸出 Token 速率(TPOT)。
優勢:KVP 通信開銷與 “Query Token 數” 而非 Cache 長度相關,適宜大上下文。
如下圖 Figure 12 是 Medha 對應的 3D 并行 Inference 系統:
PS:Meta 也在 [2411.01783] Context Parallelism for Scalable Million-Token Inference [27] 中提出了類似的方案,在 128 個 H100 GPU 實驗,LLaMA 3 405B 模型 1M Token 的 Inference 時間僅有 77s,達到 93% 的并行率以及 63% 的 FLOPs 利用率。
四、分布式 KV Cache
4.1 概述
現代大模型(如 GPT-4、Gemini、LongRoPE 等)支持越來越長的上下文窗口,序列長度可達幾十萬甚至上百萬 Token。而當前的主要 LLM 都是 Decoder Only 的 Transformer 模型,其自回歸特性需要 KV Cache 的存在,而 KV Cache 隨序列長度線性增加,單卡或單節點可能無法容納足夠的 KV Cache。
除此之外,即使 Request 結束,也可以保留熱點 KV Cache 用于 Prefix Caching 的加速,進一步加劇 KV Cache 存儲的壓力。比如,LLaMA-7B 模型,在 Context 長度為 1000K 時,其 KV Cache 可達數百 GB,遠超單 GPU 的顯存。 為了解決這個問題,分布式 KV Cache 應運而生。
4.2 阿里 DistKV-LLM
阿里在 [2401.02669] Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache [28] 中提出了 DistAttention,這是一種分布式注意力算法,可將 KV Cache 分割為更小的、可管理的單元,從而實現注意力模塊的分布式處理和存儲。基于此,作者也提出了 DistKV-LLM,它是一個分布式 LLM Serving 系統,可以動態管理 KV Cache,并有效地協調整個數據中心中所有可以訪問的 GPU 顯存和 CPU 內存,確保其可以適應各種上下文長度,提供高性能的 LLM 服務。
如下圖所示為 DistKV-LLM 的概覽,當 LLM 服務實例由于 KV 緩存增加而遇到內存不足時,DistKV-LLM 會主動識別并借用其他容量過剩的設備中的可用內存空間,這種自動化機制是通過兩個主要組件(rManger 和 gManager)的協作操作來實現的:
- gManager 充當中心化管理器,維護所有實例的全局內存信息,如下圖左圖所示。每個實例定期向 gManager 發送心跳信號,發送有關其剩余可用內存空間的更新信息。gManager 利用這些信息構建了一個詳細的表,稱作 Global Debt Ledger。
- rManger 提供統一的 API 接口,用于本地和遠程內存操作,如下圖右上所示。這些操作包括:
為新生成的 KV Cache 分配 Physical rBlock,并在不再需要時釋放。
在收到 local 或 remote 的內存分配請求時,rManager 會查詢 rBlock 表,以確定第一個可用的 Physical rBlock 空間。
在沒有足夠的空間時,rManager 會啟動從其他實例空間借用過程,如下圖右下所示。此時,如果分配請求來自 remote 實例,則 rManager 會返回錯誤響應,表示 remote 分配不可用,避免陷入死循環。
可以為分配給 remote 實例的 rBlock 數量設置上限,避免搶占 local 分配的空間。當然,該配置需要通過實驗確定,并配置為超參。
4.3 Moonshot AI Transfer-Engine
Moonshot AI 在 Github(GitHub - kvcache-ai/Mooncake: Mooncake is the serving platform for Kimi, a leading LLM service provided by Moonshot AI. [29])開源了 Transfer Engine 的具體實現,也在 Mooncake: Trading More Storage for Less Computation [30] 中有簡單的介紹。
Transfer Engine 是 MoonCake 項目的核心組件,支持高速數據傳輸,特別是在分離架構中,專注于 KV Cache 數據的復用、管理和傳輸,以提高可復用性和效率。旨在解決分布式系統中數據移動的瓶頸,提升系統整體性能。Transfer Engine 能處理多種介質,包括 Local DRAM、Local GPU VRAM、Remote DRAM、Remote GPU VRAM 和 Remote NVMe 等,滿足復雜分布式環境的需求。
Transfer Engine 圍繞兩個核心抽象設計:
- Segment:表示一個連續的地址空間,支持遠程讀寫,可分為 RAM Segment(非持久性存儲,如 DRAM 或 VRAM)和 NVMe-of Segment(通過 NVMe over Fabrics 的持久存儲)。
- BatchTransfer:封裝操作請求,負責在不同 Segment 的非連續數據空間之間同步數據,支持異步讀寫操作,類似于 AllScatter/AllGather,但更靈活。
Transfer Engine 實現了拓撲感知路徑選擇機制:各節點在處理請求前會生成拓撲矩陣并在集群廣播,該矩陣根據內存注冊時指定的類型,將網卡劃分為不同內存類型的“首選”與“備選”列表。正常工作中會優先選擇“首選”列表中的網卡進行傳輸。如下圖 Figure 4 所示,若需要將 Local 節點中 Buffer 0(分配到 CPU:0)的數據傳輸到目標節點 Buffer 1(分配到 CPU:1),Engine 首先根據 Local 節點的拓撲矩陣確定 CPU:0 的首選網卡,并選定 mlx5_1 作為 Local 網卡。同理,基于目標內存地址選定目標網卡。為了最大化帶寬利用率,單個請求的傳輸會被劃分為 16KB 的切片,各個切片可能采用不同的傳輸路徑,從而實現所有 RDMA 網卡的協同工作。
4.4 NVIDIA NIXL
我們在前面提到,NVIDIA 的 Dynamo 框架底層通過 NIXL(NVIDIA Inference Xfer Library) 實現 KV Cache 的傳輸。NVIDIA 也開源了相應實現:GitHub - ai-dynamo/nixl: NVIDIA Inference Xfer Library (NIXL) [31]。NIXL 是專為分布式 AI Inference 場景設計的高性能 P2P 通信庫,用于多節點/GPU Inference 框架中的數據傳輸。
如下圖所示,NIXL 以 Transfer Agent 為核心,每個節點上通常都會創建一個 Agent,負責管理該節點的 GPU、內存等存儲資源,并同其他 Agent 協調數據傳輸。Agent 具有全局唯一表示,并由上層 Inference 平臺分配。整個架構假定有一個上層協調進程負責 Inference 調度、內存分配、用戶請求等,并負責調用 NIXL 接口以及在 Agent 之間交換元數據。
NIXL 位于 Inference 框架和物理通信層之間,為 Inference 框架提供所需的數據平面接口,其設計包含 3 個關鍵抽象組件:
- Memory Section(內存區):對內存和存儲進行統一抽象。內存區由多個地址范圍(segment)組成,支持類型包括 CPU DRAM、GPU HBM、NVMe-of、對象存儲、文件等。每個內存區包含 Local 段信息和 Remote 訪問標識,方便不同 Agent 間的數據訪問。
- Transfer Backend Interface(傳輸后端接口):用于封裝不同通信后端(如 UCX、GPUDirect Storage,S3 等)的具體實現。每種后端在 Agent 初始化時注冊,Agent 會根據源/目標內存類型和雙方可用后端,自動選擇最優的傳輸路徑。
- Metadata Handler(元數據處理):管理 Agent 之間傳輸所需的元信息,包括每個后端的連接信息和各 Memory Section 的遠程標識符。
4.5 火山引擎-分布式 KV
前段時間火山引擎也發布了彈性極速緩存 EIC(Elastic Instant Cache),具體可以參考:推理加速新范式:火山引擎高性能分布式 KVCache (EIC)核心技術解讀 [32]。做的事情都類似,這里不再贅述。
五、調度策略
5.1 概述
在分離-分布式 Inference 系統中,調度策略對于實現高并發、低時延至關重要。在介紹上述的 Inference 系統時多多少少都有提到,這里不再贅述,只做簡單總結。
請求調度:系統需要根據請求的特點(Prompt 長度,可預測的生成長度、優先級等)將其調度到合理的節點上。比如,Decoding 階段通常采用 Continuous Batching 技術,將多個請求打包到一個 Micro-Batch 中處理,以緩解 Memory Bound 問題,提升 GPU 吞吐;然而過大的 Batch 又可能導致時延的增加,進而影響 SLO。此外,對于一些較短的請求,也可以適當提升優先級,減少尾部延遲,提高整體 SLO。最后,也需要充分考慮 Pipeline 的調度,盡量減少 Pipeline Bubble。
KV Cache 調度:在跨節點部署時,要合理管理 KV Cache 在各節點間的分布和一致性。通常使用數據親和性策略,將同一個會話或上下文路由到同一節點,這樣該節點保留了相關的KV Cache,無需頻繁跨網絡獲取;但是,也要避免節點過熱導致的負載不均衡問題,當節點繁忙或資源不足時,系統需要暫停向該節點轉發新請求,優先利用空閑節點;對于動態負載和故障,還要保證 KV Cache 的一致性與彈性。
資源調度:資源的彈性調度也是分離式 Inference 系統中老生常談的問題,比如系統中部署多少個 Prefill 節點,多少個 Decoding 節點,比例如何確定。輸入、輸出 Token 分布以及對應的模型、硬件等都會對節點的負載產生影響,因此也就需要提供詳細的 Metric 指標,進行動態的資源調整,以便在滿足 SLO 要求的情況下盡可能提升吞吐。
5.2 NVIDIA Dynamo KV Cache 路由
在大規模 Inference 系統中,KV Cache 對于減少重復計算、降低 Inference 成本至關重要。然而,KV Cache 的存在進一步增加了系統的復雜度,也使得路由機制變得更加具有挑戰性。
在 NVIDIA 的 Dynamo 框架中,KV Cache 路由機制是智能 Router 的核心,負責根據 Request 的 KV Cache 匹配度和 Worker 的負載情況,決定將 Request 發送到哪個 Worker。
Dynamo 中包括一些最基礎的路由機制:
- Random Routing:隨機調用,對應 client.generate() 和 client.random()。
- Round-Robin Routing:輪詢調用,對應 client.round_robin()。
- Direct Routing:指定特定的 Worker 調用,對應 client.direct(input, component_id)。KV Cache 路由需要使用這種方式。
將 Request 盡可能路由到已經存在 KV Cache 的 Worker,可以顯著加速生成過程,這也是 KV Cache 路由機制的目標。這使得負載均衡策略變得更加復雜,如果路由策略不感知每個 Worker 的 KV Cache,則可能出現:
- 如果路由到錯誤的 Worker,則錯過重用 KV Cache 的機會或者需要傳輸 KV Cache。
- 如果少量 Worker 接收過多請求,會導致負載不均衡,進而影響整體的吞吐。
解決這個問題的最好方式是能夠感知到全局視角的 KV Cache 和負載,以便 Router 能夠確定一個代價函數對各個 Worker 進行打分,挑選出最合適的 Worker。如下圖所示,將 KV Cache 匹配度 - 負載作為代價函數,Worker 1 得分(15%-30%=-15%),Worker 2 得分(50% - 50%=0),Worker 3 得分(75% - 80%=-5%),因此 Router 選擇 Worker 2。
在 NVIDIA Dynamo 中提供了 KvMetricsAggregator 來匯集各種指標,以便代價函數使用,比如:
- KV Cache 命中率
- 隊列中等待的 Request 數
- GPU Cache 使用率
- 負載百分比
六、參考鏈接
- ??https://arxiv.org/abs/2305.05665??
- ??https://arxiv.org/abs/2407.21783??
- ??https://ai.meta.com/blog/llama-4-multimodal-intelligence/??
- ??https://arxiv.org/abs/2406.11832??
- ??https://www.adept.ai/blog/fuyu-8b??
- ??https://arxiv.org/abs/2502.00937??
- ??https://arxiv.org/abs/2211.05102??
- ??https://arxiv.org/abs/2311.18677??
- ??https://arxiv.org/abs/2407.00079??
- ??https://arxiv.org/abs/2412.19437??
- ??https://github.com/deepseek-ai/DeepEP??
- ??https://developer.nvidia.com/dynamo??
- ??https://github.com/ai-dynamo/dynamo??
- ??https://arxiv.org/abs/2406.17565??
- ??https://arxiv.org/abs/2406.01566??
- ??https://arxiv.org/abs/2504.02263??
- ??https://developer.nvidia.com/blog/3x-faster-allreduce-with-nvswitch-and-tensorrt-llm-multishot/??
- ??https://arxiv.org/abs/2412.04964??
- ??https://arxiv.org/abs/2406.06858??
- ??https://arxiv.org/abs/2503.20313??
- ??https://github.com/ByteDance-Seed/Triton-distributed??
- ??https://arxiv.org/abs/2207.00032??
- ??https://arxiv.org/abs/2305.05920??
- ??https://arxiv.org/abs/2207.00032??
- ??https://arxiv.org/abs/2102.07988??
- ??https://arxiv.org/abs/2409.17264??
- ??https://www.arxiv.org/abs/2411.01783??
- ??https://arxiv.org/abs/2401.02669??
- ??https://github.com/kvcache-ai/Mooncake??
- ??https://www.usenix.org/system/files/fast25-qin.pdf??
- ??https://github.com/ai-dynamo/nixl??
- ??https://mp.weixin.qq.com/s/tasDqXf0Gxr3o_WCJ2IJUQ???
本文轉載自??????AI閑談??????,作者:AI閑談
