一、背景

大模型，如大語(yǔ)言模型（LLM）和大型多模態(tài)模型（LMM），正在改變自然語(yǔ)言處理和多模態(tài)任務(wù)的格局。然而，這些模型的 Inference 過(guò)程面臨大計(jì)算、大內(nèi)存、高時(shí)延等諸多挑戰(zhàn)。為了應(yīng)對(duì)這些問(wèn)題，分布式分離 Inference 系統(tǒng)應(yīng)運(yùn)而生，旨在通過(guò)將模型的不同部分分開(kāi)處理來(lái)優(yōu)化性能。

大體來(lái)說(shuō)，大模型 Inference 經(jīng)歷了從單體到分布式，再到分離式的演進(jìn)，并在繼續(xù)發(fā)展中：

1.單體 Inference 階段（2020 年前）：

• 模型完整加載至單個(gè)設(shè)備（CPU 或 GPU），Inference 過(guò)程在同一設(shè)備完成
• 適用于小型模型（如 BERT、ViT、早期 GPT-2）
• 局限性：受單設(shè)備內(nèi)存和計(jì)算能力限制

2. 分布式并行階段（2020 - 2023）：

• 采用模型并行（Tensor Parallel）和流水線并行（Pipeline Parallel）
• 模型仍作為整體處理，但計(jì)算分布在多設(shè)備

3. 分離式階段（2024 至今）：

• 按特性將模型拆分為獨(dú)立組件
• 組件可獨(dú)立部署和優(yōu)化
• 引入專門的調(diào)度和 Cache 管理

當(dāng)前的大模型正在經(jīng)歷探索突破智能上限的階段，受到硬件的制約，架構(gòu)由簡(jiǎn)到繁是迫不得已的選擇。隨著硬件的快速迭代以及模型和算法的不斷演進(jìn)，也必將經(jīng)歷由繁到簡(jiǎn)的過(guò)程。本文中我們基于之前關(guān)注的一系列工作，從幾個(gè)方面分別介紹當(dāng)前大模型 Inference 系統(tǒng)的幾個(gè)關(guān)鍵技術(shù)：

• 分離式 Inference：

模態(tài)分離：Vision/Language 單獨(dú)處理

PD 分離：Prefill 和 Decoding 分布式擴(kuò)展

Attention/MoE 分離：提升資源利用

• 分布式并行：

TP 及優(yōu)化（MultiShot，F(xiàn)lux 等）

PP/EP/SPP

• 分布式 KV Cache 存儲(chǔ)和管理：

阿里 DistKV，Transfer Engine，NIXL，EIC 等。

• 調(diào)度策略

二、分離式 Inference 系統(tǒng)

2.1 概述

• 模態(tài)分離：涉及將不同模態(tài)分開(kāi)處理，例如在 Vision-Language 模型中，Vision Encoder 可以獨(dú)立優(yōu)化。
• PD 分離：將 LLM 的 Prefill 和 Decoding 階段分開(kāi)，允許在不同硬件上運(yùn)行，從而提高吞吐量，降低成本。
• Attention & MoE 分離：在 Decoding 階段，通過(guò)將 Attention 計(jì)算和 FFN（MoE）模塊分開(kāi)，優(yōu)化 MoE 模型的資源利用率。

2.2 模態(tài)分離

2.2.1 多模態(tài)模型范式

• 左側(cè)（Decoder-Only）：Image Token 和 Text Token 一起作為 Input Token 輸入 LLM，在模型結(jié)構(gòu)上與常見(jiàn)的 LLM 沒(méi)有本質(zhì)區(qū)別。這種方式也是當(dāng)前最常見(jiàn)的。

優(yōu)勢(shì)：可以充分利用預(yù)訓(xùn)練好的 LLM 模型，適配的成本低，效果好。

劣勢(shì)：常見(jiàn)場(chǎng)景的 Image Token 比較多，并且圖像越多，分辨率越高也往往意味著越多的 Image Token，導(dǎo)致較大的 Inference 成本。

• 右側(cè)（Cross-Attention）：在 Transformer Block 中添加 Cross-Attention 模塊，Image Token 只在 Cross-Attention 模塊實(shí)現(xiàn)與 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]）中放棄了這種方案，轉(zhuǎn)向上述的早期模態(tài)融合方案。

優(yōu)勢(shì)：即使 Image Token 比較多，也不會(huì)明顯增加 Inference 的成本。

劣勢(shì)：訓(xùn)練和適配的成本較高，并且相應(yīng)的效果可能不如上述早期融合方案。

2.2.2 Vision 模態(tài)分離

• 洞察一：各階段性能異質(zhì)，Vision Encoder 成為 TTFT 瓶頸。
• 洞察二：Vision 與文本處理可以并行，解耦有巨大加速潛力。
• 洞察三：架構(gòu)不同（Decoder-Only 和 Cross-Attention），LLM 后端的 Prefill 延遲差異巨大。
• 洞察四：Batching 策略、并行度對(duì)每一階段的影響大不相同。
• 洞察五：?jiǎn)误w架構(gòu)限制了分階段獨(dú)立伸縮，導(dǎo)致大幅資源浪費(fèi)。
• 洞察六：路由與排隊(duì)策略必須模態(tài)感知，否則極易形成 HoL 阻塞和尾部延遲。
• 洞察七：動(dòng)態(tài)負(fù)載需按 Token 級(jí)吞吐速率而非簡(jiǎn)單 QPS 衡量與擴(kuò)縮容。

總體來(lái)說(shuō)，Inference 系統(tǒng)面臨如下挑戰(zhàn)：

• 多階段 Inference 流水線：典型的 LMM 推理包含 Vision 預(yù)處理、Vision 編碼和 LLM 后端生成等多個(gè)環(huán)節(jié)，計(jì)算特征與資源消耗各不相同。

預(yù)處理：包括 Image/Video 下載（網(wǎng)絡(luò)密集型）和 Image/Video 解碼和預(yù)處理（比如 Crop、Resize 等，通常使用 CPU，是 CPU 密集型）。

Vision 編碼：將 Image 和 Video 幀編碼為 Vision Token，通常使用 Vision Encoder 模型，GPU 密集型。此外，不需要緩存中間狀態(tài)（比如 KV Cache），因此對(duì) GPU 顯存的需求不大。

LLM 后端：分為 Prefill 和 Decoding 階段，計(jì)算特性稍有不同，Prefill 通常是 GPU 計(jì)算密集，而 Decoding 通常是 GPU 內(nèi)存密集，并且往往需要分布式。

• 多模態(tài)輸入的異質(zhì)性：不同請(qǐng)求間輸入內(nèi)容（圖片數(shù)、文本長(zhǎng)度）和流量模式變動(dòng)極大，會(huì)帶來(lái)負(fù)載預(yù)測(cè)和資源管理的挑戰(zhàn)。
• 現(xiàn)有 Inference 服務(wù)系統(tǒng)的局限：主流 LMM 服務(wù)往往采用“單體”架構(gòu)，即所有子模塊（Vision、Text 等）在同一進(jìn)程（設(shè)備）內(nèi)共同調(diào)度、伸縮，難以針對(duì)具體環(huán)節(jié)和突發(fā)流量實(shí)現(xiàn)精細(xì)資源分配，導(dǎo)致高延遲和資源浪費(fèi)。

為此，ModServe 論文中作者提出了對(duì)應(yīng)的 ModServe 架構(gòu)，如下圖 Figure 13 所示：

• 將 LMM Inference 分模塊化解耦，Vision 和 Text 相關(guān)計(jì)算在不同實(shí)例資源上獨(dú)立調(diào)度和優(yōu)化。
• 支持自適應(yīng)自動(dòng)伸縮、模型切分與 Batching 策略，動(dòng)態(tài)匹配真實(shí)時(shí)變流量，降低成本并精準(zhǔn)滿足延遲 SLO。
• 設(shè)計(jì)模態(tài)感知調(diào)度與路由，有效應(yīng)對(duì)圖像流量突發(fā)、請(qǐng)求異構(gòu)性，降低尾部延遲。

PS：盡管 ModServe 不是最早探討模態(tài)分離研究的，但其在系統(tǒng)層面的深入分析非常具有代表性。

