總體架構

Mooncake 的架構是非常典型的分離式架構，將單個同構 GPU 集群(qun)的(de)資(zi)源(yuan)打(da)散并重新組(zu)織成三個可(ke)以獨立彈性伸縮的(de)資(zi)源(yuan)池(chi)。其(qi)中 Prefill Pool 處理用戶輸(shu)入，主(zhu)要對 Time To First Token (TTFT) 負(fu)責。同時因為(wei) Prefill 相對計算密集，這(zhe)一部(bu)(bu)分也(ye)承擔著抬高整(zheng)體(ti)資(zi)源(yuan)利(li)用率的(de)任(ren)務。Prefill 處理完之(zhi)后對應的(de) KVCache 會被送到 Decode Pool 進行(xing) autoregression 式的(de)流式輸(shu)出(chu)。雖然(ran)我們(men)希望盡可(ke)能攢大的(de) batch 以提升 MFU，但這(zhe)一部(bu)(bu)分主(zhu)要需要對 Time Between Tokens (TBT) 負(fu)責。

這(zhe)里我們使(shi)用(yong)(yong) TBT 而(er)非另外一(yi)些工作(zuo)中常用(yong)(yong)的(de)(de)(de)(de)(de) Time Per Output Token (TPOT) 的(de)(de)(de)(de)(de)原因在于雖然(ran)其(qi)實本來(lai)這(zhe)兩(liang)者應(ying)該等價，但實際計(ji)算 TPOT 的(de)(de)(de)(de)(de)時(shi)候經常會(hui)被簡單(dan)地等價成計(ji)算 average TPOT。然(ran)后就直(zhi)接用(yong)(yong)了吞吐的(de)(de)(de)(de)(de)倒(dao)數(shu)或(huo)者說總(zong)生成時(shi)間(jian)除以(yi)(yi)總(zong)生成 token 數(shu)。TBT 的(de)(de)(de)(de)(de)定義則(ze)更(geng)加明確的(de)(de)(de)(de)(de)指定為兩(liang)個 token 生成間(jian)的(de)(de)(de)(de)(de)延(yan)遲，因此(ci)基于 TBT 上限(xian)設計(ji) Service Level Objective (SLO) 能更(geng)好(hao)的(de)(de)(de)(de)(de)反映流式交互時(shi)的(de)(de)(de)(de)(de)用(yong)(yong)戶體(ti)驗。直(zhi)觀的(de)(de)(de)(de)(de)理(li)解的(de)(de)(de)(de)(de)話(hua)大家可以(yi)(yi)想象如果 GPT-4o 的(de)(de)(de)(de)(de)語音回復是說一(yi)句(ju)停頓一(yi)段時(shi)間(jian)再說一(yi)句(ju)的(de)(de)(de)(de)(de)說話(hua)方(fang)式的(de)(de)(de)(de)(de)話(hua)，即便 average TPOT 能小一(yi)些體(ti)驗還是不夠(gou)好(hao)。

在 Prefill/Decode 之(zhi)(zhi)外我(wo)們(men)還(huan)利用(yong)(yong)每臺(tai) HGX 機器上組成了(le)一(yi)個 KVCache Pool 來進行全局(ju)的(de) Prefix Cache。相比 vLLM 當前單機的(de) Prefix Cache 通過全局(ju)的(de)調(diao)度(du)能夠(gou)大幅(fu)度(du)提升(sheng)復用(yong)(yong)率從而提升(sheng)總吞吐。由此(ci)(ci)帶(dai)來了(le)一(yi)系列如何調(diao)度(du)，分配，復制 KVCache 的(de)問題，而且必須和(he) Prefill/Decode 的(de)調(diao)度(du) co-design，因此(ci)(ci)我(wo)們(men)把 Mooncake 的(de)架構稱之(zhi)(zhi)為(wei)以 KVCache 為(wei)中心(xin)。

在此基礎上接下來(lai)我們分別(bie)討論 Prefill/KVCache/Decode Pool 的關(guan)鍵 design choice。

Prefill Pool

論文這一章我(wo)原本起(qi)了(le)個“Prefill: To Separate or Not. Is it a Question?” 的(de)(de)名字。主(zhu)(zhu)要原因在于(yu)由于(yu) Chunked Prefill 的(de)(de)提出(chu)實際(ji)上 Prefill 和 Decode 節(jie)點(dian)的(de)(de)邊界已經(jing)模糊了(le)起(qi)來。不過除非是(shi)特(te)別短的(de)(de) prompt （可以直接一次性加(jia)入到(dao) decode 的(de)(de) continues batch 里提升(sheng) MFU 但(dan)不違反 TBT SLO 的(de)(de)那種）不然我(wo)們還是(shi)傾向于(yu)使用獨立的(de)(de) Prefill 節(jie)點(dian)。原因主(zhu)(zhu)要是(shi)兩方面。

