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

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

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

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

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

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

　　在模擬場(chǎng)景下，Mooncake最高能帶來525%的吞吐量增長(zhǎng)，實(shí)際場(chǎng)景中也能多處理75%請(qǐng)求。

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

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

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

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

　　之所以這樣做，是因?yàn)閳F(tuán)隊(duì)預(yù)計(jì)KV緩存的容量會(huì)長(zhǎng)期保持高位，因此圍繞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)度器是用戶請(qǐng)求到達(dá)系統(tǒng)后的第一站，它負(fù)責(zé)接收請(qǐng)求并根據(jù)KV緩存分布和負(fù)載情況，將請(qǐng)求調(diào)度到Prefill和Decoding節(jié)點(diǎn)。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

　　提前預(yù)測(cè)負(fù)載，及時(shí)拒絕超量請(qǐng)求

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

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

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

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

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

　　具體的SLO要求是TTFT的90分位值(P90)不超過單個(gè)請(qǐng)求在空載條件下處理時(shí)間的10倍，TBT的P90值不超過5倍。

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

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

　　在階段1，Prefill節(jié)點(diǎn)和Decoding節(jié)點(diǎn)的負(fù)載都較低，此時(shí)調(diào)度器會(huì)持續(xù)接受新的請(qǐng)求，直到Prefill節(jié)點(diǎn)的負(fù)載達(dá)到上限。

　　進(jìn)入階段2后，Rrefill節(jié)點(diǎn)處理的請(qǐng)求開始進(jìn)入Decoding節(jié)點(diǎn)，導(dǎo)致其負(fù)載快速上升。當(dāng)Decoding節(jié)點(diǎn)的負(fù)載超過閾值后調(diào)度器開始拒絕新的請(qǐng)求，但此時(shí)Prefill節(jié)點(diǎn)的負(fù)載仍然很高。

　　到了階段3，由于調(diào)度器拒絕新請(qǐng)求，Prefill節(jié)點(diǎn)的負(fù)載開始下降。但此前積壓的請(qǐng)求正在Decoding階段處理，節(jié)點(diǎn)的負(fù)載仍然很高。

　　最后是階段4，Decoding節(jié)點(diǎn)的負(fù)載開始下降，因?yàn)榍懊娴恼?qǐng)求都處理完成，而新的請(qǐng)求又被拒絕了。這時(shí)調(diào)度器再次開始接受新請(qǐng)求，Prefill節(jié)點(diǎn)的負(fù)載又開始上升。

　　之后，這個(gè)過程會(huì)周期性地重復(fù)，導(dǎo)致Prefill和Decoding節(jié)點(diǎn)的負(fù)載出現(xiàn)反相位的波動(dòng)。

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

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

　　預(yù)測(cè)可以在請(qǐng)求級(jí)別和系統(tǒng)級(jí)別兩個(gè)層面進(jìn)行，請(qǐng)求級(jí)別的預(yù)測(cè)比較困難，因?yàn)橐A(yù)測(cè)單個(gè)請(qǐng)求的執(zhí)行時(shí)間;系統(tǒng)級(jí)別的預(yù)測(cè)相對(duì)容易一些，只需要預(yù)測(cè)整體的負(fù)載情況。

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

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

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

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

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

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

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

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