2.3 PD 分離

2.3.1 早期研究

Google 早在 22 年 11 月（[2211.05102] Efficiently Scaling Transformer Inference [7]）就對(duì) Prefill 和 Decoding 進(jìn)行了深入和系統(tǒng)的分析：

• Prefill 階段：輸入序列（通常為用戶輸入上下文/提示）全部一次性送入模型，可以對(duì)所有輸入 Token 同時(shí)并行計(jì)算，也就是高度并行化，通常是 Compute Bound。
• Decoding 階段：模型一步步生成新 Token，每生成一個(gè) Token，都需要依賴前面（包括Prefill 階段）已生成/輸入的所有 Token，因而必須嚴(yán)格序列化，也就是并行化很低，通常是 Memory Bound。

如下圖 Figure 1 和 Table 3 所示，Prefill 的 Cost 相比 Decoding 更低，并且 Prefill 在比較小的 Batch Size 就能達(dá)到比較高的 MFU，而 Decoding 階段往往需要非常大的 Batch Size 才能達(dá)到比較大的 MFU。如紅框所示，Prefill 在 Batch Size 16-32 左右的 MFU 和 Decoding 階段 Batch Size 512-1024 的 MFU 相當(dāng)。

不過(guò)上述工作中并沒(méi)有真的將 Prfill 和 Dcoding 進(jìn)行分離部署。直到 23 年 11 月，微軟和華盛頓大學(xué)提出 Splitwise（[2311.18677] Splitwise: Efficient generative LLM inference using phase splitting [8]），其為 LLM Inference 不同階段選擇不同的 GPU 類型。Prefill 階段為 Compute 密集型，選擇高算力 GPU，而 Decoding 階段為 Memory 密集型，使用算力不是特別強(qiáng)而訪存帶寬比較大的 GPU。同時(shí)為了兩個(gè)階段 KV Cache 的共享，需要在 GPU 間有高速的 IB 網(wǎng)絡(luò)互聯(lián)。

如下圖 Figure 10 所示為 Splitwise 的架構(gòu)概覽。Splitwise 會(huì)維護(hù)兩個(gè)獨(dú)立節(jié)點(diǎn)池：Prompt Pool 和 Token Pool，用于 Prefill 和 Decoding 的處理。第三個(gè)混合池（Mixed Pool）根據(jù)工作負(fù)載的情況來(lái)擴(kuò)展 Prompt Pool 和 Token Pool。

• 當(dāng)新的推理請(qǐng)求到達(dá)時(shí)，調(diào)度進(jìn)程會(huì)將其分配給一對(duì)節(jié)點(diǎn)（Prompt 和 Token）。Prompt 節(jié)點(diǎn)還會(huì)將 KV Cache 發(fā)送給 Token 節(jié)點(diǎn)。在 Token 節(jié)點(diǎn)上使用 Continuous Batching，以最大限度地提高其利用率。
• 為了滿足 SLO 并避免在較高負(fù)載下由于碎片而導(dǎo)致性能突降，Splitwise 維護(hù)了一個(gè)特殊的 Mixed Pool，該 Pool 根據(jù)輸入和輸出 Token 的請(qǐng)求速率和分布動(dòng)態(tài)增長(zhǎng)和收縮。Mixed Pool 中的節(jié)點(diǎn)仍然保留其作為 Prompt 節(jié)點(diǎn)或 Token 節(jié)點(diǎn)的標(biāo)記，并在其掛起隊(duì)列中沒(méi)有相反類型的任務(wù)時(shí)返回其原始 Pool。

Splitwise 使用分層的兩級(jí)調(diào)度，Cluster 級(jí)調(diào)度（CLS）進(jìn)程 1 負(fù)責(zé)將輸入請(qǐng)求路由到特定節(jié)點(diǎn)并重新調(diào)整節(jié)點(diǎn)的用戶。Machine 級(jí)調(diào)度（MLS）進(jìn)程 2 維護(hù)掛起的隊(duì)列并管理每個(gè)節(jié)點(diǎn)上的批量請(qǐng)求。在較低的請(qǐng)求速率下，目標(biāo)是在 Splitwise 中實(shí)現(xiàn)更好的延遲，而在更高的請(qǐng)求速率下，目標(biāo)是避免由于 Prompt 節(jié)點(diǎn)和 Token 節(jié)點(diǎn)池之間的碎片而導(dǎo)致的任何性能或吞吐量降低。

2.3.2 MoonCake

24 年中期，Moonshot AI 團(tuán)隊(duì)與清華大學(xué)共同提出 MoonCake（[2407.00079] Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving [9]），它是 Kimi（Moonshot 旗下 AI應(yīng)用）使用的核心 LLM Inference 系統(tǒng)。該系統(tǒng)以 KV Cache 為中心，采用解耦（Disaggregated）架構(gòu)，專為支撐多樣化、海量、具有高時(shí)延需求的 Inference 流量而設(shè)計(jì)。由于其已經(jīng)在生成環(huán)境大規(guī)模部署落地，在國(guó)內(nèi)引起了大規(guī)模討論。

MoonCake 的核心創(chuàng)新點(diǎn)包括：

• KVCache 為中心的調(diào)度與存儲(chǔ)（KV Cache-centric）：將模型 Inference 過(guò)程中的 KV Cache 單獨(dú)管理、分片、調(diào)度和分布式存儲(chǔ)，實(shí)現(xiàn)不同請(qǐng)求之間的最大化重復(fù)利用，減少重復(fù)計(jì)算。
• Prefill 與 Decoding 階段解耦：Inference 流程解耦為 Prefill 與 Decoding，分別由單獨(dú)的異構(gòu)節(jié)點(diǎn)池負(fù)責(zé)，實(shí)現(xiàn)各自資源的最優(yōu)調(diào)度。
• 分布式 KV Cache 池：充分利用 GPU 集群中未被完全利用的 CPU/DRAM/SSD 資源，通過(guò) RDMA 實(shí)現(xiàn)節(jié)點(diǎn)間 KV Cache 高速轉(zhuǎn)移與副本復(fù)制，極大提升 Cache 命中率與帶寬利用。
• Cache 感知與超載導(dǎo)向調(diào)度算法：結(jié)合 KV Cache 分布和機(jī)器負(fù)載，實(shí)現(xiàn)“最大重用+最優(yōu)負(fù)載”調(diào)度，針對(duì)高并發(fā)、高負(fù)載場(chǎng)景，還引入預(yù)測(cè)驅(qū)動(dòng)的請(qǐng)求提前拒絕機(jī)制，保障 SLO 并減少無(wú)效資源消耗。
• 流水線分塊 Prefill 與層級(jí) KV Cache 轉(zhuǎn)存：對(duì)長(zhǎng)上下文輸入，采用分塊+流水線 Prefill（Chunked Pipeline Parallelism），降低 TTFT；同時(shí) Prefill 過(guò)程中 KV Cache 分層異步加載/存儲(chǔ)與轉(zhuǎn)移，最大程度釋放寶貴 VRAM 空間。

如下圖 Figure 1 所示為其主要架構(gòu)和流程：

• Conductor（全局調(diào)度器）：根據(jù) KV Cache 分布與預(yù)估請(qǐng)求負(fù)載，智能選擇 Prefill 節(jié)點(diǎn)與 Decoding 節(jié)點(diǎn)。
• KV Cache 復(fù)用：輸入經(jīng)過(guò) Token 分塊，優(yōu)先復(fù)用已有 KV Cache 塊，需遠(yuǎn)程拉取時(shí)平衡吞吐與時(shí)延。
• Prefill 分塊流水處理：對(duì)長(zhǎng)輸入流按 Chunk 分發(fā)多節(jié)點(diǎn)串行處理，計(jì)算與 KV Cache 傳輸可并行化，顯著降低延遲。
• KV Cache 傳輸與存儲(chǔ)：KV Cache Pool 子模塊通過(guò) RDMA 在 CPU 內(nèi)存/GPU 顯存/SSD 間實(shí)現(xiàn) KV Cache 高效傳遞與副本冗余，支持冷熱數(shù)據(jù)分級(jí)與 Cache 淘汰。
• Decoding 與批處理：Prefill 結(jié)束后，KV Cache 轉(zhuǎn)到預(yù)選 Decoding 節(jié)點(diǎn)，加入 Continuous Batching，高效生成后續(xù) Token。

