Kimi論文自曝推理架構(gòu)，80%流量都靠它承擔(dān)

克雷西 2024-07-05 09:00:45 來源：量子位

克雷西 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI

月之暗面和清華KVCache.ai團(tuán)隊的最新論文，首次揭秘了Kimi背后的推理架構(gòu)！

要知道Kimi是國產(chǎn)大模型的當(dāng)紅炸子雞，火到可以說從來沒缺過流量，甚至還經(jīng)常出現(xiàn)過載。

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

Kimi背后的推理架構(gòu)名叫Mooncake（月餅），主要特點(diǎn)是采取了分離式的設(shè)計方案。

而且，Mooncake在設(shè)計之時就考慮了可能出現(xiàn)的大流量場景，并針對這種情況專門研發(fā)。

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

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

從KV緩存出發(fā)，建造分布式系統(tǒng)

整個Mooncake系統(tǒng)設(shè)計的核心，是圍繞著KV緩存展開的。

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

之所以這樣做，是因為團(tuán)隊預(yù)計KV緩存的容量會長期保持高位，因此圍繞KV緩存進(jìn)行優(yōu)化十分必要。

從結(jié)構(gòu)上看，Mooncake由全局調(diào)度器（Conductor）、Prefill節(jié)點(diǎn)集群、Decoding節(jié)點(diǎn)集群和分布式KVCache池幾部分組成，另外還有RDMA通信組件（Messenger）。

其中全局調(diào)度器是用戶請求到達(dá)系統(tǒng)后的第一站，它負(fù)責(zé)接收請求并根據(jù)KV緩存分布和負(fù)載情況，將請求調(diào)度到Prefill和Decoding節(jié)點(diǎn)。

調(diào)度器在調(diào)度時需要綜合考慮KV緩存的復(fù)用長度、負(fù)載均衡等因素，實現(xiàn)KV緩存復(fù)用的最大化。

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

同時，這種動態(tài)復(fù)制熱點(diǎn)KV緩存塊的方式，也是實現(xiàn)均衡負(fù)載的一種重要途徑。

實驗結(jié)果表明，與隨機(jī)調(diào)度和負(fù)載均衡調(diào)度相比，Mooncake的調(diào)度策略可以顯著降低TTFT（Time To First Token，首個Token延遲），提高系統(tǒng)性能。

完成調(diào)度之后，任務(wù)會分別交由Prefill和Decoding節(jié)點(diǎn)進(jìn)行運(yùn)算。

Prefill節(jié)點(diǎn)接收到調(diào)度器轉(zhuǎn)發(fā)過來的請求后，會從KV緩存池中讀取緩存，執(zhí)行預(yù)計算并生成新的KV緩存。

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

而Decoding節(jié)點(diǎn)除了接收調(diào)度器發(fā)來的請求外，還會收到Prefill階段生成的KV緩存，節(jié)點(diǎn)會對這些緩存執(zhí)行解碼并生成最終結(jié)果。

這當(dāng)中，大容量、高性能的KV緩存存儲由緩存池提供；RDMA通信組件則憑借其高帶寬、低延遲的優(yōu)勢，負(fù)責(zé)在不同節(jié)點(diǎn)之間的KV緩存?zhèn)鬏敗?/p>

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

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

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

這就導(dǎo)致了兩者在計算復(fù)雜度、內(nèi)存訪問方式、并行粒度和對延遲的敏感度上都存在差異：

所以，月之暗面團(tuán)隊對GPU集群也進(jìn)行了相應(yīng)的拆分，以便將它們分別部署在不同節(jié)點(diǎn)集群上，實現(xiàn)資源隔離和專門優(yōu)化。

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

提前預(yù)測負(fù)載，及時拒絕超量請求

不過，即使Mooncake采用了高效的分離架構(gòu)，但實際環(huán)境中的超大流量，對系統(tǒng)仍然是一個考驗。

對此，作者也提出了新的應(yīng)對策略。

在過載場景下，調(diào)度的關(guān)鍵是決定是否接受新的請求。

由于Mooncake采用的是分離式架構(gòu)，可以采取早期拒絕策略，在Prefill階段就根據(jù)Decoding節(jié)點(diǎn)的負(fù)載情況，提前拒絕請求。

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

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

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

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

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

針對這一問題，月之暗面團(tuán)隊對這種簡單的早期拒絕策略進(jìn)行了修正，提出了基于預(yù)測的早期拒絕策略，從而降低節(jié)點(diǎn)負(fù)載的波動。

這種策略的核心思想是對一段時間后的Decoding節(jié)點(diǎn)負(fù)載進(jìn)行預(yù)測，并基于預(yù)測結(jié)果決定是否拒絕請求。

預(yù)測可以在請求級別和系統(tǒng)級別兩個層面進(jìn)行，請求級別的預(yù)測比較困難，因為要預(yù)測單個請求的執(zhí)行時間；系統(tǒng)級別的預(yù)測相對容易一些，只需要預(yù)測整體的負(fù)載情況。

Mooncake采用的是一種簡化的系統(tǒng)級別預(yù)測方法，假設(shè)每個請求的執(zhí)行時間服從某個固定分布，據(jù)此預(yù)測未來一段時間內(nèi)的負(fù)載情況。

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

最終，端到端性能評估結(jié)果表明，Mooncake的架構(gòu)設(shè)計和優(yōu)化策略，有效提高了推理服務(wù)性能，尤其在長上下文和真實場景下優(yōu)勢更加顯著。

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

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

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

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

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

論文地址：
https://arxiv.org/pdf/2407.00079
GitHub：
https://github.com/kvcache-ai/Mooncake
參考鏈接：
[1]https://zhuanlan.zhihu.com/p/705910725
[2]https://zhuanlan.zhihu.com/p/706204757