2月26日，Mooncake的論文獲得「計算機存儲頂會 FAST 2025」Best Paper，這也是國內連續第三年拿到FAST Best Paper。同時，Mooncake 團隊宣布和 vLLM 團隊已經合作制定了一個多階段路線圖。這次整合將為 vLLM 引入 P/D（Prefill/Decode）分解和全局 KVCache 設計。
Mooncake：
GitHub - kvcache-ai/Mooncake: Mooncake is the serving platform for Kimi, a leading LLM service provi?github.com/kvcache-ai/Mooncake
大模型訓練/推理 高性能優化
原始論文發表于 2024年，論文地址：Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving
分離式部署的動機是什么？
大模型推理分成?prefill?和?decode?兩個階段，prefill 階段是計算密集型，decode 階段是訪存密集型且受限于顯存，兩者的資源利用率不一致。decode 階段在顯存滿足的情況下應給與更多的 tokens，來提高資源利用率，分離式部署就可以很好的滿足這點需求。
通過分離式部署，prefill 和 decode 可以采取對各自有利的分布式配置、可以采取對各自有利的 GPU 架構，來利用有限的 GPU 資源最大化吞吐。
但是分離式部署，需要將 prefill 的?kv cache?傳輸到 decode，引入了通信開銷；常用的解決辦法有如下幾點：
- prefill 階段按 chunk 或者 layer 來傳遞，計算完一個就傳遞一個
- kv cache 量化，采樣 int 4 或者 int8，在精度達標的情況下，減少了通信量
分離式部署一般需要做什么？
- 針對 prefill instance，盡可能提高 kv cache 命中率，來減少重復計算
- prefill 按 chunk/layer 計算，【計算】和【傳遞 kv cache 給 decode 】流式并行
- decode 加載 kv cache，并且將請求添加到 continuous batching 進行處理
chunked pipeline parallelism
針對每個請求，都會被劃分為 chunks，同一請求的不同 chunk 在不同的節點上面進行計算，這樣多個請求就可以并行起來，從而減少 TTFT。
prefill 階段的優化目標
prefill 階段需要盡可能提高 kv cache 命中率，減少重復計算。這就會涉及到針對常用的 kv blocks 需要被拷貝到 prefill 節點，不常用的 kv blocks 需要釋放掉。在常見的分離式部署框架中，一般由 KV Router 模塊來管理。
decode 階段的優化目標
decode 階段需要集合盡可能多的 tokens 到一個 batch 來進行推理，從而提高 MFU（Model Flops Utilization），這就會涉及到顯存的問題，可能顯存會不夠用，同時還不能過多的影響 TBT（Time Between Tokens）。TBT 是用戶能比較直接感受到的一個指標，TBT 過高，用戶體驗會較差。
Mooncake 是為了解決什么？
Mooncake 主要是遇到了 GPU 資源緊張，在實際應用場景，有限的資源情況下，如果處理過載調度（overload-oriented scheduling）的問題。
Mooncake 提出了哪些方案？
Mooncake 提出了一種?early rejection?機制，通過預測將來的負載來提前拒絕某些請求，從而減少計算資源浪費。因為實際高峰場景中會遇到某個請求完成 prefill 計算后，decoding 階段并沒有資源槽來做，還不如提前結束反饋給客戶端，避免資源的進一步擁堵。
(1).?KVCache Pool
Mooncake 的 KVCache 緩存池是在 CPU 上的。上面的流程圖解釋的比較清晰，先看命中了哪些 prefix cache，然后 prefill instance 進行加載，再計算 incremental cache blocks，然后將 prefix cache blocks + incremental cache blocks 傳輸到 decode 節點上。
(2). Inference Workflow
上圖是一個完整推理的流程圖，可以發現 mooncake 使用了 prefill computation 和 kv cache load and store 并行來減少傳輸 overhead。由于 kv cache 主要以 CPU 作為中介存儲點，decode 階段也是計算和加載異步執行，邊加載邊計算。
kv cache 放在 CPU 上因為線上并發比較高，GPU 顯存有限，而 CPU 存儲空間大。實際場景 CPU 內存也是不夠用的，還需要用到硬盤，涉及到 kv cache offload 相關的進一步調度。
(3). KVCache-centric Scheduling
Mooncake 關于 KVCache 調度邏輯見上圖。
針對 prefill instance 的選擇，會考慮 prefix cache 命中長度和 KVCache blocks 重復使用的分布情況。新進來一個請求時，input tokens 會被切分成幾個 blocks，針對每個 block 計算一個 hash key，然后將這些 block keys 與每個 prefill instance cache keys 去匹配，來識別 prefix match length（prefix_len）。
調度器會根據每個請求的長度和 prefix_len 來估計 prefill 需要執行的時間（使用離線測試集制作的預測模型），然后加上預計需要的等待時間（隊列里面所有請求的 prefill 時長之和），就是這個請求的預估 TTFT。
由于大模型結構的一致性，prefill 時間還是比較好預估的，難預估的是 transfer 時間，因為 transfer 是異步的。
總結
Mooncake 這篇論文核心是為了解決 GPU 資源有限且請求文本較長的場景，提出的 early-rejection 方案可以參考一下。具體實際應用場景下的大模型推理分布式部署，還需要結合業務情況來具體優化，但是常見的優化手段，如 kv cache 量化、P 到 D 按 layer 傳輸、prefix cache 命中率提升這些都是必須的。另外包括異構場景下，P 和 D 不同的分布式配置，會引入什么問題，又會有哪些優化，論文并沒有涉及。
發布于 2025-05-12 21:46