2.3.3 DeepSeek

24 年底，DeepSeek V3（[2412.19437] DeepSeek-V3 Technical Report [10]）發(fā)布，其中的 PD 分離部署及專家并行（Expert Parallelism）優(yōu)化進(jìn)一步引起大家對(duì)分離式部署方案的討論。其在 H800 集群上部署 DeepSeek-V3，為了同時(shí)確保在線服務(wù)的 SLO 和高吞吐量，采用 Prefill 階段和 Decoding 階段分離部署。

Prefill 階段：最小部署單元由 4 個(gè)節(jié)點(diǎn)組成，共 32 個(gè) H800 GPU。

• Attention 部分采用 4 TP 結(jié)合 SP（Sequence Parallelism），并與 8 DP（Data Parallelism）相結(jié)合。其較小的 TP 規(guī)模（4）限制了 TP（Tensor Parallelism）通信的開(kāi)銷。
• MoE 部分采用 32 EP，確保每個(gè)專家處理足夠大的 Batch Size，從而提升計(jì)算效率。

為了實(shí)現(xiàn) MoE 部分不同 Expert 間的負(fù)載均衡，需確保每個(gè) GPU 處理大致相同數(shù)量的 Token。引入了冗余 Expert 部署策略，即對(duì)高負(fù)載 Expert 進(jìn)行復(fù)制并冗余部署。高負(fù)載 Expert 基于在線部署期間收集的統(tǒng)計(jì)數(shù)據(jù)檢測(cè)并定期調(diào)整（如每 10 分鐘一次）。確定冗余 Expert 集合后，根據(jù)觀察到的負(fù)載情況，在節(jié)點(diǎn)內(nèi)的 GPU 間精心重新安排 Expert，力求在不增加跨節(jié)點(diǎn) All2All 通信開(kāi)銷的前提下，盡可能平衡各 GPU 的負(fù)載。對(duì)于 DeepSeek-V3 的部署，作者為 Prefill 階段設(shè)置了 32 個(gè)冗余 Expert。每個(gè) GPU 除了原本承載的 8 個(gè) Expert 外（256/32），還將額外承載一個(gè)冗余 Expert。

此外，在 Prefill 階段，為了提高吞吐量并隱藏 All2All 和 TP 通信的開(kāi)銷，同時(shí)處理兩個(gè)計(jì)算工作量相近的 Micro Batch，將一個(gè) Micro Batch 的 Attention 和 MoE 計(jì)算與另一個(gè) Micro Batch的 Dispatching 和 Combining Overlap 執(zhí)行。

Decoding 階段：將共享 Expert 視為路由 Expert 之一。由此視角出發(fā)，每個(gè) Token 在路由時(shí)會(huì)選擇 9 個(gè) Expert，其中共享 Expert 被視為高負(fù)載 Expert，始終被選中。Decoding 階段的最小部署單元由 40 個(gè)節(jié)點(diǎn)組成，共 320 H800 GPU。

• Attention 部分采用 4 TP 結(jié)合 SP，并與 80 DP 協(xié)同工作，而 MoE 部分則采用 320 EP。
• MoE 部分，每個(gè) GPU 僅承載一位 Expert，其中 64 個(gè) GPU 負(fù)責(zé)承載冗余 Expert 及共享 Expert。Dispatching 與 Combining 部分的 All2All 通信通過(guò) IB Direct P2P 傳輸實(shí)現(xiàn)，以實(shí)現(xiàn)低時(shí)延。此外，還利用 IBGDA 技術(shù)進(jìn)一步降低時(shí)延，提升通信效率（PS：DeepSeek 開(kāi)源了對(duì)應(yīng)的代碼庫(kù) DeepEP: an efficient expert-parallel communication library [11]）。

與 Prefill 階段類似，基于在線服務(wù)的 Expert 負(fù)載統(tǒng)計(jì)數(shù)據(jù)，在特定間隔內(nèi)定期確定冗余專家集合。然而，由于每個(gè) GPU 僅承載一位 Expert，因此無(wú)需重新安排 Expert。

2.3.4 NVIDIA Dynamo

Dynamo（Dynamo Inference Framework | NVIDIA Developer [12]） 是 NVIDIA 新開(kāi)源的 Inference 服務(wù)框架（GitHub - ai-dynamo/dynamo: A Datacenter Scale Distributed Inference Serving Framework [13]），專為生成式 AI 和大模型的大規(guī)模部署而設(shè)計(jì)。Dynamo 同樣支持 Prefill 和 Decoding 的分離式部署，可以無(wú)縫兼容常見(jiàn)的 Inference Engine，比如 TensorRT-LLM、vLLM 和 SGLang 等。其核心模塊采用 Rust 編寫(xiě)以獲得高性能，同時(shí)提供 Python 接口以便靈活控制。

如下圖 Figure 1 所示為 Dynamo 的架構(gòu)示意圖，其中包含多個(gè)組件，每個(gè)組件都可以獨(dú)立伸縮：

• 包含 4 層：

API Server：接收用戶請(qǐng)求，兼容 OpenAI API 等主流接口，也會(huì)提供相應(yīng)的 Metric 信息。

Smart Router：根據(jù)實(shí)際指標(biāo)（比如每個(gè) Worker 的 KV Cache 命中率和負(fù)載）對(duì)請(qǐng)求進(jìn)行調(diào)度，盡可能平衡 KV Cache 復(fù)用和負(fù)載均衡。

Disaggregated Serving：真正的執(zhí)行引擎，分為 Prefill Worker 和 Decoding Worker（PD 分離，也許后續(xù)還會(huì)有 Vision Worker）。

NVIDIA Inference Transfer Engine：也就是 NIXL（NVIDIA Inference Xfer Library），這是一個(gè)專門為 Inference 場(chǎng)景設(shè)計(jì)的高性能異步數(shù)據(jù)傳輸引擎。Prefill Worker 和 Decoding Worker 之間的 KV Cache 傳輸需要通過(guò) NIXL。

• 其他組件：
• Planner 和 Event Plane 組件負(fù)責(zé)動(dòng)態(tài)資源調(diào)度和監(jiān)控。系統(tǒng)通過(guò)一個(gè)事件總線收集實(shí)時(shí)指標(biāo)（例如請(qǐng)求隊(duì)列長(zhǎng)度、Worker 利用率等），供 Planner 參考。當(dāng)檢測(cè)到負(fù)載變化時(shí)，Planner 可實(shí)時(shí)增加或減少 Prefill/Decoding Worker 的數(shù)量，實(shí)現(xiàn)彈性擴(kuò)展。
• KV Cache Manager：Dynamo 還提供了一個(gè)全局 KV Cache Manager，負(fù)責(zé)跨節(jié)點(diǎn)維護(hù)和調(diào)度所有 KV Cache。KV Cache Manager 記錄全局的 Cache 命中信息，并支持多層次內(nèi)存卸載：當(dāng) GPU 顯存不足時(shí)，將冷數(shù)據(jù)卸載到 CPU 內(nèi)存、SSD 或遠(yuǎn)程存儲(chǔ)等更便宜的存儲(chǔ)介質(zhì)上。這種分級(jí) Cache 策略使得系統(tǒng)可以支持比單機(jī)顯存容量更長(zhǎng)的上下文長(zhǎng)度，同時(shí)保持較低的時(shí)延。

2.3.5 其他

PD 分離的推理系統(tǒng)還很多，這里就不一一介紹，其思路基本類似，可以參考：

• [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 模型在擴(kuò)展 LLM 能力方面展現(xiàn)出巨大潛力，既能提升性能又可降低計(jì)算復(fù)雜度。然而，其稀疏激活架構(gòu)使得 Inference 過(guò)程中 FFN 從 Compute Bound 轉(zhuǎn)變?yōu)?Memory Bound，導(dǎo)致 GPU 利用率顯著下降并增加運(yùn)營(yíng)成本。

