月之暗面kimi底層推理系統(tǒng)方案揭秘(二)
共 21301字,需瀏覽 43分鐘
·
2024-07-11 22:00
極市導(dǎo)讀
深入探討了月之暗面的底層推理平臺—Mooncake。文章基于作者在知乎上的討論和分享,特別是清華大學(xué)助理教授zhangmingxing的觀點(diǎn),詳細(xì)介紹了Mooncake的分離式架構(gòu)設(shè)計(jì)及其在實(shí)際應(yīng)用中的優(yōu)勢和挑戰(zhàn)。 >>加入極市CV技術(shù)交流群,走在計(jì)算機(jī)視覺的最前沿
既上一篇許欣然的月之暗面kimi底層推理系統(tǒng)方案揭秘,本篇繼續(xù)。
作者分享在知乎上引起了廣泛討論,很多system方向的大佬炸出來了,本文根據(jù)作者清華助理教授zhangmingxing的一些分享整理,歡迎大家去圍觀知乎討論區(qū):
李博杰(華為天才少年):2020年我開始做分離式內(nèi)存,做了非常高速的網(wǎng)絡(luò)互聯(lián)總線和很大的內(nèi)存池(不能透露更多技術(shù)指標(biāo)),當(dāng)時(shí)搞了存儲、數(shù)據(jù)庫、AI和云混合部署等幾個(gè)場景,感覺Persistent KV Cache會是disaggregated memory一個(gè)很重要的應(yīng)用場景。很高興看到Moonshot把它實(shí)現(xiàn)了。Prefill和Decoding分離,算力型硬件和內(nèi)存帶寬型硬件解耦,也是我非常喜歡的架構(gòu)設(shè)計(jì)。
作者:昨天華為云的workshop上還聊到這個(gè),分離式內(nèi)存終于算是有了一個(gè)非常核心的落地場景。不過主要是帶寬bound,iops bound 的場景想要落地還需要更多的探索,也包括 cxl/ub 之類硬件的進(jìn)一步成熟
《Mooncake (1): 在月之暗面做月餅,Kimi 以 KVCache 為中心的分離式推理架構(gòu)》
傳送門:https://zhuanlan.zhihu.com/p/705754254
作者Disclaimer:和論文不同本文夾帶過量私貨個(gè)人觀點(diǎn),不代表本人單位更不代表 Moonshot,純技術(shù)討論。
TL;DR Mooncake 是由月之暗面創(chuàng)建的超人氣智能助手 Kimi的底層推理平臺。
本系列則是對應(yīng)技術(shù)報(bào)告的濃縮版介紹。和強(qiáng)調(diào) inclusion 四平八穩(wěn)啥都講一點(diǎn)的論文格式不同,這里更多的是想討論一些當(dāng)前還未形成共識的一些 design choice (私貨警告)。特別的本篇主要討論 Mooncake 的分離式架構(gòu),討論點(diǎn)包括但不限于 TBT or TPOT,Prefill 節(jié)點(diǎn)應(yīng)不應(yīng)當(dāng)獨(dú)立存在以及獨(dú)立的話如何多節(jié)點(diǎn)并行處理,KVCache cache 全局調(diào)度的原理和調(diào)度策略,Decode 還能不能進(jìn)一步分離等等。
總體架構(gòu)
Mooncake 的架構(gòu)是非常典型的分離式架構(gòu),將單個(gè)同構(gòu) GPU 集群的資源打散并重新組織成三個(gè)可以獨(dú)立彈性伸縮的資源池。其中 Prefill Pool 處理用戶輸入,主要對 Time To First Token (TTFT) 負(fù)責(zé)。同時(shí)因?yàn)?Prefill 相對計(jì)算密集,這一部分也承擔(dān)著抬高整體資源利用率的任務(wù)。Prefill 處理完之后對應(yīng)的 KVCache 會被送到 Decode Pool 進(jìn)行 autoregression 式的流式輸出。雖然我們希望盡可能攢大的 batch 以提升 MFU,但這一部分主要需要對 Time Between Tokens (TBT) 負(fù)責(zé)。
這里我們使用 TBT 而非另外一些工作中常用的 Time Per Output Token (TPOT) 的原因在于雖然其實(shí)本來這兩者應(yīng)該等價(jià),但實(shí)際計(jì)算 TPOT 的時(shí)候經(jīng)常會被簡單地等價(jià)成計(jì)算 average TPOT。然后就直接用了吞吐的倒數(shù)或者說總生成時(shí)間除以總生成 token 數(shù)。TBT 的定義則更加明確的指定為兩個(gè) token 生成間的延遲,因此基于 TBT 上限設(shè)計(jì) Service Level Objective (SLO) 能更好的反映流式交互時(shí)的用戶體驗(yàn)。直觀的理解的話大家可以想象如果 GPT-4o 的語音回復(fù)是說一句停頓一段時(shí)間再說一句的說話方式的話,即便 average TPOT 能小一些體驗(yàn)還是不夠好。
在 Prefill/Decode 之外我們還利用每臺 HGX 機(jī)器上組成了一個(gè) KVCache Pool 來進(jìn)行全局的 Prefix Cache。相比 vLLM 當(dāng)前單機(jī)的 Prefix Cache 通過全局的調(diào)度能夠大幅度提升復(fù)用率從而提升總吞吐。由此帶來了一系列如何調(diào)度,分配,復(fù)制 KVCache 的問題,而且必須和 Prefill/Decode 的調(diào)度 co-design,因此我們把 Mooncake 的架構(gòu)稱之為以 KVCache 為中心。
在此基礎(chǔ)上接下來我們分別討論 Prefill/KVCache/Decode Pool 的關(guān)鍵 design choice。
Prefill Pool
論文這一章我原本起了個(gè)“Prefill: To Separate or Not. Is it a Question?” 的名字。主要原因在于由于 Chunked Prefill 的提出實(shí)際上 Prefill 和 Decode 節(jié)點(diǎn)的邊界已經(jīng)模糊了起來。不過除非是特別短的 prompt (可以直接一次性加入到 decode 的 continues batch 里提升 MFU 但不違反 TBT SLO 的那種)不然我們還是傾向于使用獨(dú)立的 Prefill 節(jié)點(diǎn)。原因主要是兩方面。
首先我們定義了一個(gè) VRAM occupation cost 的概念,即 KVCache 大小乘以其在顯存中留駐的時(shí)間。使用 chunked prefill 的話由于每個(gè) chunk 還需要和其他 decode request batch 在一起處理因此是顯著增加了對應(yīng) KVCache 在顯存中留駐的時(shí)間的,因此等效的降低了寶貴顯存空間的利用率。當(dāng)然顯存空出來了有沒有地方用就是另外一個(gè)問題了。這里我們在論文 4.2 節(jié)中主要討論了和未來 batch job 結(jié)合的可能性。
另外更重要的一個(gè)原因則主要是在于長文本場景下我們需要在 Prefill 集群中使用特殊的多節(jié)點(diǎn)分布式劃分方法來壓低 TTFT。這里我們討論了相比現(xiàn)在被討論的更多的 SP 策略直接使用 Chunked PP 策略在調(diào)度和減少通訊量方面的優(yōu)越性,前者大幅度簡化架構(gòu),后者則能夠省出寶貴的傳輸帶寬來搬運(yùn) KVCache。具體策略其實(shí)就是 TeraPipe 的推理簡化版本,因?yàn)橹挥?forward 沒有 backword 的話不需要啥動態(tài)規(guī)劃來平衡每個(gè) Chunk 的計(jì)算量,只要給每個(gè) Chunk 設(shè)置一個(gè)最小的值就行。
KVCache Pool
都 2024.6 了單節(jié)點(diǎn)的 Prefix Cache 已經(jīng)不是什么新鮮東西了。這里主要討論一下全局調(diào)度能夠成立的底層原因:對于每 X byte 的 KVCache,其生成所需的算力正比于 X*hd 再乘以一個(gè)較大的常數(shù),這里 hd 對應(yīng)模型的 hidden dimension。因此只要每卡算力比如 A100 的 220TFLOPS 和每卡通訊帶寬比如 CX7 的 100Gbps 的比值小于 hd 乘以這個(gè)常數(shù),那么從遠(yuǎn)端傳輸 KVCache 相比原地重算不僅僅減少了計(jì)算量還減少了 TTFT。又省錢用戶體驗(yàn)還更好了自然是何樂而不為。
但是上述公式中具體的 hd,常數(shù)是和模型結(jié)構(gòu)強(qiáng)相關(guān)的所以不是說一定可以,只能說原則上比較大的模型應(yīng)該都行。另外雖然網(wǎng)卡有 100Gbps,但是如果因?yàn)楸热?KVCache 下沉到慢速存儲上去了,網(wǎng)絡(luò)擁擠了之類的原因跑不到這么高那么自然也就不一定成立。所以還是需要以 KVCache 為中心針對不同請求的 TTFT SLO 分別去調(diào)度。在論文的 5.2 節(jié)中我們主要是提供了一種簡單的基于 heuristic 的熱點(diǎn)識別和復(fù)制方法。倒不是說有多么精妙,但重點(diǎn)是簡單好實(shí)現(xiàn),不用特別去預(yù)測未來的 KVCache 使用模式,而且在實(shí)際的使用中效果不錯(cuò)的方案。未來會嘗試更多的調(diào)度算法不過這個(gè)可以作為一個(gè) strong baseline 來使用。
順帶這里感慨一句,目前標(biāo)準(zhǔn)的 8 卡 HGX 服務(wù)器每臺會配 9~10 張 100Gbps 乃至 200Gbps 的 RDMA 網(wǎng)卡,至少 2TB 的內(nèi)存和一大堆的高速 SSD。假設(shè)一千臺的話就是近萬張網(wǎng)卡互連起來的數(shù) PB 共享內(nèi)存池。做 Disaggregated Memory 挺久了第一次見這樣的富裕仗。可惜作為一個(gè) bandwidth bound 場景在架構(gòu)上能設(shè)計(jì)的東西沒那么多特殊的奇淫技巧可用,更多還是工程實(shí)現(xiàn)。
Decode Pool
Decode Pool 本身業(yè)界的實(shí)踐是最多的。特別是開源有極快速地在迭代的 vLLM,該有的不是已經(jīng)有了也是在 roadmap 里。這里主要是夾帶一下 future work 中討論的增加面向大容量大帶寬設(shè)計(jì)的高性價(jià)比設(shè)備,形成異構(gòu)集群的私貨。從去年開始就斷斷續(xù)續(xù)和很多人討論這個(gè)話題,雖然聽聞了很多的風(fēng)聲,不過實(shí)際正用上的設(shè)備目測還得有個(gè)兩年。而且硬件畢竟迭代慢,出來了算法變了(比如 MLA)或者被老黃的大力出奇跡創(chuàng)飛了就很尷尬。
不過實(shí)際上我們幾個(gè)月前就掛出但沒多少人看的另一篇論文也提到,其實(shí)不需要等新的硬件,只要組合基于 GDDR 的設(shè)備其實(shí)就能將 Decode 部分進(jìn)一步拆分成 attention 算子和 Linear 算子兩個(gè)池子從而進(jìn)一步提升整體性能。不過目前云上的 GDDR 設(shè)備一般不會和旗艦顯卡集群放在一起,而且為了節(jié)省成本只會配很差的 CPU 和網(wǎng)卡,所以想要實(shí)際部署這一套方案還有很多挑戰(zhàn)。而且當(dāng)前設(shè)備畢竟不是專門為這個(gè)場景設(shè)計(jì)的,attention offload 策略性價(jià)比的提升的潛力沒有辦法充分發(fā)揮出來,為此進(jìn)一步將架構(gòu)復(fù)雜化而且增大一定的 TBT 不一定很合適。所以歸根結(jié)底還是得看看有沒有硬件廠商愿意賭一把了。
其它
恭喜若愚今天正式以優(yōu)秀畢業(yè)生的身份畢業(yè),本科就這么厲害博士記得帶我起飛~
《Mooncake (2):Kimi “潑天的流量”怎么接,分離架構(gòu)下基于預(yù)測的調(diào)度策略 》
傳送門:https://zhuanlan.zhihu.com/p/706204757
作者Disclaimer:和論文不同本文夾帶過量私貨個(gè)人觀點(diǎn),不代表本人單位更不代表 Moonshot,純技術(shù)討論。
TL;DR 正如 @許欣然 所說,本文公開的目的之一就是推動硬件廠商和云廠商向前篇提到的分離式,乃至未來異構(gòu)分離的方向演化。作為我個(gè)人已知范圍內(nèi)承載了最大規(guī)模在線流量的分離式推理系統(tǒng),Mooncake 也積攢了一些獨(dú)特的實(shí)踐經(jīng)驗(yàn)。下篇這里主要是想要討論其中非常關(guān)鍵的一項(xiàng):Kimi 在承接急速增長的需求時(shí),面向過載(Overload-oriented)場景的調(diào)度策略。
Kimi 也會累
和一般按照峰值容量規(guī)劃,調(diào)度主要解決如何盡可能利用閑置資源的傳統(tǒng)工作不同,Kimi 自從破圈之后面臨的是完全不同的,如何在不斷擴(kuò)容還天天過載的情況下盡可能保障用戶體驗(yàn)的問題。雖然是一個(gè)非常凡爾賽的甜蜜煩惱,同時(shí) infra 也已經(jīng)在最大程度避免,但在高峰期不可避免地還是需要偶爾向大家喊累。
這里主要是我們發(fā)現(xiàn)由于當(dāng)前的負(fù)載情況已經(jīng)無法按照 SLO 提供響應(yīng),所以與其提供一個(gè)完全無保障的長時(shí)間排隊(duì)(違反 TTFT SLO)或者一字一頓的慢速回答(違反 TBT SLO)這樣極大傷害體驗(yàn)的服務(wù),不如先暫停一小部分服務(wù)。至于為什么是喊累而不是顯示 429 Too Many Requests,當(dāng)然是因?yàn)?Kimi 是大家的智能_伙伴_,伙伴累了要休息一下希望大家可以體諒。
分離式架構(gòu)在過載場景下的挑戰(zhàn)
也正是在這樣的環(huán)境下正式上線了 Mooncake。雖然 Mooncake 大幅提升了 Kimi 能夠承載的總吞吐,但 Prefill 和 Decode 分離的架構(gòu)在過載調(diào)度上也的確引入了一些新的挑戰(zhàn)。這里最重要的問題就是 Prefill 和 Decode 是分別對 TTFT 和 TBT 負(fù)責(zé)的,而且有一個(gè)關(guān)鍵的時(shí)間差。所以在壞的情況下 Prefill 前調(diào)度器覺得可以同時(shí)滿足 TTFT 和 TBT,但是 Prefill 后由于 Decode 集群的過載 TBT 無法滿足了。這樣就進(jìn)入了一個(gè)是違反 SLO 還是浪費(fèi)已經(jīng)花費(fèi)了的 Prefill 資源的兩難問題。
為了解決上述問題,一個(gè)自然的,也是被 Mooncake 采用的解決方案就是同時(shí)綜合 Prefill 和 Decode 兩個(gè)集群的情況,然后以兩者中更高負(fù)載更危險(xiǎn)的那個(gè)集群為去判定是否可以接受服務(wù)。由于會因?yàn)樵?Prefill 集群有空閑的情況下由于未來 Decode 集群的負(fù)載問題提前拒絕一些服務(wù),這樣的策略被我們稱之為 Early Reject。
看起來很直接,但實(shí)際部署之后我們觀測到集群負(fù)載出現(xiàn)了奇怪的顛簸現(xiàn)象。可以看到 Prefill 和 Decode 集群的負(fù)載就和蹺蹺板一樣一邊上去一邊下來,然后交替。
仔細(xì)分析之后發(fā)現(xiàn)的原因下圖有一個(gè)很直觀的展示。由于 Prefill 和 Decode 集群負(fù)載之間的時(shí)間差,如果簡單的參考當(dāng)前 Decode 集群負(fù)載去拒絕請求的話會導(dǎo)致 Decode 集群負(fù)載被消化的時(shí)候 Prefill 沒有及時(shí)跟上,由此產(chǎn)生了蹺蹺板效應(yīng)。
基于預(yù)測的調(diào)度
為了解決上述的問題,我們進(jìn)一步設(shè)計(jì)和上線了基于預(yù)測的調(diào)度策略。原理也很直觀,能夠預(yù)測未來(特別是新加入的請求完成 Prefill 階段的時(shí)刻)Decode 集群負(fù)載的話自然就可以打消時(shí)間差的問題平滑整個(gè)系統(tǒng)的負(fù)載。
具體來說,對于未來的某個(gè)時(shí)刻 t,首先我們將現(xiàn)有的 Prefill 集群中 t 時(shí)已完成的請求加入 Decode 集群,這一點(diǎn)可以通過預(yù)測 Prefill 時(shí)間和排隊(duì)時(shí)間來實(shí)現(xiàn)。然后,我們假設(shè)每個(gè)請求的 decode 時(shí)間均為 t_d,將 Decode 集群中 t 時(shí)已結(jié)束(即decode 時(shí)間超過 t_d)的請求移除。這一步我們也在嘗試通過預(yù)測每個(gè)請求實(shí)際會輸出的 token 數(shù)做一個(gè)更精確的預(yù)測。最后,我們利用事先擬合的模型根據(jù) Decode 集群中 t 時(shí)仍在處理的請求預(yù)測 TBT,作為 Decode 集群的預(yù)測負(fù)載并根據(jù)這個(gè)數(shù)據(jù)決定是否需要 Early Reject。
未來:調(diào)度會越來越復(fù)雜,也會越來越重要 從目前的趨勢來看,未來 LLM Serving 的負(fù)載會愈發(fā)的復(fù)雜和多元化。原本同等地位的請求會在廠商和用戶的驅(qū)動下基于多種多樣的訴求進(jìn)一步分化。比如有些批處理請求可以完全不考慮 TTFT/TBT 僅僅需要保障 Job Completion Time (JCT) SLO,有些高級/付費(fèi)用戶無疑需要保證更加可靠和更高的 SLO。另外外部基于 Agent 的任務(wù)也不僅僅看單個(gè) request,而是看整個(gè) Workflow 的綜合 SLO,加上 Code Executor 等外部環(huán)境的交互后也會進(jìn)一步復(fù)雜化整個(gè)系統(tǒng)。
在這些領(lǐng)域上我們也是剛剛起步,探索了一些,但是需要研究的更多,也希望借此機(jī)會和大家更多的探討這一部分未來可能的優(yōu)化空間。
技術(shù)報(bào)告原文原文(powered by kimi/互相賦能就讓它自己宣傳自己吧)
1 引言
1.1 動機(jī):開發(fā) Mooncake 的原因
隨著大型語言模型(LLMs)在各種場景中的迅速采用,對LLM服務(wù)的工作負(fù)載已經(jīng)變得顯著多樣化。這些工作負(fù)載在輸入/輸出長度、到達(dá)頻率和分布上有所不同,最重要的是,它們需要不同種類的服務(wù)級別目標(biāo)(SLOs)。作為一個(gè)模型即服務(wù)(MaaS)提供商,Kimi [5] 的一個(gè)主要目標(biāo)是解決一個(gè)具有多個(gè)復(fù)雜約束的優(yōu)化問題。優(yōu)化目標(biāo)是在滿足不同級別的SLOs的約束下,最大化整體有效吞吐量,這直接影響收入。
為了實(shí)現(xiàn)這一目標(biāo),一個(gè)先決條件是充分利用GPU集群中可用的各種資源。具體來說,盡管當(dāng)前GPU服務(wù)器以高度集成的節(jié)點(diǎn)形式提供(例如,DGX/HGX超級計(jì)算機(jī) [6]),但有必要將它們解耦并重新結(jié)構(gòu)化為幾個(gè)分離的資源池,每個(gè)池針對不同但協(xié)作的目標(biāo)進(jìn)行優(yōu)化。例如,許多研究人員 [7, 8, 9] 已經(jīng)建議將預(yù)填充服務(wù)器(prefill servers)與解碼服務(wù)器(decoding servers)分開,因?yàn)長LM服務(wù)的這兩個(gè)階段具有非常不同的計(jì)算特性,在請求從預(yù)填充服務(wù)器轉(zhuǎn)移到解碼服務(wù)器時(shí),KVCache會發(fā)生變化。
基于這個(gè)想法,我們發(fā)現(xiàn)KVCache的調(diào)度是LLM服務(wù)調(diào)度的核心。為了提高整體吞吐量,通常有兩種一般方法:1)盡可能多地重用KVCache以減少所需的計(jì)算資源;2)最大化每個(gè)批次中的token數(shù)量以提高模型FLOPs利用率(MFU)。然而,從遠(yuǎn)程位置重用KVCache會延長首次token的時(shí)間(TTFT),而較大的批次大小會導(dǎo)致更大的token間時(shí)間(TBT)。因此,這兩種面向吞吐量的優(yōu)化的使用可能會導(dǎo)致違反與延遲相關(guān)的SLOs。
根據(jù)上述指導(dǎo)方針,我們提出了一個(gè)以KVCache為中心的分離設(shè)計(jì),用于調(diào)度和優(yōu)化。圖1展示了我們當(dāng)前的以KVCache為中心的分離架構(gòu),名為Mooncake。對于每個(gè)請求,全局調(diào)度器(Conductor)需要選擇一對預(yù)填充和解碼實(shí)例,并按以下步驟調(diào)度請求:1)盡可能多地將可重用的KVCache轉(zhuǎn)移到選定的預(yù)填充實(shí)例;2)以塊/層的方式完成預(yù)填充階段,并將輸出的KVCache連續(xù)地流式傳輸?shù)较鄳?yīng)的解碼實(shí)例;3)在解碼實(shí)例中加載KVCache,并將請求添加到連續(xù)批處理過程中以生成請求輸出。
盡管這個(gè)過程看起來簡單,但由于許多限制,選擇策略卻相當(dāng)復(fù)雜。在預(yù)填充階段,主要目標(biāo)是盡可能多地重用KVCache以避免冗余計(jì)算。然而,等待存儲在低級存儲上的KVCache可能會導(dǎo)致違反TTFT SLO。此外,對KVCache服務(wù)器的高需求可能會導(dǎo)致網(wǎng)絡(luò)擁塞,延長等待時(shí)間。因此,Conductor還負(fù)責(zé)預(yù)測KVCache塊的未來使用情況,并相應(yīng)地執(zhí)行調(diào)度操作,例如交換和復(fù)制。最熱門的塊應(yīng)該復(fù)制到多個(gè)節(jié)點(diǎn)以避免獲取擁塞,而最冷的塊應(yīng)該被交換出去以減少保留成本。預(yù)填充調(diào)度還受到預(yù)填充節(jié)點(diǎn)上DRAM空間可用性的限制,特別是當(dāng)大量內(nèi)存被保留用于全局KVCache池時(shí)。
相比之下,解碼階段有不同的優(yōu)化目標(biāo)和約束條件。目標(biāo)是將盡可能多的token聚合到一個(gè)解碼批次中以提高M(jìn)FU。然而,這一目標(biāo)不僅受到TBT SLO的限制,還受到可以包含在VRAM中的聚合KVCache總大小的限制。
更重要的是,現(xiàn)有的LLM服務(wù)研究假設(shè)資源充足,并專注于提高資源利用率。相比之下,當(dāng)前的GPU/加速器供應(yīng)有限,許多MaaS提供商面臨嚴(yán)重的過載問題,尤其是在高峰時(shí)段。在這種情況下的調(diào)度提出了現(xiàn)有工作尚未探索的獨(dú)特挑戰(zhàn)。例如,我們需要預(yù)測未來的負(fù)載,并在預(yù)填充階段后如果沒有可用的解碼插槽,則盡早拒絕某些請求,以節(jié)省浪費(fèi)的計(jì)算資源。然而,這種早期拒絕策略的直接實(shí)現(xiàn)令人驚訝地導(dǎo)致了過載的波動。這促使我們致力于預(yù)測特定查詢的生成長度,并在短期內(nèi)進(jìn)行整體負(fù)載預(yù)測,以實(shí)施更好的拒絕策略。同樣必要的是,對不同請求優(yōu)先級進(jìn)行分類,以實(shí)施基于優(yōu)先級的調(diào)度。在本文中,我們將這些問題總結(jié)為面向過載的調(diào)度,并展示了我們的初步研究結(jié)果。
1.2 Mooncake 的設(shè)計(jì)與結(jié)果
在本文的后續(xù)部分,我們首先介紹 Mooncake 架構(gòu)的概覽,包括其主要組件和處理請求的典型工作流程(§3)。然后,我們描述了在實(shí)現(xiàn)過程中做出的主要設(shè)計(jì)選擇,特別是那些當(dāng)前研究中未涵蓋的部分。
首先,在§4中,我們討論了如何實(shí)現(xiàn)一個(gè)單獨(dú)的預(yù)填充節(jié)點(diǎn)池,以無縫處理上下文長度的動態(tài)分配。我們采用了分塊流水線并行機(jī)制(Chunked Pipeline Parallelism, CPP)來擴(kuò)展單個(gè)請求在多個(gè)節(jié)點(diǎn)上的處理,這對于減少長上下文輸入的首次Token時(shí)間(TTFT)是必要的。與傳統(tǒng)的基于序列并行(Sequence Parallelism, SP)解決方案相比,CPP減少了網(wǎng)絡(luò)消耗,并簡化了對頻繁彈性擴(kuò)展的依賴。這種機(jī)制通過層級預(yù)填充進(jìn)一步補(bǔ)充,使KVCache的流式傳輸能夠重疊延遲。
接下來,在§5中,我們詳細(xì)闡述了我們的以KVCache為中心的請求調(diào)度算法,它平衡了實(shí)例負(fù)載和通過首次Token時(shí)間(TTFT)和Token間時(shí)間(TBT)SLOs測量的用戶體驗(yàn)。這包括一個(gè)基于啟發(fā)式的自動熱點(diǎn)遷移方案,它在不需要精確預(yù)測未來KVCache使用情況的情況下復(fù)制熱門KVCache塊。實(shí)驗(yàn)結(jié)果表明,我們的緩存感知調(diào)度可以顯著降低現(xiàn)實(shí)世界場景中的TTFT。在使用公共數(shù)據(jù)集、模擬數(shù)據(jù)和真實(shí)工作負(fù)載的端到端實(shí)驗(yàn)中,Mooncake 在長上下文場景中表現(xiàn)出色。與基線方法相比,Mooncake 能夠在滿足SLOs的同時(shí)實(shí)現(xiàn)高達(dá)525%的吞吐量增加。在真實(shí)工作負(fù)載下,Mooncake 使 Kimi 能夠處理75% 的更多請求。
最后,與假設(shè)所有請求都將被處理的現(xiàn)有LLM服務(wù)工作不同,由于Kimi 用戶請求的快速增長,Mooncake 持續(xù)面臨過載問題。因此,Mooncake 的調(diào)度涉及根據(jù)系統(tǒng)負(fù)載決定是否接受或拒絕傳入請求。在§6中,我們討論了我們獨(dú)特的早期拒絕策略的實(shí)現(xiàn),它減少了過載場景中浪費(fèi)的計(jì)算資源。我們進(jìn)一步探討了由簡單早期拒絕引起的負(fù)載波動問題,以及如何通過預(yù)測未來負(fù)載來緩解這個(gè)問題。
Mooncake 目前是服務(wù) Kimi 的主要平臺,并已成功處理了呈指數(shù)級增長的工作負(fù)載,證明了其在擴(kuò)展到大型和高度過載工作負(fù)載方面的有效性。然而,還有更多問題需要探索,這些未來的發(fā)展方向也包括在本文中。
為了保護(hù)專有信息并促進(jìn)可復(fù)制性,本文報(bào)告的所有實(shí)驗(yàn)結(jié)果都是基于真實(shí)工作負(fù)載的重放跟蹤,但使用的是一個(gè)遵循與LLaMA2-70B相同架構(gòu)的虛擬模型。該跟蹤僅包括請求到達(dá)的時(shí)間、輸入token的數(shù)量和輸出token的數(shù)量,不包含任何真實(shí)用戶內(nèi)容。在遵循某些內(nèi)部程序后,該跟蹤將稍后開源。
2 預(yù)備知識和問題定義
現(xiàn)代大型語言模型(LLMs)基于Transformer架構(gòu),該架構(gòu)利用注意力機(jī)制和多層感知器(MLP)來處理輸入。基于流行的Transformer模型,例如GPT [10] 和 LLaMA [11],采用的是僅解碼器結(jié)構(gòu)。每個(gè)推理請求在邏輯上被劃分為兩個(gè)階段:預(yù)填充階段和解碼階段。
在預(yù)填充階段,所有輸入token并行處理。此階段生成第一個(gè)輸出token,同時(shí)存儲計(jì)算出的中間結(jié)果,這些中間結(jié)果被稱為KVCache。解碼階段隨后使用這個(gè)KVCache來自回歸地生成新的token,將新計(jì)算出的鍵和值添加到KVCache中。預(yù)填充階段能夠同時(shí)處理輸入token,通常使其在計(jì)算上非常密集,除了短請求之外。由于注意力網(wǎng)絡(luò)的計(jì)算復(fù)雜度與輸入長度呈二次方增長,如圖2左側(cè)所示,預(yù)填充階段的計(jì)算時(shí)間通常隨著輸入長度的增加而超線性增長。
與此相反,解碼階段由于自回歸生成的限制,每個(gè)批次只處理一個(gè)token。這使得它受內(nèi)存限制,并且計(jì)算時(shí)間隨著批次大小的增加而次線性增長,如圖2右側(cè)所示。在解碼階段廣泛使用的優(yōu)化是連續(xù)批處理 [12, 13]。在每次迭代之前,調(diào)度器檢查所有請求的狀態(tài),將新到達(dá)的請求添加到批次的預(yù)填充階段,同時(shí)移除已完成的請求。
由于預(yù)填充階段和解碼階段的不同特性,MaaS提供商設(shè)置了不同的度量標(biāo)準(zhǔn)來衡量它們各自的服務(wù)級別目標(biāo)(SLOs)。具體來說,預(yù)填充階段主要關(guān)注請求到達(dá)和生成第一個(gè)token之間的延遲,即首次Token時(shí)間(TTFT)。另一方面,解碼階段側(cè)重于同一請求連續(xù)token生成之間的延遲,即Token間時(shí)間(TBT)。
作為一個(gè)MaaS提供商,通過滿足服務(wù)協(xié)議中定義的SLO度量標(biāo)準(zhǔn)來確保質(zhì)量保證至關(guān)重要。例如,TTFTP 90 = 4×這樣的度量標(biāo)準(zhǔn)表明,在相同條件下無干擾地運(yùn)行的單個(gè)請求的TTFT的90%不超過四倍。具體來說,在本文的端到端實(shí)驗(yàn)(§7.1)中,我們設(shè)置了TTFTP 90 = 10×和TBTP 90 = 5×。在真實(shí)部署中,我們設(shè)置了固定的TTFT和TBT SLOs。如果監(jiān)控檢測到未滿足的SLOs,我們要么增加推理資源,要么拒絕一些傳入的請求。
然而,由于當(dāng)前GPU的有限供應(yīng),通常不可行地?cái)U(kuò)展推理集群。因此,在面向過載的調(diào)度中,決定拒絕哪些請求成為核心問題。我們的主要目標(biāo)是在遵守SLOs的同時(shí)最大化整體吞吐量,這一概念在其他研究中被稱為goodput [8, 14]。我們的方法不同之處在于,只有完全完成執(zhí)行的請求才會計(jì)入goodput的度量。否則,所有之前消耗/生成的token都不會被計(jì)算,相應(yīng)的資源就會被浪費(fèi)。換句話說,如果一個(gè)請求在SLO下無法完成其全部執(zhí)行,則應(yīng)盡早拒絕該請求。
實(shí)現(xiàn)這一目標(biāo)不僅需要優(yōu)化預(yù)填充和解碼階段的架構(gòu),還需要開發(fā)預(yù)測短期未來負(fù)載的能力。
3 Mooncake 分離架構(gòu)的概述
正如圖1所示,Mooncake 采用了一種分離架構(gòu),它不僅將預(yù)填充(prefill)節(jié)點(diǎn)與解碼(decoding)節(jié)點(diǎn)分開,而且還將GPU集群的CPU、DRAM、SSD和RDMA資源分組,實(shí)現(xiàn)了一個(gè)分離的KVCache。這種分離緩存利用了未充分利用的資源,提供了充足的緩存容量和傳輸帶寬,使得在不增加額外成本的情況下,能夠高效地實(shí)現(xiàn)接近GPU的前綴緩存。
KVCache池在CPU內(nèi)存中的存儲和傳輸邏輯如圖3所示。在CPU內(nèi)存中,KVCache被存儲為分頁塊。根據(jù)請求模式,它可以采用如最近最少使用(LRU)、最少頻繁使用(LFU)等緩存逐出算法,或基于請求特征的算法。這些KVCache塊在CPU和GPU之間的傳輸由一個(gè)稱為Messenger的獨(dú)立(GPUDirect) RDMA基礎(chǔ)組件處理。此架構(gòu)還使我們能夠?yàn)橥獠坑脩籼峁┥舷挛木彺鍭PI,以實(shí)現(xiàn)KVCache的更高重用。
為了調(diào)度所有這些分離組件,在Mooncake的中心,實(shí)現(xiàn)了一個(gè)名為Conductor的全局調(diào)度器。Conductor負(fù)責(zé)根據(jù)當(dāng)前的KVCache分布和工作負(fù)載來調(diào)度請求。如果對未來推理有利,它還會復(fù)制或交換某些KVCache塊。具體來說,圖4展示了一個(gè)請求的典型工作流程。一旦完成令牌化,Conductor就選擇一對預(yù)填充節(jié)點(diǎn)和一個(gè)解碼節(jié)點(diǎn),并啟動一個(gè)包含四個(gè)步驟的工作流程:
1) KVCache重用: 被選定的預(yù)填充節(jié)點(diǎn)(組)接收一個(gè)請求,該請求包括原始輸入、可以重用的前綴緩存塊ID,以及分配給請求的完整緩存塊ID。它根據(jù)前綴緩存塊ID從遠(yuǎn)程CPU內(nèi)存加載前綴緩存到GPU內(nèi)存以啟動請求。如果沒有前綴緩存存在,則跳過此步驟。此選擇平衡了三個(gè)目標(biāo):盡可能重用KVCache、平衡不同預(yù)填充節(jié)點(diǎn)的工作負(fù)載,并保證TTFT SLO。這導(dǎo)致了將進(jìn)一步討論的以KVCache為中心的調(diào)度。
2) 增量預(yù)填充: 預(yù)填充節(jié)點(diǎn)(組)使用前綴緩存完成預(yù)填充階段,并將新生成的增量KVCache存儲回CPU內(nèi)存。如果未緩存的輸入token數(shù)量超過某個(gè)閾值(prefill_chunk),預(yù)填充階段將被分成多個(gè)塊,并以流水線方式執(zhí)行。此閾值被選擇以充分利用相應(yīng)GPU的計(jì)算能力,通常大于1000個(gè)token。使用分塊但仍然分離的預(yù)填充節(jié)點(diǎn)的原因在§4.1中解釋。
3) KVCache傳輸: 前述的Messenger服務(wù)在每個(gè)節(jié)點(diǎn)中部署,用于管理和傳輸這些緩存。每個(gè)Messenger在其各自的推理實(shí)例中作為一個(gè)獨(dú)立進(jìn)程運(yùn)行,接收信號以促進(jìn)高速跨機(jī)器KVCache傳輸。這一步是異步執(zhí)行的,并與上述增量預(yù)填充步驟重疊,將每個(gè)模型層生成的KVCache流式傳輸?shù)侥繕?biāo)解碼節(jié)點(diǎn)的CPU內(nèi)存中,以減少等待時(shí)間。
4) 解碼: 在解碼節(jié)點(diǎn)的CPU DRAM中接收到所有KVCache后,請求以連續(xù)批處理的方式加入下一個(gè)批次。Conductor根據(jù)其當(dāng)前負(fù)載預(yù)先選擇解碼節(jié)點(diǎn),以確保它不會違反TBT SLO。然而,這種SLO會由本地調(diào)度器再次檢查,因?yàn)轭A(yù)計(jì)的負(fù)載可能在預(yù)填充階段之后發(fā)生了變化。這種雙重檢查可能導(dǎo)致請求被拒絕,在這種情況下,相應(yīng)的預(yù)填充成本就會被浪費(fèi)。
4 預(yù)填充池的實(shí)現(xiàn)
與不可違反的解碼節(jié)點(diǎn)不同,設(shè)計(jì)一個(gè)單獨(dú)且彈性的預(yù)填充池的必要性和最佳實(shí)踐仍然存在爭議。例如,盡管許多研究人員[7, 8, 9]與我們有相同的直覺,即使用分離架構(gòu),但在引入分塊預(yù)填充[15]后,分離是否仍然必要仍然值得討論。分塊預(yù)填充將輸入token分成多個(gè)小塊,這些小塊加入連續(xù)批處理過程。這種方法有兩個(gè)明顯的好處:1) 無需分離,所有節(jié)點(diǎn)被視為平等的,使調(diào)度變得更容易;2) 將分塊預(yù)填充內(nèi)聯(lián)到解碼批處理中可以提高解碼批處理的計(jì)算強(qiáng)度,從而實(shí)現(xiàn)更好的模型FLOPs利用率(MFU)。
然而,在仔細(xì)考慮后,我們決定保持Mooncake的分離架構(gòu)。只有在不需要分塊且不影響TBT SLO的情況下,請求的預(yù)填充才會內(nèi)聯(lián)到解碼批處理中。做出此決定有兩個(gè)主要原因:1) 預(yù)填充節(jié)點(diǎn)需要不同的跨節(jié)點(diǎn)并行設(shè)置來處理長上下文(§4.1);2) 它提供了節(jié)省VRAM的獨(dú)特機(jī)會(§4.2)。
4.1 多節(jié)點(diǎn)預(yù)填充
最近LLMs的可用上下文長度正在迅速增加,從8k到128K甚至1M[16]。通常,對于這樣的長上下文請求,輸入token可能是輸出token的10到100倍,這使得優(yōu)化TTFT變得至關(guān)重要。由于長上下文預(yù)填充中存在大量的并行性,因此使用超過一個(gè)8x GPU節(jié)點(diǎn)來并行處理是可取的。然而,將張量并行性(TP)擴(kuò)展到超過一個(gè)節(jié)點(diǎn)需要每層進(jìn)行兩次昂貴的基于RDMA的all-reduce操作,顯著降低了預(yù)填充節(jié)點(diǎn)的MFU。
最近,許多工作提出了序列并行性(SP)[17, 18, 19, 20, 21, 22, 23]。SP通過在不同節(jié)點(diǎn)上分配請求的輸入序列來實(shí)現(xiàn)加速。這些SP方法利用了注意力運(yùn)算符的關(guān)聯(lián)屬性,并且至少需要在每層實(shí)現(xiàn)Ring Attention[18]或Striped Attention[19]期間進(jìn)行一次跨節(jié)點(diǎn)通信。這大大減少了網(wǎng)絡(luò)消耗并提高了MFU。
然而,即使采用SP,與僅使用單節(jié)點(diǎn)TP相比,MFU仍然更差。理想的部署是將預(yù)填充節(jié)點(diǎn)組織成兩組:一組僅使用TP,另一組使用SP。僅在必要時(shí)為滿足TTFT SLO時(shí)才將請求調(diào)度到SP組。這種進(jìn)一步的分離導(dǎo)致了動態(tài)調(diào)整每個(gè)組中節(jié)點(diǎn)數(shù)量的問題,因?yàn)殪o態(tài)并行設(shè)置可能導(dǎo)致集群的低利用率。最近的研究[14]提出了彈性序列并行性來動態(tài)擴(kuò)展或縮減SP組。盡管這是可能的,但它為我們的架構(gòu)增加了復(fù)雜性。例如,它需要提前建立一個(gè)全局通信組,并在考慮如緩存重用利用率和SLO要求違規(guī)等指標(biāo)時(shí),使Conductor的設(shè)計(jì)復(fù)雜化。這使得我們的系統(tǒng)在需要頻繁即時(shí)可伸縮性的情況下面臨挑戰(zhàn)。此外,SP仍然需要頻繁的跨節(jié)點(diǎn)通信,這降低了MFU,并與跨節(jié)點(diǎn)傳輸KVCache的網(wǎng)絡(luò)資源競爭。
為了解決這個(gè)問題,Mooncake利用了僅解碼器變換器的自回歸屬性,并為長上下文預(yù)填充實(shí)現(xiàn)了分塊流水線并行性(Chunked Pipeline Parallelism, CPP)。我們將預(yù)填充集群中的每X個(gè)節(jié)點(diǎn)分組為一個(gè)流水線預(yù)填充節(jié)點(diǎn)組。對于每個(gè)請求,其輸入token被劃分為塊,每個(gè)塊的長度不超過prefill_chunk。相同請求的不同塊可以由不同的節(jié)點(diǎn)同時(shí)處理,從而并行化處理并減少TTFT。
CPP提供了兩個(gè)主要好處:1) 與訓(xùn)練中的流水線并行性類似,它僅在每個(gè)流水線階段的邊界處需要跨節(jié)點(diǎn)通信,這可以容易地與計(jì)算重疊。這帶來了更好的MFU和更少的網(wǎng)絡(luò)資源爭奪與KVCache傳輸。2) 它自然適應(yīng)短上下文和長上下文,為短上下文預(yù)填充帶來了沒有顯著開銷的好處,并避免了頻繁的動態(tài)節(jié)點(diǎn)劃分調(diào)整。這種基于流水線的加速方法已經(jīng)在訓(xùn)練系統(tǒng)中被探索[24],但據(jù)我們所知,這是在推理階段的首次應(yīng)用,因?yàn)殚L上下文推理直到最近才出現(xiàn)。
4.2 逐層預(yù)填充
除了計(jì)算能力外,VRAM的有限大小也是寶貴的資源,我們的目標(biāo)是通過狀態(tài),主要是KVCache,最小化VRAM的占用。理論上,如果一個(gè)請求的KVCache大小為S,處理時(shí)間為T,則其占用成本為S*T。如果一個(gè)請求被分塊,并且每個(gè)塊的處理與解碼請求在分塊預(yù)填充中內(nèi)聯(lián),則T將增加,導(dǎo)致更大的占用成本。
此外,由于預(yù)填充是逐層處理并且受計(jì)算限制,因此可以將KVCache的傳輸和轉(zhuǎn)儲與計(jì)算重疊,進(jìn)一步降低其占用成本。在Mooncake中,KVCache的加載和存儲通過啟動和等待操作異步執(zhí)行。在每層的注意力計(jì)算開始之前,模型等待該層的KVCache異步加載完成,并觸發(fā)下一層的異步KVCache加載。在注意力計(jì)算完成后,啟動該層KVCache的異步存儲。一旦所有層的計(jì)算完成,進(jìn)程等待所有異步存儲操作的完成。傳輸重疊允許預(yù)填充實(shí)例的執(zhí)行時(shí)間大致等同于KVCache加載時(shí)間或標(biāo)準(zhǔn)預(yù)填充時(shí)間,這取決于相對于輸入長度的前綴緩存比例。
實(shí)驗(yàn)結(jié)果,如圖5所示,KVCache存儲延遲表明,逐層預(yù)填充可以有效地減少長上下文請求的延遲。這種重疊效果的主要優(yōu)點(diǎn)是,它使我們能夠在預(yù)填充調(diào)度中忽略可用VRAM的大小,只要它能夠容納一個(gè)請求。如圖1所示,預(yù)填充節(jié)點(diǎn)的調(diào)度只考慮KVCache的分布和可用的DRAM大小。
將來,我們打算探索這種空閑VRAM的更多用途。例如,OpenAI最近提出了使用批量API[25],它使用戶能夠以50%的低成本發(fā)送異步請求組,但只有明確的24小時(shí)周轉(zhuǎn)時(shí)間。此服務(wù)非常適合處理不需要立即響應(yīng)的工作。由于這些批量請求沒有嚴(yán)格的TBT,我們可以在有足夠的VRAM空間容納相應(yīng)的KVCache的情況下,甚至將這些請求的解碼階段內(nèi)聯(lián)到預(yù)填充處理中以獲得更好的MFU。
5 以KVCache為中心的調(diào)度
在本節(jié)中,我們主要討論在正常條件下,Conductor如何調(diào)度請求和KVCache塊,下一部分將討論過載場景下的討論。
5.1 預(yù)填充全局調(diào)度
以往關(guān)于LLM服務(wù)的研究通常使用基于每個(gè)實(shí)例分配請求數(shù)量的負(fù)載均衡策略。然而,在Mooncake中,選擇預(yù)填充實(shí)例時(shí)會考慮額外的因素——不僅是負(fù)載,還有前綴緩存命中長度和可重用KVCache塊的分布。雖然傾向于將請求路由到具有較長前綴緩存長度的預(yù)填充實(shí)例以減少計(jì)算成本,但有時(shí)為了確保整體系統(tǒng)平衡和滿足TTFT SLO,將它們調(diào)度到其他節(jié)點(diǎn)可能是有益的。為了解決這些復(fù)雜性,我們提出了一個(gè)考慮前綴緩存導(dǎo)致的預(yù)填充時(shí)間和實(shí)例負(fù)載相關(guān)的排隊(duì)時(shí)間的緩存感知全局調(diào)度算法。
算法1詳細(xì)說明了我們的緩存感知預(yù)填充調(diào)度機(jī)制。對于每個(gè)新請求,其輸入token被分成幾個(gè)塊,并為每個(gè)塊計(jì)算一個(gè)哈希鍵。這涉及到為一個(gè)塊中的token生成一個(gè)哈希鍵,并將該哈希鍵與前一個(gè)塊的哈希鍵(如果有)連接起來。然后,將請求的塊鍵與每個(gè)預(yù)填充實(shí)例的緩存鍵逐一比較,以確定前綴匹配長度(prefix_len)。類似的重用邏輯已經(jīng)在vLLM中實(shí)現(xiàn),但vLLM的開源版本僅支持本地KVCache緩存。
有了這些匹配信息,Conductor根據(jù)請求長度和prefix_len(因?qū)嵗悾┕烙?jì)相應(yīng)的執(zhí)行時(shí)間。然后,它將該請求的估計(jì)等待時(shí)間加到該實(shí)例的TTFT上。最后,Conductor將請求分配給具有最短TTFT的實(shí)例,并相應(yīng)地更新該實(shí)例的緩存和隊(duì)列時(shí)間。如果無法達(dá)到SLO,Conductor直接向上層返回HTTP 429 Too Many Requests響應(yīng)狀態(tài)碼。
這個(gè)調(diào)度框架的核心是直接的,但是各種組件的工程實(shí)現(xiàn)中隱藏著復(fù)雜性。例如,為了預(yù)測請求的預(yù)填充階段的計(jì)算時(shí)間,我們采用了一個(gè)基于離線測試數(shù)據(jù)的預(yù)測模型。該模型根據(jù)請求的長度和前綴緩存命中長度估計(jì)預(yù)填充持續(xù)時(shí)間。由于Transformer的計(jì)算模式是規(guī)則的,只要有足夠的離線數(shù)據(jù),這個(gè)預(yù)測的誤差范圍就很小。一個(gè)請求的排隊(duì)時(shí)間是通過聚合所有排隊(duì)請求的預(yù)填充時(shí)間來計(jì)算的。在實(shí)際實(shí)現(xiàn)中,TTFTs是并行計(jì)算的,使得處理時(shí)間與推理時(shí)間相比可以忽略不計(jì)。
預(yù)測傳輸時(shí)間更加困難,因?yàn)樗粌H由傳輸數(shù)據(jù)的大小決定,還由當(dāng)前的網(wǎng)絡(luò)狀態(tài)決定,特別是發(fā)送節(jié)點(diǎn)是否處于擁塞狀態(tài)。這也需要復(fù)制熱門KVCache塊,這將在下一節(jié)中討論。
5.2 緩存負(fù)載均衡
在我們的Mooncake集群中,每臺預(yù)填充機(jī)器管理自己的一組本地前綴緩存。這些緩存的使用頻率差異很大。例如,系統(tǒng)提示幾乎被每個(gè)請求訪問,而存儲本地長文檔內(nèi)容的緩存可能只被一個(gè)用戶使用。如§5.1所討論的,Conductor在實(shí)現(xiàn)緩存匹配和實(shí)例負(fù)載之間的最佳平衡中起著至關(guān)重要的作用。因此,從分布式緩存系統(tǒng)的角度看,負(fù)載均衡也起著重要作用。具體來說,它涉及策略制定,如何備份緩存以確保全局預(yù)填充調(diào)度能夠?qū)崿F(xiàn)高緩存命中率和低負(fù)載。
解決這個(gè)KVCache調(diào)度問題的一個(gè)初步解決方案可能是收集每個(gè)塊的全局使用情況,使用預(yù)測模型預(yù)測它們未來的使用情況,并據(jù)此做出調(diào)度決策。然而,與預(yù)填充時(shí)間的估計(jì)不同,工作負(fù)載是高度動態(tài)的,并且隨時(shí)間顯著變化。特別是對于用戶基數(shù)快速增長的MaaS提供商來說,準(zhǔn)確預(yù)測未來的使用情況是不可能的。因此,我們提出了一個(gè)基于啟發(fā)式的自動熱點(diǎn)遷移方案來增強(qiáng)緩存負(fù)載均衡。
正如前所述,由于高實(shí)例負(fù)載,請求可能不會總是被定向到具有最長前綴緩存長度的預(yù)填充實(shí)例。在這種情況下,Conductor將緩存的位置和請求轉(zhuǎn)發(fā)到一個(gè)替代實(shí)例,如果估計(jì)的額外預(yù)填充時(shí)間短于傳輸時(shí)間。這個(gè)實(shí)例主動從持有者檢索KVCache并將其存儲在本地。更重要的是,我們更傾向于計(jì)算輸入token,如果最佳的遠(yuǎn)程前綴匹配長度不大于當(dāng)前本地可重用前綴乘以一個(gè)閾值。這兩種策略不僅減少了請求的預(yù)填充時(shí)間,還促進(jìn)了熱點(diǎn)緩存的自動復(fù)制,允許它們在多臺機(jī)器上更廣泛地分布。
6 面向過載的調(diào)度
大多數(shù)現(xiàn)有的關(guān)于LLM服務(wù)的研究都假設(shè)所有請求都將被處理,相應(yīng)地優(yōu)化吞吐量或請求的TTFT和TBT。然而,在實(shí)際情況下,處理每一個(gè)傳入的請求既不經(jīng)濟(jì)也不現(xiàn)實(shí)。對于面臨用戶請求量迅速增長的商業(yè)推理服務(wù)來說,集群的推理資源增長速度遠(yuǎn)遠(yuǎn)落后于傳入請求的增加。因此,過載是當(dāng)前LLM服務(wù)中常見的問題,尤其是在高峰時(shí)段。
為了平衡成本和用戶體驗(yàn),系統(tǒng)應(yīng)該盡可能多地處理請求,直到系統(tǒng)負(fù)載達(dá)到一個(gè)預(yù)定義的閾值。達(dá)到這一點(diǎn)后,剩余的請求將被直接拒絕或推遲到以后再重試。作為實(shí)現(xiàn)為分離式推理系統(tǒng)的Mooncake,允許更靈活的調(diào)度策略,但也面臨非分離式系統(tǒng)所沒有的獨(dú)特調(diào)度挑戰(zhàn),且在以前的工作中沒有提及。
在這一部分,我們描述了一個(gè)為分離式架構(gòu)特別設(shè)計(jì)的早期拒絕策略,并解決了這種方法引起的負(fù)載波動問題。然后,我們探討了預(yù)測生成長度對于緩解這些問題的必要性。
6.1 過載場景中的調(diào)度
在系統(tǒng)過載發(fā)生的情況下,調(diào)度涉及到根據(jù)系統(tǒng)負(fù)載決定是否接受或拒絕傳入的請求。這一過程中一個(gè)關(guān)鍵方面是定義什么構(gòu)成“系統(tǒng)負(fù)載”,因?yàn)檫@一定義影響請求被拒絕的閾值。在傳統(tǒng)的耦合系統(tǒng)中,由于預(yù)填充和解碼階段之間的干擾,預(yù)測TTFT和TBT可能變得復(fù)雜。因此,負(fù)載通常簡單地通過正在處理的請求數(shù)量與系統(tǒng)最大容量的比率來衡量。
與此相反,Mooncake以其分離式架構(gòu)處理預(yù)填充和解碼階段的獨(dú)立性。因此,我們使用SLO滿足度作為直接的負(fù)載測量。具體來說,我們定義了lttft和ltbt分別為請求的TTFT和TBT SLO約束。然后,預(yù)填充和解碼實(shí)例的負(fù)載通過將實(shí)例上預(yù)測的最大TTFT和TBT與lttft和ltbt進(jìn)行比較來確定。有了這兩個(gè)標(biāo)準(zhǔn),Mooncake的調(diào)度需要做出兩個(gè)關(guān)鍵決策:第一,根據(jù)預(yù)填充實(shí)例的負(fù)載是否接受預(yù)填充階段;第二,根據(jù)解碼實(shí)例的負(fù)載決定是否繼續(xù)解碼階段。
6.2 早期拒絕
在實(shí)踐中,預(yù)填充或解碼實(shí)例的個(gè)別負(fù)載并不準(zhǔn)確地反映系統(tǒng)處理的實(shí)際請求數(shù)量。這種差異源于調(diào)度單個(gè)請求的預(yù)填充和解碼實(shí)例之間存在時(shí)間滯后。如果一個(gè)請求在預(yù)填充階段完成后由于解碼實(shí)例的高負(fù)載而被拒絕,那么在預(yù)填充階段花費(fèi)的計(jì)算資源就被浪費(fèi)了。因此,預(yù)填充階段成功處理的請求實(shí)際數(shù)量少于負(fù)載指標(biāo)所指示的數(shù)量。
為解決這個(gè)問題,很自然地將解碼實(shí)例的負(fù)載評估提前到預(yù)填充階段開始之前。我們將此策略稱為早期拒絕。在請求到達(dá)時(shí),Conductor根據(jù)預(yù)填充和解碼池之間的較大負(fù)載評估是否接受請求。早期拒絕顯著減少了被拒絕請求的無效計(jì)算,并提高了負(fù)載均衡。
6.3 早期拒絕引起的負(fù)載波動
然而,早期拒絕引入了新的挑戰(zhàn)。圖7顯示了在使用基于預(yù)測的早期拒絕策略后,在20分鐘內(nèi)觀察到的20臺機(jī)器集群的預(yù)填充和解碼實(shí)例負(fù)載。它突出了預(yù)填充和解碼機(jī)器之間的顯著反相波動。當(dāng)預(yù)填充機(jī)器較少,以及預(yù)填充階段較長的場景中,這種現(xiàn)象更加明顯。
經(jīng)進(jìn)一步探索,我們發(fā)現(xiàn)這種負(fù)載波動問題根源于預(yù)測解碼負(fù)載與實(shí)際執(zhí)行之間的時(shí)間差。基于當(dāng)前解碼負(fù)載的調(diào)度本質(zhì)上存在延遲。這個(gè)延遲導(dǎo)致預(yù)填充和解碼實(shí)例的負(fù)載出現(xiàn)波動和相位錯(cuò)位,如圖8a中描述的理論示例所示。綠色曲線代表預(yù)填充實(shí)例的負(fù)載(按0到1的比例縮放),黃色曲線代表解碼實(shí)例的負(fù)載。
在第一階段,預(yù)填充和解碼實(shí)例的負(fù)載都很低,因此Conductor接受了大量的請求,直到預(yù)填充實(shí)例的負(fù)載達(dá)到其限制。在第二階段,由預(yù)填充實(shí)例處理的請求被調(diào)度到解碼實(shí)例,導(dǎo)致解碼實(shí)例的負(fù)載很高。因此,Conductor拒絕傳入的請求,導(dǎo)致預(yù)填充實(shí)例的負(fù)載降低。在第三階段,沒有新的請求進(jìn)入解碼階段,導(dǎo)致負(fù)載減少。此時(shí),Conductor再次接受大量的請求,直到預(yù)填充實(shí)例完全負(fù)載。在第四階段,隨著解碼實(shí)例的負(fù)載增加,Conductor拒絕請求,導(dǎo)致預(yù)填充實(shí)例的負(fù)載降低。這種預(yù)填充和解碼實(shí)例之間負(fù)載的嚴(yán)重波動導(dǎo)致推理集群的資源利用率低下。
6.4 基于預(yù)測的早期拒絕
為解決負(fù)載波動問題,我們提出了一個(gè)面向過載場景的分離式LLM服務(wù)系統(tǒng)(如Mooncake)的基于預(yù)測的早期拒絕框架。如圖8b所示,該框架預(yù)測了預(yù)填充階段后傳入請求的解碼負(fù)載,并使用此預(yù)測來決定是否接受請求,這有助于緩解波動問題。該策略的核心組件是準(zhǔn)確預(yù)測后續(xù)期間的解碼負(fù)載。我們?yōu)榇艘肓藘煞N方法:
請求級別:以前的工作強(qiáng)調(diào)了預(yù)測LLM服務(wù)負(fù)載的一個(gè)重大挑戰(zhàn):每個(gè)請求的未知輸出長度。如果我們能夠提前確定輸出長度,就有可能更準(zhǔn)確地估計(jì)TTFT和TBT。這反過來有助于預(yù)測解碼實(shí)例可以完成的請求數(shù)量,以及在指定時(shí)間后將添加的新請求數(shù)量,從而獲得該時(shí)間點(diǎn)的負(fù)載。然而,由于成本高昂[9]或準(zhǔn)確度低,尤其是在資源稀缺和需要準(zhǔn)確預(yù)測的過載條件下,使得請求級別的預(yù)測特別困難。
系統(tǒng)級別:與請求級別的預(yù)測不同,系統(tǒng)級別的預(yù)測不嘗試預(yù)測單個(gè)請求的完成時(shí)間。相反,它們估計(jì)在指定時(shí)間后實(shí)例的整體批次計(jì)數(shù)或TBT狀態(tài)。這種類型的預(yù)測是持續(xù)的,并且需要的精度較低,使其更適合過載場景。
在Mooncake中,我們目前使用系統(tǒng)級預(yù)測策略:我們假設(shè)每個(gè)請求的解碼階段需要統(tǒng)一的時(shí)間td。首先,對于給定時(shí)刻t,可以由t時(shí)刻的預(yù)填充實(shí)例完成的請求被添加到統(tǒng)一的解碼實(shí)例中。接下來,在t之前可以完成(即,執(zhí)行時(shí)間超過td)的請求從解碼實(shí)例中移除。最后,計(jì)算所有解碼實(shí)例的平均TBT比率與ltbt的比值來預(yù)測負(fù)載。將請求級預(yù)測的探索留作未來的工作。
7 評估
本節(jié)評估了Mooncake在不同數(shù)據(jù)集和各種工作負(fù)載下的端到端性能。如前所述,為了保護(hù)專有信息并促進(jìn)可復(fù)制性,本文報(bào)告的所有實(shí)驗(yàn)結(jié)果都是基于一個(gè)虛擬模型,該模型遵循與LLaMA2-70B相同的架構(gòu)。
7.1 端到端性能
表1: 在端到端實(shí)驗(yàn)中使用的數(shù)據(jù)集。
| 數(shù)據(jù)集 | 平均輸入長度 | 平均輸出長度 | 緩存比率 | 到達(dá)模式 |
|---|---|---|---|---|
| ArXiv摘要[26] | 8088 | 229 | ~0% | 泊松過程 |
| L-Eval[27] | 19019 | 72 | >80% | 泊松過程 |
| 模擬數(shù)據(jù) | 16k, 32k, 64k, 128k | 512 | 50% | 泊松過程 |
| 真實(shí)數(shù)據(jù) | 7955 | 194 | ~50% | 基于時(shí)間戳 |
7.1.1 公共數(shù)據(jù)集
本節(jié)評估了Mooncake和vLLM在公共數(shù)據(jù)集上的端到端測試性能,使用了ArXiv摘要和L-Eval。我們使用四個(gè)vLLM實(shí)例作為基線,表示為vLLM-[4M]。相比之下,Mooncake有兩種不同的配置:一個(gè)集群由三個(gè)預(yù)填充實(shí)例和一個(gè)解碼實(shí)例組成,標(biāo)記為Mooncake-[3P+1D];另一個(gè)有兩個(gè)預(yù)填充和兩個(gè)解碼實(shí)例,標(biāo)記為Mooncake-[2P+2D]。圖9顯示,在ArXiv摘要和L-Eval數(shù)據(jù)集上,Mooncake-[3P+1D]在滿足SLO的同時(shí),分別比vLLM-[4M]實(shí)現(xiàn)了20%和40%的吞吐量提升。此外,L-Eval數(shù)據(jù)集上Mooncake的吞吐量通過前綴緩存進(jìn)一步提高,顯著減少了預(yù)填充時(shí)間。然而,盡管具有較低的TBT延遲,Mooncake-[2P+2D]在TTFT指標(biāo)上的表現(xiàn)不如Mooncake-[3P+1D]和vLLM-[4M]。這種差異源于預(yù)填充和解碼實(shí)例之間的負(fù)載不平衡。在現(xiàn)實(shí)世界的集群中,預(yù)填充和解碼實(shí)例的需求通常在一定時(shí)期內(nèi)保持穩(wěn)定,只有輕微的暫時(shí)性不平衡。因此,預(yù)填充和解碼實(shí)例的比例可以預(yù)設(shè)。未來的研究將探索更靈活的部署和轉(zhuǎn)換方法。
7.1.2 模擬數(shù)據(jù)
本節(jié)使用模擬數(shù)據(jù)進(jìn)行端到端實(shí)驗(yàn)。集群配置與§7.1.1相同,使用Mooncake配置[3P+1D]、[2P+2D]和vLLM-[4M]。值得注意的是,模擬數(shù)據(jù)中的長上下文請求顯著破壞了vLLM的解碼階段。為了對抗這一點(diǎn),vLLM單獨(dú)處理請求,而不是批量處理。實(shí)驗(yàn)結(jié)果如圖10所示。盡管Mooncake采用批量處理,但其兩階段分離設(shè)計(jì)有效地最小化了預(yù)填充階段對解碼階段的影響,確保它從未違反TBT SLO。Mooncake展示了顯著更高的吞吐量,提升范圍從50%到525%,同時(shí)遵守與vLLM相同的TTFT和TBT SLO約束。
7.1.3 真實(shí)工作負(fù)載
我們進(jìn)一步利用10個(gè)預(yù)填充實(shí)例和10個(gè)解碼實(shí)例,標(biāo)記為Mooncake-[10P+10D],以及20個(gè)vLLM實(shí)例,稱為vLLM-[20M],重放真實(shí)請求跟蹤并在Mooncake和vLLM上進(jìn)行負(fù)載測試。在這個(gè)實(shí)驗(yàn)設(shè)置中,TTFT的上限設(shè)置為30秒,而TBT閾值限制在每個(gè)token 0.1秒。圖11展示了兩個(gè)系統(tǒng)在TTFT和TBT上的累積分布函數(shù)(CDF)圖。Mooncake-[10P+10D]和vLLM-[20M]的TTFT分布幾乎相同,幾乎所有請求都滿足了TTFT SLO。然而,盡管Mooncake-[10P+10D]的所有請求都滿足了TBT SLO,但vLLM-[20M]的請求中只有57%滿足了這一標(biāo)準(zhǔn),一些請求顯示出極高的TBT。在這個(gè)實(shí)驗(yàn)中,Mooncake可以在遵守SLO的同時(shí)處理大約75%的更多請求。
7.2 過載場景下的性能
本節(jié)評估了在過載場景下的性能,重點(diǎn)關(guān)注系統(tǒng)可以處理的最大請求數(shù)量,如§6所述。基線策略在兩個(gè)階段開始之前基于負(fù)載拒絕請求,導(dǎo)致資源浪費(fèi),因?yàn)橐呀?jīng)處理過的預(yù)填充階段的請求被拒絕。與此相反,我們提出了早期拒絕和基于預(yù)測的早期拒絕策略,分別在§6.2和§6.4中詳細(xì)說明。這些策略綜合考慮了系統(tǒng)的負(fù)載,從而減少了不必要的請求拒絕。
具體來說,我們構(gòu)建了一個(gè)具有8個(gè)預(yù)填充實(shí)例和8個(gè)解碼實(shí)例的Mooncake集群,并使用23,000個(gè)請求的真實(shí)跟蹤對其進(jìn)行了測試。為了模擬過載場景,我們將重放速度提高到2倍。
表2顯示了Mooncake在不同策略下的系統(tǒng)拒絕請求的數(shù)量。使用基線策略時(shí),系統(tǒng)拒絕了4183個(gè)請求。相比之下,在早期拒絕和基于預(yù)測的早期拒絕策略下,Mooncake分別拒絕了3771和3589個(gè)請求。這表明,通過早期拒絕請求,Mooncake可以避免不必要的預(yù)填充計(jì)算,從而提高系統(tǒng)資源的有效利用率。此外,通過預(yù)測解碼實(shí)例的負(fù)載,Mooncake可以緩解負(fù)載波動,增加請求處理能力。
| 策略類型 | 拒絕請求的數(shù)量 |
|---|---|
| 基線(Baseline) | 4183 |
| 早期拒絕(Early Rejection) | 3771 |
| 基于預(yù)測的早期拒絕(Early Rejection based on Prediction) | 3589 |
公眾號后臺回復(fù)“數(shù)據(jù)集”獲取100+深度學(xué)習(xí)各方向資源整理
極市干貨
點(diǎn)擊閱讀原文進(jìn)入CV社區(qū)
收獲更多技術(shù)干貨