首先我們定義了一(yi)個 VRAM occupation cost 的(de)(de)概念，即(ji) KVCache 大(da)小乘以其在顯存(cun)中留駐的(de)(de)時間。使(shi)用(yong) chunked prefill 的(de)(de)話由于每個 chunk 還需要(yao)和(he)其他 decode request batch 在一(yi)起處理因此是(shi)顯著(zhu)增加了對應(ying) KVCache 在顯存(cun)中留駐的(de)(de)時間的(de)(de)，因此等效的(de)(de)降低了寶貴顯存(cun)空間的(de)(de)利用(yong)率。當然顯存(cun)空出(chu)來(lai)了有沒有地方用(yong)就是(shi)另(ling)外一(yi)個問題了。這(zhe)里我們在論文 4.2 節(jie)中主要(yao)討論了和(he)未來(lai) batch job 結合的(de)(de)可能性。

另外更(geng)重要的一個原因(yin)則主要是在于(yu)長文(wen)本場景下我們需要在 Prefill 集(ji)群(qun)中使用(yong)特殊的多節點分(fen)布(bu)式劃(hua)分(fen)方(fang)法來壓低 TTFT。這(zhe)里我們討論了相比現在被(bei)討論的更(geng)多的 SP 策(ce)略(lve)直(zhi)接使用(yong) Chunked PP 策(ce)略(lve)在調度和減少(shao)通訊量方(fang)面的優越性，前者大(da)幅度簡化架構，后(hou)者則能夠省出寶貴的傳輸帶寬來搬運(yun) KVCache。具體策(ce)略(lve)其實就是 TeraPipe 的推(tui)理簡化版本，因(yin)為(wei)只(zhi)有 forward 沒有 backword 的話不需要啥(sha)動態(tai)規(gui)劃(hua)來平衡每(mei)(mei)個 Chunk 的計(ji)算量，只(zhi)要給(gei)每(mei)(mei)個 Chunk 設置一個最小(xiao)的值就行。

KVCache Pool

都 2024.6 了單(dan)節點的(de)(de) Prefix Cache 已經不是(shi)什(shen)么新鮮東西(xi)了。這里(li)(li)主要討(tao)論一(yi)下全局調度能夠成立的(de)(de)底層原因(yin)：對于每(mei) X byte 的(de)(de) KVCache，其(qi)生成所需(xu)的(de)(de)算力正比于 X*hd 再乘以(yi)一(yi)個較大的(de)(de)常(chang)數，這里(li)(li) hd 對應(ying)模型的(de)(de) hidden dimension。因(yin)此只要每(mei)卡(ka)算力比如 A100 的(de)(de) 220TFLOPS 和每(mei)卡(ka)通(tong)訊帶寬比如 CX7 的(de)(de) 100Gbps 的(de)(de)比值小于 hd 乘以(yi)這個常(chang)數，那(nei)么從遠端傳(chuan)輸 KVCache 相比原地重算不僅僅減少了計算量還減少了 TTFT。又省錢(qian)用戶體驗(yan)還更好了自然是(shi)何(he)樂(le)而不為(wei)。

但是上述公式中具體的 hd，常數是和模型結構強相關的所以不是說一定可以，只能說原則上比較大的模型應該都行。另外雖然網卡有 100Gbps，但是如果因為比如 KVCache 下沉到慢速存儲上去了，網絡擁擠了之類的原因跑不到這么高那么自然也就不一定成立。所以還是需要以 KVCache 為中心針對不同請求的 TTFT SLO 分別去調度。
在論文的(de)(de)(de) 5.2 節中我們(men)主(zhu)要(yao)是(shi)(shi)提供了(le)一種(zhong)簡(jian)單(dan)的(de)(de)(de)基于 heuristic 的(de)(de)(de)熱點(dian)識別(bie)和復制(zhi)方(fang)法。倒不(bu)是(shi)(shi)說有多么精妙，但重點(dian)是(shi)(shi)簡(jian)單(dan)好(hao)實現，不(bu)用(yong)特別(bie)去預(yu)測未(wei)來的(de)(de)(de) KVCache 使用(yong)模(mo)式(shi)，而且在實際的(de)(de)(de)使用(yong)中效(xiao)果不(bu)錯的(de)(de)(de)方(fang)案。未(wei)來會嘗試更多的(de)(de)(de)調度算(suan)法不(bu)過這個可以(yi)作為一個 strong baseline 來使用(yong)。

順帶這里感慨(kai)一句，目前(qian)標準的 8 卡 HGX 服務器每臺(tai)會配 9~10 張 100Gbps 乃至 200Gbps 的 RDMA 網(wang)卡，至少 2TB 的內存和一大堆的高速 SSD。假設一千(qian)臺(tai)的話就是(shi)近萬張網(wang)卡互連起來的數 PB 共享內存池。做(zuo) Disaggregated Memory 挺久了第一次見這樣的富(fu)裕仗(zhang)。可惜(xi)作為一個 bandwidth bound 場景在架構上(shang)能設計的東西沒那么多(duo)特殊的技巧可用，更多(duo)還是(shi)工程實(shi)現。

Decode Pool

Decode Pool 本身(shen)業界的(de)實踐是(shi)(shi)最(zui)多的(de)。特別是(shi)(shi)開源有(you)極快(kuai)速地在迭代的(de) vLLM，該有(you)的(de)不是(shi)(shi)已經有(you)了也是(shi)(shi)在 roadmap 里。