字節(jié)跳動(dòng)提出 MegaScale-Infer（[2504.02263] MegaScale-Infer: Serving Mixture-of-Experts at Scale with Disaggregated Expert Parallelism [16]）：一個(gè)高效且經(jīng)濟(jì)的 MoE 模型服務(wù)系統(tǒng)。該系統(tǒng)在 LLM 的 Decoding 階段，對(duì)各模型層中的 Attention 模塊與 Expert 模塊解耦，實(shí)現(xiàn)獨(dú)立擴(kuò)展、定制化并行策略及異構(gòu)部署。

除此之外，MegaScale-Infer 創(chuàng)新性地采用 Ping-Pong 流水線并行技術(shù)，將請(qǐng)求 Batch 劃分為 Micro-Batch 并在 Attention 與 Expert 模塊間動(dòng)態(tài)調(diào)度以完成 Inference。結(jié)合各模塊專屬的模型并行策略，該系統(tǒng)可以有效隱藏通信開(kāi)銷并最大化 GPU 利用率。針對(duì)解耦后的 Attention 與 Expert 模塊特性，為最小化數(shù)據(jù)傳輸開(kāi)銷（如 Token Dispatch），MegaScale-Infer 實(shí)現(xiàn)了高性能 M2N 通信庫(kù)，消除了不必要的 GPU-CPU 數(shù)據(jù)拷貝、Group 初始化開(kāi)銷及 GPU 同步等待開(kāi)銷。

如下圖 Figure 3 所示為 MegaScale-Infer 的架構(gòu)，作者依然會(huì)采用常見(jiàn)的 PD 分離方案，也就是 Prefill 和 Decoding 有單獨(dú)的集群。本文的重點(diǎn)是 Decoding 階段，進(jìn)一步對(duì) Decoding 集群進(jìn)行切分，包含 Attention 節(jié)點(diǎn)和 Expert 節(jié)點(diǎn)：

• Attention 節(jié)點(diǎn)：

每個(gè)節(jié)點(diǎn)包含全部 Attention 參數(shù)，采用 TP（Tensor Parallelism）。

• Expert 節(jié)點(diǎn)：

每個(gè)節(jié)點(diǎn)多個(gè) GPU 存儲(chǔ)一個(gè) Expert，采用 TP。

所有專家節(jié)點(diǎn)一起組成一個(gè) EP（Exert Parallelism）組。

• TP：節(jié)點(diǎn)內(nèi) TP 可以充分利用高帶寬的 NVLink 通信。

為了解決資源浪費(fèi)問(wèn)題，作者引入了 Ping-Pong 流水線并行方案，與分布式訓(xùn)練中的流水線并行思路完全類似。如下圖 Figure 4 所示，類似 Interleaved 1F1B，不過(guò)只有兩個(gè)設(shè)備，共 2*L 個(gè) Stage（L 表示層數(shù)，每層都是 2 個(gè) Stage）；此外，Stage 間的通信方式也稍有不同。

采用分離式架構(gòu)后，Attention 和 Expert 是相互獨(dú)立的節(jié)點(diǎn)，并且數(shù)量可能不同。假設(shè) M 個(gè) Attention 節(jié)點(diǎn)，N 個(gè) Exert 節(jié)點(diǎn)，則 Dispatch（A2E） 變成 M2N 操作，而 Combine（E2A）變成 N2M 操作。為此，作者也開(kāi)發(fā)了高效的 M2N 通信庫(kù)。

三、分布式并行 Inference 系統(tǒng)

3.1 概述

在單體 Inference 系統(tǒng)中，由于模型較大，使用單一設(shè)備進(jìn)行 Inference 往往會(huì)導(dǎo)致 Latency 過(guò)高、性能不加，甚至出現(xiàn)無(wú)法部署的情況。而分布式并行技術(shù)可以有效的緩解上述問(wèn)題，通過(guò)合理的并行策略，可以充分利用分布式計(jì)算資源，提高系統(tǒng)的吞吐量和響應(yīng)速度。

雖然分離式 Inference 系統(tǒng)可以在一定程度上緩解了上述問(wèn)題，但依然無(wú)法完全避免，這也是為什么在上述的分離式系統(tǒng)中還往往穿插著各種分布式并行技術(shù)。

3.2 Tensor Parallelism

3.2.1 原理

TP 是大模型 Inference 中最常使用的分布式并行策略。其將模型權(quán)重矩陣沿特定維度切分，這樣 Tensor 計(jì)算可以分散到多個(gè)設(shè)備上并行執(zhí)行，這也是其能降低時(shí)延的最主要原因。如下圖所示為 MLP 和 Self-Attention 中 TP 的常見(jiàn)切分方案：

使用 TP 也會(huì)帶來(lái)一定的通信開(kāi)銷，在 TP 的特定階段通常需要引入額外的 AllReduce 通信。對(duì)于標(biāo)準(zhǔn)的 Transformer Layer 而言，如果只是采用 TP 分布式 Inference，則每一層會(huì)引入 2 次 AllReduce 通信，并且常見(jiàn)的實(shí)現(xiàn)中這些 AllReduce 通信并不能被很好的 Overlap。也有一些列工作在優(yōu)化這個(gè)問(wèn)題。

3.2.2 NVIDIA MultiShot

NVIDIA 在 TensorRT-LLM 中利用 NVSwitch 的 MultiCast 能力對(duì) AllReduce 進(jìn)行優(yōu)化（3x Faster AllReduce with NVSwitch and TensorRT-LLM MultiShot | NVIDIA Technical Blog [17]），可以有效降低 AllReduce 操作的時(shí)延（降低時(shí)延和增加吞吐是非常不一樣的場(chǎng)景，NCCL 中的 AllReduce 更關(guān)注吞吐，而 LLM Inference 中的 AllReduce 更希望降低時(shí)延）。

TensorRT-LLM 中的 MultiShot 實(shí)際上是真正的將 AllReduce 分成 ReduceScatter + AllGather，并且都進(jìn)行了相應(yīng)的優(yōu)化。（PS：NVIDIA 的 Blog 中沒(méi)有詳細(xì)介紹這個(gè)優(yōu)化，下面是我們的推測(cè)）

如下圖所示為 ReduceScatter 的優(yōu)化：

如下圖所示為 AllGather 的優(yōu)化：

3.2.3 美團(tuán) Flash Communication

美團(tuán)在 [2412.04964] Flash Communication: Reducing Tensor Parallelization Bottleneck for Fast Large Language Model Inference [18] 中提出一種兩步量化策略，以替代傳統(tǒng)的 Ring AllReduce 方法，稱為 Flash AllReduce。

如下圖 Figure 6 展示了 Flash Communication 的通信原理：

• 首先，將每個(gè) GPU 上的激活值按 Rank 的數(shù)量進(jìn)行劃分。
• 在激活值上進(jìn)行細(xì)粒度量化后，執(zhí)行 All2All 通信（與我們猜測(cè)的 TRT-LLM 的 MultiShot 實(shí)現(xiàn)類似），使得每個(gè)設(shè)備接收其規(guī)約所需的計(jì)算負(fù)載。當(dāng)然，接收后也需要反量化操作。
• 在設(shè)備內(nèi)完成 Reduce，不涉及通信操作。
• 對(duì)得到的結(jié)果再次進(jìn)行量化以加速傳輸。然后進(jìn)行 AllGather 以匯總所有結(jié)果，并在每個(gè)設(shè)備上進(jìn)行反量化以恢復(fù)浮點(diǎn)數(shù)值。

當(dāng)然，量化可能影響精度，這也是該方案的最大問(wèn)題。

3.2.4 字節(jié)跳動(dòng) Flux

字節(jié)跳動(dòng)在 [2406.06858] FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion [19] 中提出 Flux，其將通信和計(jì)算操作分解為更細(xì)粒度的操作，并進(jìn)一步融合成更大的 Kernel，從而在不損害 Kernel 效率的前提下有效隱藏通信。在融合 Kernel 的情況下，F(xiàn)lux 有望重疊高達(dá) 96% 的通信時(shí)間。

