月之暗面和清華KVCache.ai團隊的最新論文，首次揭秘了Kimi背后的推理架構！

要知道Kimi是國產大模型的當紅炸子雞，火到可以說從來沒缺過流量，甚至還經常出現過載。

而隨著論文的發布，這潑天的流量到底如何被Kimi接住的問題，也有了答案。

Kimi背后的推理架構名叫Mooncake（月餅），主要特點是采取了分離式的設計方案。

而且，Mooncake在設計之時就考慮了可能出現的大流量場景，并針對這種情況專門研發。

在模擬場景下，Mooncake最高能帶來525%的吞吐量增長，實際場景中也能多處理75%請求。

另據月之暗面工程副總裁許欣然的一篇知乎文章介紹，Kimi有80%以上的流量，都是由該系統承接。

從KV緩存出發，建造分布式系統

整個Mooncake系統設計的核心，是圍繞著KV緩存展開的。

（KV緩存用于存儲鍵-值對（Key-Value Pairs），主要優勢在于可以簡單高效地訪問和檢索數據，在大模型當中可以提高推理速度并減少計算資源消耗。）

之所以這樣做，是因為團隊預計KV緩存的容量會長期保持高位，因此圍繞KV緩存進行優化十分必要。

從結構上看，Mooncake由全局調度器（Conductor）、Prefill節點集群、Decoding節點集群和分布式KVCache池幾部分組成，另外還有RDMA通信組件（Messenger）。

其中全局調度器是用戶請求到達系統后的第一站，它負責接收請求并根據KV緩存分布和負載情況，將請求調度到Prefill和Decoding節點。

調度器在調度時需要綜合考慮KV緩存的復用長度、負載均衡等因素，實現KV緩存復用的最大化。

具體到Mooncake，它采用了一種啟發式的自動熱點遷移策略，可以在不需要精確預測未來訪問的情況下自動復制熱點KV緩存塊。

同時，這種動態復制熱點KV緩存塊的方式，也是實現均衡負載的一種重要途徑。

實驗結果表明，與隨機調度和負載均衡調度相比，Mooncake的調度策略可以顯著降低TTFT（Time To First Token，首個Token延遲），提高系統性能。

完成調度之后，任務會分別交由Prefill和Decoding節點進行運算。

Prefill節點接收到調度器轉發過來的請求后，會從KV緩存池中讀取緩存，執行預計算并生成新的KV緩存。

對于長上下文請求，Mooncake還會分塊流水并行的方式，使用多個節點并行處理來降低延遲。

而Decoding節點除了接收調度器發來的請求外，還會收到Prefill階段生成的KV緩存，節點會對這些緩存執行解碼并生成最終結果。

這當中，大容量、高性能的KV緩存存儲由緩存池提供；RDMA通信組件則憑借其高帶寬、低延遲的優勢，負責在不同節點之間的KV緩存傳輸。

除了采取以KV緩存為中心的工作流程外，Mooncake還有另一個重要特點——分離式的架構。

采取分離式架構的重要因素之一，是在于Prefill和Decoding兩個階段的計算特性差異很大。

具體來說，它們分別要對TTFT和TBT（Time Between Tokens，Token間延遲）負責。

這就導致了兩者在計算復雜度、內存訪問方式、并行粒度和對延遲的敏感度上都存在差異：

所以，月之暗面團隊對GPU集群也進行了相應的拆分，以便將它們分別部署在不同節點集群上，實現資源隔離和專門優化。

另外，Mooncake中的KV緩存池也是分布式的，同時充分利用了GPU集群中空閑的CPU、DRAM和SSD資源，實現了大容量、高帶寬的KV緩存存儲和傳輸，同時也減少了閑置資源的浪費。

提前預測負載，及時拒絕超量請求

不過，即使Mooncake采用了高效的分離架構，但實際環境中的超大流量，對系統仍然是一個考驗。

對此，作者也提出了新的應對策略。

在過載場景下，調度的關鍵是決定是否接受新的請求。

由于Mooncake采用的是分離式架構，可以采取早期拒絕策略，在Prefill階段就根據Decoding節點的負載情況，提前拒絕請求。

Mooncake使用TTFT和TBT的SLO（Service Level Objective，服務等級目標）滿足情況作為負載的度量指標。

具體的SLO要求是TTFT的90分位值（P90）不超過單個請求在空載條件下處理時間的10倍，TBT的P90值不超過5倍。

這種早期拒絕策略可以顯著減少無效的Prefill計算,提高資源利用率，但同時也帶來了新的問題——Prefill和Decoding節點負載的波動，導致資源利用率下降、影響系統性能。

這是由于早期拒絕策略中，系統做出請求拒絕的決策時存在滯后性，如下圖所示：

• 在階段1，Prefill節點和Decoding節點的負載都較低，此時調度器會持續接受新的請求，直到Prefill節點的負載達到上限。
• 進入階段2后，Rrefill節點處理的請求開始進入Decoding節點，導致其負載快速上升。當Decoding節點的負載超過閾值后調度器開始拒絕新的請求，但此時Prefill節點的負載仍然很高。
• 到了階段3，由于調度器拒絕新請求，Prefill節點的負載開始下降。但此前積壓的請求正在Decoding階段處理，節點的負載仍然很高。
• 最后是階段4，Decoding節點的負載開始下降，因為前面的請求都處理完成，而新的請求又被拒絕了。這時調度器再次開始接受新請求，Prefill節點的負載又開始上升。
• 之后，這個過程會周期性地重復，導致Prefill和Decoding節點的負載出現反相位的波動。

針對這一問題，月之暗面團隊對這種簡單的早期拒絕策略進行了修正，提出了基于預測的早期拒絕策略，從而降低節點負載的波動。

這種策略的核心思想是對一段時間后的Decoding節點負載進行預測，并基于預測結果決定是否拒絕請求。

預測可以在請求級別和系統級別兩個層面進行，請求級別的預測比較困難，因為要預測單個請求的執行時間；系統級別的預測相對容易一些，只需要預測整體的負載情況。

Mooncake采用的是一種簡化的系統級別預測方法，假設每個請求的執行時間服從某個固定分布，據此預測未來一段時間內的負載情況。

實驗結果表明，這種基于預測的早期拒絕策略，可以有效緩解負載波動問題。

最終，端到端性能評估結果表明，Mooncake的架構設計和優化策略，有效提高了推理服務性能，尤其在長上下文和真實場景下優勢更加顯著。

在ArXiv Summarization和L-Eval數據集上，Mooncake的吞吐量比baseline方法vLLM分別提高了20%和40%。

在模擬數據集上，Mooncake的吞吐量最高可達525%，在真實數據集上也可以比vLLM多處理約75%的請求。

過載場景下的性能評估結果則顯示，使用基于預測的早期拒絕策略時，拒絕的請求數量從baseline的4183個減少到了3589個，說明系統的請求處理能力得到了提高。

針對未來的發展，論文的另一位作者、清華大學計算機系助理教授章明星表示，從目前的趨勢來看，大模型服務的負載會愈發的復雜和多元化，調度會越來越復雜，也會越來越重要。

而對于月之暗面的發展方向，則是由許欣然做了解答——分布式策略的實施，也意味著未來月之暗面的整個系統，將往“算力/$”和“帶寬/$”兩個方向獨立發展，從而對硬件優化更加友好。

論文地址：https://arxiv.org/pdf/2407.00079
GitHub：https://github.com/kvcache-ai/Mooncake