在 Flux 中，同樣是將 ReduceScatter 與 GEMM 進(jìn)行 Overlap 和 Kernel 融合。ReduceScatter 操作可以分解為一個(gè) AlltoAll 操作和一個(gè) local 的 Reduce 操作。這里只有 AlltoAll 需要設(shè)備間通信，因此，將 All2All 融合到 GEMM 的尾部通常足以 Overlap 通信。與 ReduceScatter 不同，AllGather 的實(shí)現(xiàn)采用首部融合方式，直接嵌入 GEMM Kernel 中。具體而言，AllGather 的信號(hào)檢查功能被融合至 GEMM 內(nèi)核的前序階段。

如下圖 Figure 5 展示了 ReduceScatter 中 Overlap 之間的主要差異：

• 現(xiàn)有 Overlap 方案 Tm 理論上可能比原始方法 Tc 執(zhí)行得更快，但通常情況下，Tm 仍慢于原始 GEMM 操作時(shí)間 Tg。主要原因在于，將一個(gè) GEMM Kernel 拆分為一系列較小的 GEMM Kernel 會(huì)降低 GPU GEMM 的執(zhí)行效率。GEMM 通常需要合理大小的矩陣才能充分利用 GPU 的計(jì)算能力。這些具有數(shù)據(jù)依賴性的小型 GEMM 操作序列進(jìn)一步阻礙了 GEMM Kernel 通過(guò) GPU 多路復(fù)用技術(shù)并行運(yùn)行，因此，Tensor 并行度越高，GPU 上的 GEMM 效率越低。
• 相比之下，作者提出的技術(shù)不存在上述限制。其 Tf 能夠在極小開(kāi)銷下實(shí)現(xiàn)與原始 GEMM 操作 Tg 相當(dāng)?shù)男阅?。其?xì)粒度分解策略完美契合現(xiàn)代 GPU 設(shè)計(jì)特性，即通過(guò)上下文切換的 Warp 和數(shù)百個(gè)在 SM 間并發(fā)活躍的 Warp 來(lái)隱藏延遲，如下圖 Figure 5 底部所示。最終，作者的方法在不影響 GEMM 計(jì)算效率的前提下，僅在執(zhí)行末尾引入少量通信開(kāi)銷。

3.2.5 字節(jié)跳動(dòng) TileLink

在上述的 Flux 中需要手寫(xiě) Kernel，成本比較高。為此，字節(jié)跳動(dòng)進(jìn)一步提出了 TileLink（[2503.20313] TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitives [20]），旨在高效編譯并生成計(jì)算-通信 Overlap 執(zhí)行的 Kernel。TileLink 由前端（Frontend）和后端（Backend）構(gòu)成：

• 在前端，系統(tǒng)通過(guò)以 Tile 為中心的原語(yǔ)將通信與計(jì)算的設(shè)計(jì)空間解耦并建立關(guān)聯(lián)。
• 在后端，將這些原語(yǔ)轉(zhuǎn)換為底層指令，整合通信與計(jì)算組件以實(shí)現(xiàn) Overlap 執(zhí)行。

我們?cè)谥暗奈恼轮幸呀?jīng)詳細(xì)介紹過(guò)，這里不再贅述。最近作者也開(kāi)源了相應(yīng)的代碼，具體可以參考：GitHub - ByteDance-Seed/Triton-distributed [21]。

3.3 Pipeline Parallelism

3.3.1 原理

PP 通過(guò)將模型按層切分到不同設(shè)備上，形成 Inference 流水線，是處理超大規(guī)模模型的另一種分布式并行策略。

然而，對(duì)于每個(gè) Query/Batch 而言，PP 中的不同層仍是串行執(zhí)行，在系統(tǒng)沒(méi)有明顯瓶頸時(shí)并不會(huì)降低時(shí)延，因而在 Online 場(chǎng)景用用的并不是特別多，反而由于切分方式簡(jiǎn)單、通信開(kāi)銷少等優(yōu)勢(shì)在 Offline 場(chǎng)景廣泛使用。

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 的優(yōu)劣：

• TP

優(yōu)點(diǎn)：高帶寬利用率，相鄰設(shè)備間通信延遲低（例如，單機(jī)多卡通過(guò) NVLink 通信）

局限：跨節(jié)點(diǎn)大規(guī)模擴(kuò)展時(shí)，受限于帶寬瓶頸和 AllReduce 通信；無(wú)法單獨(dú)用來(lái)突破單節(jié)點(diǎn)內(nèi)存限制。

• PP

優(yōu)點(diǎn)：適合多節(jié)點(diǎn)大規(guī)模擴(kuò)展，通信只發(fā)生在相鄰 Pipeline Stage 間，通訊量小，易于橫向擴(kuò)展；可以突破單節(jié)點(diǎn)內(nèi)存上限。

局限：推理時(shí)存在 Pipeline Bubble（氣泡），尤其在自回歸生成場(chǎng)景，每步都要同步。自回歸依賴導(dǎo)致不可完全并行，Pipeline 難以100% 飽和；需要對(duì) Batch size/Micro-batch 數(shù)量做精細(xì)調(diào)度才可最大化利用率，否則延遲較高。

3.3.3 TP + PP 混合——北大 FastServe

北大在 FastServe（[2305.05920] Fast Distributed Inference Serving for Large Language Models [23]）中也探索了 TP + PP 進(jìn)行混合并行 Inference 的方案。其沒(méi)有通過(guò)實(shí)驗(yàn)對(duì) TP 和 PP 的性能進(jìn)行詳細(xì)的對(duì)比，不過(guò)也確實(shí)提到了各自的優(yōu)劣，也重點(diǎn)說(shuō)明當(dāng)不得不進(jìn)行多機(jī)多卡分布式 Inference 時(shí)，考慮到跨節(jié)點(diǎn)通信成本很高，可以采用 TP + PP 的混合方案。

3.4 Expert Parallelism

3.4.1 原理

MoE 模型顯著提升了參數(shù)規(guī)模和表達(dá)能力，其關(guān)鍵做法是將 FFN 層由多個(gè) “Expert（FFN）” 的組合，每次只激活其中少數(shù) K 個(gè) Expert，“稀疏” 地利用更大的參數(shù)空間。然而，MoE 模型參數(shù)極大，單一 DP 或 TP 已無(wú)法高效擴(kuò)展和利用資源。

MoE 結(jié)構(gòu)進(jìn)一步進(jìn)一步放大了 Decoding階段的 Memory Bound 問(wèn)題。這主要是每次只激活個(gè)別專家，假設(shè)激活專家與總專家的比例是 1/16：

• 對(duì)于 Dense 模型，可能在 Batch Size 為 156（A100：312T FLOPS / 2TB/s = 156）即可以達(dá)到 Compute Bound。
• 對(duì)于 MoE 模型，需要 Batch Size 為 156*16=2496  才能達(dá)到 Compute Bound，并且由于負(fù)載不均，可能會(huì)更加嚴(yán)重。?

EP 就是為了解決上述問(wèn)題，其把每個(gè) Expert 分別放在不同 GPU 上。每個(gè) GPU 只負(fù)責(zé)一部分 Expert。這樣一個(gè) Batch 的 Token 只被路由（分配）到需要激活的 Expert，從而減少顯存需求，同時(shí)每個(gè)設(shè)備都有更大 Batch Size 的可能，幫助提升 GPU 利用率。

3.4.2 DeepSpeed Inference

微軟在 DeepSpeed Inference（[2207.00032] DeepSpeed Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale [24]）中創(chuàng)新性地將 EP 與 TP 和 DP 結(jié)合，通過(guò)新的通信和 Kernel 優(yōu)化，實(shí)現(xiàn)稀疏大模型的極致擴(kuò)展。

3.5 Sequence Pipeline Parallelism

3.5.1 背景

隨著 LLM 應(yīng)用拓展，如論文、書(shū)籍總結(jié)、長(zhǎng)電影分析等， Inference 時(shí)對(duì)上下文長(zhǎng)度的需求已提升到百萬(wàn)級(jí)甚至千萬(wàn)級(jí) Token。然而，現(xiàn)有的大部分 Inference 系統(tǒng)在面對(duì)這種場(chǎng)景時(shí)存在顯著瓶頸：

• 自注意力計(jì)算復(fù)雜度為 O(n^2)，計(jì)算和內(nèi)存需求都飆升。
• 現(xiàn)有 Inference 系統(tǒng)（如 LoongServe）容易出現(xiàn)頭阻塞（Head-of-Line, HOL）：短請(qǐng)求被長(zhǎng)請(qǐng)求嚴(yán)重阻塞。
• 訓(xùn)練時(shí)用到的一些長(zhǎng)上下文并行技術(shù)（如 Ring Attention/Context Parallelism）在 Inference 階段并不適用，特別是面對(duì)變長(zhǎng)請(qǐng)求和對(duì)時(shí)延 SLO 嚴(yán)格要求時(shí)。

為了解決上述問(wèn)題，一系列工作將序列并行（Sequence Parallelism）的思路引入到了 Inference 中，并將其與 Pipeline Parallelism 的思路結(jié)合，提出了 Sequence Pipeline Parallelism（也稱作 Chunked Pipeline Parallelism）的方案。

3.5.2 MoonCake Chunked Pipeline Parallelism

我們?cè)谥暗?MoonCake PD 分離中提到，MoonCake 實(shí)現(xiàn)了 Chunked Pipeline Parallelism，也就是對(duì)長(zhǎng)上下文輸入，采用分塊+流水線 Prefill，降低 TTFT。具體過(guò)程如下圖所示（C 表示 Chunk，L 表示 Layer，這里表示模型包含 6 個(gè) Layer，Request 分為 4 個(gè) Chunk）：

• Request 切分：將長(zhǎng)上下文 Request 的輸入 Token 按照預(yù)設(shè)的 Block 大?。ㄈ?1024）切分成多個(gè) Chunk。
• 節(jié)點(diǎn)分組、流水線分配：將 Prefill 節(jié)點(diǎn)劃分為一組，采用流水線方式處理各 Chunk。假如分為 4 個(gè)節(jié)點(diǎn)，第一個(gè)節(jié)點(diǎn)處理第一個(gè) chunk，處理結(jié)束傳至第二節(jié)點(diǎn)，以此類推，每個(gè)節(jié)點(diǎn)之間只有 Chunk 交接點(diǎn)存在通信。
• 并行執(zhí)行、流水線推進(jìn)：各節(jié)點(diǎn)可并行處理屬于自身的 Chunk，有效利用集群的并行能力并最大化吞吐。。

其中每個(gè) Chunk 的前向計(jì)算需要依賴前一個(gè) Chunk 同層的 KV Cache，因?yàn)樽曰貧w模型生成某 Token 的輸出時(shí)，Attention 的 “key/value” 都必須包含前序所有 Token 的上下文。所以第 i 個(gè) Chunk 進(jìn)入第 L 層時(shí)，必須等第 i?1 個(gè) Chunk 在同一層 L 的 KV Cache 計(jì)算完成。換言之，流水線是以 Chunk 維度并行，但跨 Chunk 內(nèi)有 “層級(jí)同步” 依賴，如上圖中的箭頭所示。雖然需要 “層級(jí)同步”，但是其成本要比標(biāo)準(zhǔn)序列并行中的全局同步成本要低得多，并且可以和相鄰層的計(jì)算充分 Overlap。

當(dāng)然 MoonCake 不是首先提出這種方案的，早在 21 年 2 月的 [2102.07988] TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models [25] 中就提到過(guò)類似思路。

3.5.3 MoonCake Sequence Pipeline Parallelism

Microsoft 在 [2409.17264] Medha: Efficiently Serving Multi-Million Context Length LLM Inference Requests Without Approximations [26] 中也提出了類似的方案。具體來(lái)說(shuō)，Medha 系統(tǒng)提出三大核心創(chuàng)新，打造了 3D 并行 Inference 架構(gòu)，實(shí)現(xiàn)了千萬(wàn)級(jí) Token Inference 的高效率和低時(shí)延。

• 自適應(yīng)分塊（Adaptive Chunking）+ 松散調(diào)度（Slack-aware Scheduling）

核心思想：把大的 Prefill 請(qǐng)求拆成 Chunk，使用 Slack-aware 優(yōu)先級(jí)調(diào)度算法排隊(duì)處理，能實(shí)現(xiàn)在長(zhǎng)請(qǐng)求處理中靈活 “穿插” 短請(qǐng)求，不再出現(xiàn)長(zhǎng)請(qǐng)求堵死管道。

Slack-aware（緊迫度優(yōu)先）：優(yōu)先處理 “剩余時(shí)間比最少” 的請(qǐng)求，保障 SLO。

• 序列流水并行（Sequence Pipeline Parallelism, SPP）

與傳統(tǒng) PP 相比：Medha 發(fā)現(xiàn)長(zhǎng)請(qǐng)求 Prefill 各 Chunk 間實(shí)際上是“無(wú)依賴”的（不像 Decoding 有自回歸依賴），因此可以對(duì) Prefill 各 Chunk 做 Pipeline “密集排布”，讓每個(gè) Chunk 在流程中一階段結(jié)束后，馬上用新 Chunk 補(bǔ)位，因而 Prefill 總時(shí)延幾乎能隨 Pipeline 深度線下縮短，顯著降低首 Token 時(shí)延 (TTFT)。

效能顯著：在 128 張 H100 上 Prefill 百萬(wàn) Token，TTFT 降至 15s 以內(nèi)。

• KV Cache 并行（KV Cache Parallelism, KVP）

核心思想：Decode 階段主要瓶頸在于大量 KV Cache 讀取帶寬。Medha 提出 KVP，把 KV Cache 沿序列切分，各 Server 持有部分 KV，每次 Decode 將 Q 片段廣播后，聚合 Partial Attention 結(jié)果，實(shí)現(xiàn)并行 Decoding，加速輸出 Token 速率（TPOT）。

優(yōu)勢(shì)：KVP 通信開(kāi)銷與 “Query Token 數(shù)” 而非 Cache 長(zhǎng)度相關(guān)，適宜大上下文。

如下圖 Figure 12 是 Medha 對(duì)應(yīng)的 3D 并行 Inference 系統(tǒng)：

PS：Meta 也在 [2411.01783] Context Parallelism for Scalable Million-Token Inference [27] 中提出了類似的方案，在 128 個(gè) H100 GPU 實(shí)驗(yàn)，LLaMA 3 405B 模型 1M Token 的 Inference 時(shí)間僅有 77s，達(dá)到 93% 的并行率以及 63% 的 FLOPs 利用率。

四、分布式 KV Cache

4.1 概述

現(xiàn)代大模型（如 GPT-4、Gemini、LongRoPE 等）支持越來(lái)越長(zhǎng)的上下文窗口，序列長(zhǎng)度可達(dá)幾十萬(wàn)甚至上百萬(wàn) Token。而當(dāng)前的主要 LLM 都是 Decoder Only 的 Transformer 模型，其自回歸特性需要 KV Cache 的存在，而 KV Cache 隨序列長(zhǎng)度線性增加，單卡或單節(jié)點(diǎn)可能無(wú)法容納足夠的 KV Cache。

除此之外，即使 Request 結(jié)束，也可以保留熱點(diǎn) KV Cache 用于 Prefix Caching 的加速，進(jìn)一步加劇 KV Cache 存儲(chǔ)的壓力。比如，LLaMA-7B 模型，在 Context 長(zhǎng)度為 1000K 時(shí)，其 KV Cache 可達(dá)數(shù)百 GB，遠(yuǎn)超單 GPU 的顯存。 為了解決這個(gè)問(wèn)題，分布式 KV Cache 應(yīng)運(yùn)而生。

4.2 阿里 DistKV-LLM

阿里在 [2401.02669] Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache [28] 中提出了 DistAttention，這是一種分布式注意力算法，可將 KV Cache 分割為更小的、可管理的單元，從而實(shí)現(xiàn)注意力模塊的分布式處理和存儲(chǔ)。基于此，作者也提出了 DistKV-LLM，它是一個(gè)分布式 LLM Serving 系統(tǒng)，可以動(dòng)態(tài)管理 KV Cache，并有效地協(xié)調(diào)整個(gè)數(shù)據(jù)中心中所有可以訪問(wèn)的 GPU 顯存和 CPU 內(nèi)存，確保其可以適應(yīng)各種上下文長(zhǎng)度，提供高性能的 LLM 服務(wù)。

如下圖所示為 DistKV-LLM 的概覽，當(dāng) LLM 服務(wù)實(shí)例由于 KV 緩存增加而遇到內(nèi)存不足時(shí)，DistKV-LLM 會(huì)主動(dòng)識(shí)別并借用其他容量過(guò)剩的設(shè)備中的可用內(nèi)存空間，這種自動(dòng)化機(jī)制是通過(guò)兩個(gè)主要組件（rManger 和 gManager）的協(xié)作操作來(lái)實(shí)現(xiàn)的：

• gManager 充當(dāng)中心化管理器，維護(hù)所有實(shí)例的全局內(nèi)存信息，如下圖左圖所示。每個(gè)實(shí)例定期向 gManager 發(fā)送心跳信號(hào)，發(fā)送有關(guān)其剩余可用內(nèi)存空間的更新信息。gManager 利用這些信息構(gòu)建了一個(gè)詳細(xì)的表，稱作 Global Debt Ledger。
• rManger 提供統(tǒng)一的 API 接口，用于本地和遠(yuǎn)程內(nèi)存操作，如下圖右上所示。這些操作包括：

為新生成的 KV Cache 分配 Physical rBlock，并在不再需要時(shí)釋放。

在收到 local 或 remote 的內(nèi)存分配請(qǐng)求時(shí)，rManager 會(huì)查詢 rBlock 表，以確定第一個(gè)可用的 Physical rBlock 空間。

在沒(méi)有足夠的空間時(shí)，rManager 會(huì)啟動(dòng)從其他實(shí)例空間借用過(guò)程，如下圖右下所示。此時(shí)，如果分配請(qǐng)求來(lái)自 remote 實(shí)例，則 rManager 會(huì)返回錯(cuò)誤響應(yīng)，表示 remote 分配不可用，避免陷入死循環(huán)。

可以為分配給 remote 實(shí)例的 rBlock 數(shù)量設(shè)置上限，避免搶占 local 分配的空間。當(dāng)然，該配置需要通過(guò)實(shí)驗(yàn)確定，并配置為超參。

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]）開(kāi)源了 Transfer Engine 的具體實(shí)現(xiàn)，也在 Mooncake: Trading More Storage for Less Computation [30] 中有簡(jiǎn)單的介紹。

Transfer Engine 是 MoonCake 項(xiàng)目的核心組件，支持高速數(shù)據(jù)傳輸，特別是在分離架構(gòu)中，專注于 KV Cache 數(shù)據(jù)的復(fù)用、管理和傳輸，以提高可復(fù)用性和效率。旨在解決分布式系統(tǒng)中數(shù)據(jù)移動(dòng)的瓶頸，提升系統(tǒng)整體性能。Transfer Engine 能處理多種介質(zhì)，包括 Local DRAM、Local GPU VRAM、Remote DRAM、Remote GPU VRAM 和 Remote NVMe 等，滿足復(fù)雜分布式環(huán)境的需求。

Transfer Engine 圍繞兩個(gè)核心抽象設(shè)計(jì)：

• Segment：表示一個(gè)連續(xù)的地址空間，支持遠(yuǎn)程讀寫(xiě)，可分為 RAM Segment（非持久性存儲(chǔ)，如 DRAM 或 VRAM）和 NVMe-of Segment（通過(guò) NVMe over Fabrics 的持久存儲(chǔ)）。
• BatchTransfer：封裝操作請(qǐng)求，負(fù)責(zé)在不同 Segment 的非連續(xù)數(shù)據(jù)空間之間同步數(shù)據(jù)，支持異步讀寫(xiě)操作，類似于 AllScatter/AllGather，但更靈活。

Transfer Engine 實(shí)現(xiàn)了拓?fù)涓兄窂竭x擇機(jī)制：各節(jié)點(diǎn)在處理請(qǐng)求前會(huì)生成拓?fù)渚仃嚥⒃诩簭V播，該矩陣根據(jù)內(nèi)存注冊(cè)時(shí)指定的類型，將網(wǎng)卡劃分為不同內(nèi)存類型的“首選”與“備選”列表。正常工作中會(huì)優(yōu)先選擇“首選”列表中的網(wǎng)卡進(jìn)行傳輸。如下圖 Figure 4 所示，若需要將 Local 節(jié)點(diǎn)中 Buffer 0（分配到 CPU:0）的數(shù)據(jù)傳輸?shù)侥繕?biāo)節(jié)點(diǎn) Buffer 1（分配到 CPU:1），Engine 首先根據(jù) Local 節(jié)點(diǎn)的拓?fù)渚仃嚧_定 CPU:0 的首選網(wǎng)卡，并選定 mlx5_1 作為 Local 網(wǎng)卡。同理，基于目標(biāo)內(nèi)存地址選定目標(biāo)網(wǎng)卡。為了最大化帶寬利用率，單個(gè)請(qǐng)求的傳輸會(huì)被劃分為 16KB 的切片，各個(gè)切片可能采用不同的傳輸路徑，從而實(shí)現(xiàn)所有 RDMA 網(wǎng)卡的協(xié)同工作。

4.4 NVIDIA NIXL

我們?cè)谇懊嫣岬?，NVIDIA 的 Dynamo 框架底層通過(guò) NIXL（NVIDIA Inference Xfer Library） 實(shí)現(xiàn) KV Cache 的傳輸。NVIDIA 也開(kāi)源了相應(yīng)實(shí)現(xiàn)：GitHub - ai-dynamo/nixl: NVIDIA Inference Xfer Library (NIXL) [31]。NIXL 是專為分布式 AI Inference 場(chǎng)景設(shè)計(jì)的高性能 P2P 通信庫(kù)，用于多節(jié)點(diǎn)/GPU Inference 框架中的數(shù)據(jù)傳輸。 

如下圖所示，NIXL 以 Transfer Agent 為核心，每個(gè)節(jié)點(diǎn)上通常都會(huì)創(chuàng)建一個(gè) Agent，負(fù)責(zé)管理該節(jié)點(diǎn)的 GPU、內(nèi)存等存儲(chǔ)資源，并同其他 Agent 協(xié)調(diào)數(shù)據(jù)傳輸。Agent 具有全局唯一表示，并由上層 Inference 平臺(tái)分配。整個(gè)架構(gòu)假定有一個(gè)上層協(xié)調(diào)進(jìn)程負(fù)責(zé) Inference 調(diào)度、內(nèi)存分配、用戶請(qǐng)求等，并負(fù)責(zé)調(diào)用 NIXL 接口以及在 Agent 之間交換元數(shù)據(jù)。

NIXL 位于 Inference 框架和物理通信層之間，為 Inference 框架提供所需的數(shù)據(jù)平面接口，其設(shè)計(jì)包含 3 個(gè)關(guān)鍵抽象組件：

• Memory Section（內(nèi)存區(qū)）：對(duì)內(nèi)存和存儲(chǔ)進(jìn)行統(tǒng)一抽象。內(nèi)存區(qū)由多個(gè)地址范圍（segment）組成，支持類型包括 CPU DRAM、GPU HBM、NVMe-of、對(duì)象存儲(chǔ)、文件等。每個(gè)內(nèi)存區(qū)包含 Local 段信息和 Remote 訪問(wèn)標(biāo)識(shí)，方便不同 Agent 間的數(shù)據(jù)訪問(wèn)。
• Transfer Backend Interface（傳輸后端接口）：用于封裝不同通信后端（如 UCX、GPUDirect Storage，S3 等）的具體實(shí)現(xiàn)。每種后端在 Agent 初始化時(shí)注冊(cè)，Agent 會(huì)根據(jù)源/目標(biāo)內(nèi)存類型和雙方可用后端，自動(dòng)選擇最優(yōu)的傳輸路徑。
• Metadata Handler（元數(shù)據(jù)處理）：管理 Agent 之間傳輸所需的元信息，包括每個(gè)后端的連接信息和各 Memory Section 的遠(yuǎn)程標(biāo)識(shí)符。

4.5 火山引擎-分布式 KV

前段時(shí)間火山引擎也發(fā)布了彈性極速緩存 EIC（Elastic Instant Cache），具體可以參考：推理加速新范式：火山引擎高性能分布式 KVCache （EIC）核心技術(shù)解讀 [32]。做的事情都類似，這里不再贅述。

五、調(diào)度策略

5.1 概述

在分離-分布式 Inference 系統(tǒng)中，調(diào)度策略對(duì)于實(shí)現(xiàn)高并發(fā)、低時(shí)延至關(guān)重要。在介紹上述的 Inference 系統(tǒng)時(shí)多多少少都有提到，這里不再贅述，只做簡(jiǎn)單總結(jié)。

請(qǐng)求調(diào)度：系統(tǒng)需要根據(jù)請(qǐng)求的特點(diǎn)（Prompt 長(zhǎng)度，可預(yù)測(cè)的生成長(zhǎng)度、優(yōu)先級(jí)等）將其調(diào)度到合理的節(jié)點(diǎn)上。比如，Decoding 階段通常采用 Continuous Batching 技術(shù)，將多個(gè)請(qǐng)求打包到一個(gè) Micro-Batch 中處理，以緩解 Memory Bound 問(wèn)題，提升 GPU 吞吐；然而過(guò)大的 Batch 又可能導(dǎo)致時(shí)延的增加，進(jìn)而影響 SLO。此外，對(duì)于一些較短的請(qǐng)求，也可以適當(dāng)提升優(yōu)先級(jí)，減少尾部延遲，提高整體 SLO。最后，也需要充分考慮 Pipeline 的調(diào)度，盡量減少 Pipeline Bubble。

KV Cache 調(diào)度：在跨節(jié)點(diǎn)部署時(shí)，要合理管理 KV Cache 在各節(jié)點(diǎn)間的分布和一致性。通常使用數(shù)據(jù)親和性策略，將同一個(gè)會(huì)話或上下文路由到同一節(jié)點(diǎn)，這樣該節(jié)點(diǎn)保留了相關(guān)的KV Cache，無(wú)需頻繁跨網(wǎng)絡(luò)獲??；但是，也要避免節(jié)點(diǎn)過(guò)熱導(dǎo)致的負(fù)載不均衡問(wèn)題，當(dāng)節(jié)點(diǎn)繁忙或資源不足時(shí)，系統(tǒng)需要暫停向該節(jié)點(diǎn)轉(zhuǎn)發(fā)新請(qǐng)求，優(yōu)先利用空閑節(jié)點(diǎn)；對(duì)于動(dòng)態(tài)負(fù)載和故障，還要保證 KV Cache 的一致性與彈性。

資源調(diào)度：資源的彈性調(diào)度也是分離式 Inference 系統(tǒng)中老生常談的問(wèn)題，比如系統(tǒng)中部署多少個(gè) Prefill 節(jié)點(diǎn)，多少個(gè) Decoding 節(jié)點(diǎn)，比例如何確定。輸入、輸出 Token 分布以及對(duì)應(yīng)的模型、硬件等都會(huì)對(duì)節(jié)點(diǎn)的負(fù)載產(chǎn)生影響，因此也就需要提供詳細(xì)的 Metric 指標(biāo)，進(jìn)行動(dòng)態(tài)的資源調(diào)整，以便在滿足 SLO 要求的情況下盡可能提升吞吐。

5.2 NVIDIA Dynamo KV Cache 路由

在大規(guī)模 Inference 系統(tǒng)中，KV Cache 對(duì)于減少重復(fù)計(jì)算、降低 Inference 成本至關(guān)重要。然而，KV Cache 的存在進(jìn)一步增加了系統(tǒng)的復(fù)雜度，也使得路由機(jī)制變得更加具有挑戰(zhàn)性。

在 NVIDIA 的 Dynamo 框架中，KV Cache 路由機(jī)制是智能 Router 的核心，負(fù)責(zé)根據(jù) Request 的 KV Cache 匹配度和 Worker 的負(fù)載情況，決定將 Request 發(fā)送到哪個(gè) Worker。

Dynamo 中包括一些最基礎(chǔ)的路由機(jī)制：

• Random Routing：隨機(jī)調(diào)用，對(duì)應(yīng) client.generate() 和 client.random()。
• Round-Robin Routing：輪詢調(diào)用，對(duì)應(yīng) client.round_robin()。
• Direct Routing：指定特定的 Worker 調(diào)用，對(duì)應(yīng) client.direct(input, component_id)。KV Cache 路由需要使用這種方式。

將 Request 盡可能路由到已經(jīng)存在 KV Cache 的 Worker，可以顯著加速生成過(guò)程，這也是 KV Cache 路由機(jī)制的目標(biāo)。這使得負(fù)載均衡策略變得更加復(fù)雜，如果路由策略不感知每個(gè) Worker 的 KV Cache，則可能出現(xiàn)：

• 如果路由到錯(cuò)誤的 Worker，則錯(cuò)過(guò)重用 KV Cache 的機(jī)會(huì)或者需要傳輸 KV Cache。
• 如果少量 Worker 接收過(guò)多請(qǐng)求，會(huì)導(dǎo)致負(fù)載不均衡，進(jìn)而影響整體的吞吐。

解決這個(gè)問(wèn)題的最好方式是能夠感知到全局視角的 KV Cache 和負(fù)載，以便 Router 能夠確定一個(gè)代價(jià)函數(shù)對(duì)各個(gè) Worker 進(jìn)行打分，挑選出最合適的 Worker。如下圖所示，將 KV Cache 匹配度 - 負(fù)載作為代價(jià)函數(shù)，Worker 1 得分（15%-30%=-15%），Worker 2 得分（50% - 50%=0），Worker 3 得分（75% - 80%=-5%），因此 Router 選擇 Worker 2。

在 NVIDIA Dynamo 中提供了 KvMetricsAggregator 來(lái)匯集各種指標(biāo)，以便代價(jià)函數(shù)使用，比如：

• KV Cache 命中率
• 隊(duì)列中等待的 Request 數(shù)
• GPU Cache 使用率
• 負(fù)載百分比

六、參考鏈接

1. ??https://arxiv.org/abs/2305.05665??
2. ??https://arxiv.org/abs/2407.21783??
3. ??https://ai.meta.com/blog/llama-4-multimodal-intelligence/??
4. ??https://arxiv.org/abs/2406.11832??
5. ??https://www.adept.ai/blog/fuyu-8b??
6. ??https://arxiv.org/abs/2502.00937??
7. ??https://arxiv.org/abs/2211.05102??
8. ??https://arxiv.org/abs/2311.18677??
9. ??https://arxiv.org/abs/2407.00079??
10. ??https://arxiv.org/abs/2412.19437??
11. ??https://github.com/deepseek-ai/DeepEP??
12. ??https://developer.nvidia.com/dynamo??
13. ??https://github.com/ai-dynamo/dynamo??
14. ??https://arxiv.org/abs/2406.17565??
15. ??https://arxiv.org/abs/2406.01566??
16. ??https://arxiv.org/abs/2504.02263??
17. ??https://developer.nvidia.com/blog/3x-faster-allreduce-with-nvswitch-and-tensorrt-llm-multishot/??
18. ??https://arxiv.org/abs/2412.04964??
19. ??https://arxiv.org/abs/2406.06858??
20. ??https://arxiv.org/abs/2503.20313??
21. ??https://github.com/ByteDance-Seed/Triton-distributed??
22. ??https://arxiv.org/abs/2207.00032??
23. ??https://arxiv.org/abs/2305.05920??
24. ??https://arxiv.org/abs/2207.00032??
25. ??https://arxiv.org/abs/2102.07988??
26. ??https://arxiv.org/abs/2409.17264??
27. ??https://www.arxiv.org/abs/2411.01783??
28. ??https://arxiv.org/abs/2401.02669??
29. ??https://github.com/kvcache-ai/Mooncake??
30. ??https://www.usenix.org/system/files/fast25-qin.pdf??
31. ??https://github.com/ai-dynamo/nixl??
32. ??https://mp.weixin.qq.com/s/tasDqXf0Gxr3o_WCJ2IJUQ???

本文轉(zhuǎn)載自??????AI閑談??????，作者：AI閑談