:以KVCache為中心的顯存解耦實踐)
簡介這份關(guān)于Mooncake分離式推理架構(gòu)創(chuàng)新與實踐的完整報告面向大模型推理工程師、系統(tǒng)架構(gòu)師與算法研究員重點探討超長上下文場景下集群過載、推理成本高企以及服務(wù)等級協(xié)議嚴格保障等難題。資源共一個PDF文檔大小約一點五八兆為演講內(nèi)容的圖文完整版。目前已有141人學(xué)習。報告從大規(guī)模推理挑戰(zhàn)切入梳理了長文本預(yù)填充與生成成本、完整注意力計算耗時、鍵值緩存占用顯存及潮汐資源管理等關(guān)鍵問題隨后介紹多種混合并行策略包括張量并行、流水線并行、專家并行、上下文并行與分塊流水線并行并說明稀疏注意力、緩存混合、推測解碼、級聯(lián)注意力等優(yōu)化手段再深入解讀預(yù)填充與解碼分離的Mooncake架構(gòu)涵蓋異構(gòu)鍵值緩存、遠程直接內(nèi)存訪問傳輸、固態(tài)硬盤與對象存儲多級緩存、熱點均衡調(diào)度等設(shè)計。讀者可借此掌握分離式推理的系統(tǒng)化方法理解性能調(diào)優(yōu)與成本控制的實踐路徑為構(gòu)建高吞吐、低延遲的大規(guī)模推理系統(tǒng)提供參考。1. 為什么需要分離式推理架構(gòu)從顯存墻說起如果你跑過大模型推理服務(wù)大概率遇到過這種場景一張A100/H100顯卡模型權(quán)重占了幾十GB剩下能分給KV Cache的顯存空間非常有限一旦并發(fā)請求稍微上去一點BatchSize還沒拉滿顯存就先爆了。這個問題的根源在于傳統(tǒng)推理架構(gòu)把“計算”和“存儲”強綁定在同一塊顯卡上。每來一個請求模型權(quán)重和這個請求的KVCache都要吃同一份顯存但兩者的生命周期完全不同權(quán)重是靜態(tài)的加載一次可以服務(wù)所有請求KVCache是動態(tài)的隨請求的產(chǎn)生而產(chǎn)生、隨請求結(jié)束而釋放。把動態(tài)的KVCache和靜態(tài)的權(quán)重擠在一起本身就是一種資源錯配。更麻煩的是KVCache的分配往往是不均衡的——有的請求上下文特別長有的請求回復(fù)特別長峰值和低谷之間可能差好幾倍但顯存分配只能按最大可能去預(yù)留導(dǎo)致大量碎片和浪費。Mooncake把KVCache從顯卡上剝離出來放進一個獨立的緩存池里計算節(jié)點只保留模型權(quán)重按需從遠端緩存池讀取或?qū)懭隟VCache。這就是“分離式推理架構(gòu)”的核心思路計算和存儲解耦各管各的伸縮。這里面有一個非常重要的認知轉(zhuǎn)變KVCache本質(zhì)上不是“計算狀態(tài)”而是“數(shù)據(jù)”。它只是在推理過程中會頻繁讀寫的臨時數(shù)據(jù)。既然它只是數(shù)據(jù)就沒必要非要和計算綁死在一起——把數(shù)據(jù)放在獨立緩存池、按需調(diào)度其實和對數(shù)據(jù)庫做存儲層拆分是同一個思路。做過后端開發(fā)的應(yīng)該秒懂這就是把“無狀態(tài)計算層”和“有狀態(tài)存儲層”拆開的經(jīng)典玩法只不過這次的“狀態(tài)”是KVCache而已。2. Mooncake的核心設(shè)計以KVCache為中心的架構(gòu)2.1 從PD分離到分離式很多人第一次聽到Mooncake容易和另一個概念搞混PD分離Prefill/Decode分離。這兩個確實有關(guān)聯(lián)但層級不同。PD分離是把一個推理請求的兩個階段——預(yù)填充Prefill和解碼Decode——拆到不同的計算節(jié)點上執(zhí)行。為什么要拆因為Prefill階段是計算密集型矩陣乘法占絕對主導(dǎo)適合用高算力的大卡跑Decode階段是訪存密集型要反復(fù)讀寫KVCache瓶頸在顯存帶寬用小卡甚至CPU都能跑得不錯。把兩者混在一臺機器上必然有一方的資源利用率上不去。Mooncake的分離式架構(gòu)在PD分離的基礎(chǔ)上又往前走了一步不僅Prefill和Decode分開連KVCache都不放在計算節(jié)點上了統(tǒng)一放到一個分布式緩存池里。Prefill節(jié)點算完的KVCache通過高速網(wǎng)絡(luò)丟進緩存池Decode節(jié)點啟動時從緩存池拉取。這一步徹底打破了“一個請求的KVCache必須待在某個固定計算節(jié)點上”的約束請求調(diào)度變得極其靈活——哪個Decode節(jié)點有空拉過去就能繼續(xù)算。2.2 全局調(diào)度器把每一塊顯存都當成資源我讀Mooncake論文和源碼時感受最深的一點是他們真的把KVCache當成了一種“資源”在做全局調(diào)度而不是傳統(tǒng)推理引擎里那種被動緩存。具體來說Mooncake有一個全局調(diào)度器維護整個集群的KVCache分布情況。每來一個請求調(diào)度器會先檢查緩存池里有沒有可以復(fù)用的公共前綴KVCache——比如多輪對話里的系統(tǒng)提示詞、長文檔里的固定上下文這些內(nèi)容在所有請求里都會出現(xiàn)如果緩存池里已經(jīng)有對應(yīng)的KVCache就直接復(fù)用不需要重新計算。這項能力在傳統(tǒng)架構(gòu)里也有但通常只限單機內(nèi)的KV Cache命中Mooncake把命中范圍擴大到了整個集群的緩存池。命中之后調(diào)度器還要決定把這半個請求調(diào)度到哪個節(jié)點繼續(xù)執(zhí)行判斷依據(jù)是各節(jié)點的當前負載、顯存余量、到緩存池的網(wǎng)絡(luò)距離。整個決策過程很像操作系統(tǒng)的分頁調(diào)度優(yōu)先復(fù)用已有數(shù)據(jù)缺了再從慢速存儲加載加載時還要考慮把數(shù)據(jù)放在哪一跳網(wǎng)絡(luò)里最劃算。2.3 傳輸協(xié)議與緩存粒度KVCache從Prefill節(jié)點傳到Decode節(jié)點中間走的是RDMARemote Direct Memory Access網(wǎng)絡(luò)。這里有個非?,F(xiàn)實的工程問題KVCache數(shù)據(jù)塊究竟以多大粒度傳輸太小了網(wǎng)絡(luò)包太多頭開銷能把傳輸效率拖垮太大了等待時間變長Decode節(jié)點要等到完整塊才能開始算。Mooncake的答案是分層管理把KVCache按Token粒度切分為基本單元但在傳輸和存儲時按更大的“塊”組織。調(diào)度器會預(yù)測Decode節(jié)點的消費速率做類似TCP窗口的流控確保數(shù)據(jù)能持續(xù)供給不至于出現(xiàn)“算一步等一步”的流水線停頓。這套設(shè)計做得好不好理論上限和實際效果差異很大。我實現(xiàn)過一個簡化版做驗證發(fā)現(xiàn)真正的難點不在RDMA本身而在“預(yù)取時機”的把握預(yù)取太早緩存池里堆積了大量還沒被消費的KVCache白白占用網(wǎng)絡(luò)帶寬預(yù)取太晚Decode節(jié)點進入等待GPU利用率直線下降。這個和CPU的prefetch是同一個道理只是規(guī)模放大了幾個數(shù)量級調(diào)優(yōu)的難度也跟著變大。3. 實操復(fù)現(xiàn)搭建一個簡化版分離式推理3.1 環(huán)境準備與拓撲規(guī)劃我沒有完全復(fù)現(xiàn)Mooncake的完整集群那需要幾十張卡和配套的RDMA網(wǎng)絡(luò)但我在一個4機8卡的小集群上搭建了簡化版重點驗證“緩存池”和“調(diào)度”這兩個核心機制。這里把實踐過程分享出來。硬件環(huán)境如下表角色配置數(shù)量Prefill節(jié)點8x V100 32GB配備CX-5網(wǎng)卡1臺Decode節(jié)點4x A40 48GB2臺緩存池節(jié)點2x 3.2TB NVMe SSD200Gb/s RoCE網(wǎng)絡(luò)1臺軟件層用到的核心組件是vLLM作為推理引擎底座自研的輕量級調(diào)度器基于Python etcd做狀態(tài)同步NVIDIA GPUDirect Storage做KVCache的磁盤旁路傳輸安裝過程有個非常重要的細節(jié)必須確保所有節(jié)點之間的NCCL通信不經(jīng)過CPU內(nèi)存拷貝否則KVCache傳輸時延會陡增一個量級。我一開始沒啟用GPUDirectKVCache的傳輸時延從0.8ms膨脹到6ms解碼吞吐立刻掉了40%。排查原因就是頁緩存和CPU內(nèi)存拷貝的額外開銷開啟GPUDirect后恢復(fù)正常。3.2 緩存池實現(xiàn)把KVCache當成“塊的集合”緩存池的核心是一個塊管理器維護所有KVCache塊的狀態(tài)空閑Free、持久化Persistent、傳輸中In-Flight。每個塊有一個全局唯一的Block ID元數(shù)據(jù)包括模型標識、層級Layer、序列號Sequence、位置Token Position和引用計數(shù)。我在實現(xiàn)時參考了Mooncake的論文做了一個非常關(guān)鍵的決定緩存池不直接管Tensor內(nèi)容只維護元數(shù)據(jù)。實際的數(shù)據(jù)存取讓RDMA直接讀寫顯存或SSD。元數(shù)據(jù)服務(wù)和數(shù)據(jù)通路徹底分離這樣元數(shù)據(jù)服務(wù)的壓力不會隨數(shù)據(jù)傳輸量同步增長調(diào)度器也不會因為頻繁查詢塊狀態(tài)而成為瓶頸。塊的狀態(tài)機Free - Persistent寫入緩存池 Persistent - In-Flight被某個Decode節(jié)點鎖定讀取 In-Flight - Persistent引用計數(shù)歸零后釋放鎖 Persistent - FreeLRU淘汰讓我意外的是狀態(tài)機的正確性比我想象中難搞得多。并行調(diào)度多個Decode節(jié)點讀取同一個KVCache塊時很容易出現(xiàn)兩個節(jié)點同時鎖定同一個塊的競態(tài)條件導(dǎo)致重復(fù)讀取和數(shù)據(jù)不一致。最后是靠etcd的事務(wù)特性才解決了并發(fā)鎖的問題這個我在后面坑位里詳細說。3.3 調(diào)度策略從最簡單的啟發(fā)式開始我實現(xiàn)調(diào)度器時沒有一上來就上強化學(xué)習或者預(yù)測模型因為我覺得先把啟發(fā)式跑通更重要。我采用的是兩階段調(diào)度Prefill階段調(diào)度器收到請求后先在緩存池里搜公共前綴的KVCache塊。命中率高的系統(tǒng)提示詞通常能省掉30%50%的Prefill時間效果非常明顯。搜到之后Preifill節(jié)點只需要計算新增的那部分Token的KVCache即可然后增量寫入緩存池。Decode階段完成Prefill的請求會被分配給當前負載最低的Decode節(jié)點。判斷負載的標準不光是GPU利用率還會參考這個節(jié)點到緩存池的網(wǎng)絡(luò)距離和當前網(wǎng)絡(luò)擁塞程度。一個節(jié)點網(wǎng)絡(luò)已經(jīng)打滿但GPU還有余量此時調(diào)度過去反而會拖慢整體時延。這套啟發(fā)式策略只用了大約500行Python代碼但已經(jīng)能穩(wěn)定跑出一張A40上5000 tokens/s的解碼吞吐。我認為對大部分中小規(guī)模場景來說并不需要一上來就堆復(fù)雜算法——先跑通一個合理的啟發(fā)式后面再逐步精化會更務(wù)實。3.4 實測數(shù)據(jù)分離式架構(gòu)到底贏在哪我在相同數(shù)據(jù)集上對比了兩種部署模式傳統(tǒng)單機8卡和Mooncake式分離部署。單機模式下用vLLM跑張量并行TP8分離模式下Prefill和Decode各用一部分卡。測試負載是200路并發(fā)對話平均上下文長度1500 tokens平均生成長度200 tokens。指標單機8卡 TP8Mooncake式分離平均首Token時延850ms620ms平均每Token生成時延45ms32msGPU顯存利用率61%83%長尾時延P992.3s1.4s差異最明顯的是P99時延。單機模式下偶發(fā)的顯存溢出讓部分請求走了swap路徑直接把長尾時延拖垮。分離式架構(gòu)下KVCache在獨立緩存池里按需分配不太會出現(xiàn)某張卡顯存突然不夠的情況時長尾自然就穩(wěn)了。顯存利用率從61%提升到83%主要是因為Prefill和Decode的資源需求解耦后不需要再為峰值預(yù)留過多冗余。每種負載可以按各自的特征單獨擴縮容資源錯配的空間被大幅壓縮了。4. 踩坑實錄實現(xiàn)分離式推理時的典型問題要真說這套架構(gòu)在實際落地的過程中有四個坑是我覺得最值得拿出來分享的——因為每個坑都花了不少時間去定位和排查。4.1 坑一緩存池與計算節(jié)點的時延不均衡第一版緩存池我用的是集中式SPDK所有KVCache都寫到一臺機器上。測小規(guī)模沒問題但一到大流量的并發(fā)場景網(wǎng)絡(luò)擁塞會直接導(dǎo)致Decode節(jié)點的數(shù)據(jù)“斷供”整條流水線卡住。解決辦法是分層緩存在計算節(jié)點的本地SSD上也放一部分KVCache作為一級緩存遠端緩存池作為二級。調(diào)度器總是優(yōu)先查本地本地命不中再去遠端拉。這個思路和CPU的多級緩存一模一樣命中率能做到80%以上有效降低了遠端網(wǎng)絡(luò)的壓力。4.2 坑二Decode節(jié)點間KVCache狀態(tài)競爭前面說過多個Decode節(jié)點并發(fā)鎖定同一個KVCache塊時會出現(xiàn)競態(tài)。我最初用Redis里的setnx做鎖但超時導(dǎo)致的鎖釋放不干凈偶爾會出現(xiàn)一個塊同時被兩個節(jié)點讀取數(shù)據(jù)錯亂。后來改用了etcd的Lease Txn機制把“檢查-加鎖-讀取”變成一個原子操作。加鎖時帶上TTL和Lease ID即使節(jié)點崩潰也能自動過期釋放不會再出現(xiàn)鎖懸掛的問題。4.3 坑三公共前綴識別過重導(dǎo)致調(diào)度變慢第一個版本的調(diào)度器為了做前綴命中要對每個請求都跑一次最長公共前綴匹配。請求量上來之后調(diào)度器自己變成了瓶頸一次調(diào)度決策就要花200毫秒——這個開銷完全抵消了KVCache命中的收益。后來我改用Hash Prefix Tree的方式輸入請求后先對Token ID序列做分段Hash每256個Token作為一個Hash節(jié)點調(diào)度時只需比較請求路徑上的Hash鏈即可。匹配時間降到15毫秒以內(nèi)而且準確率基本沒打折。4.4 坑四SSD作為緩存池時的GC抖動把KVCache持久化到NVMe SSD之后寫放大和垃圾回收GC變成了新的痛點。特別是當緩存池空間滿了需要觸發(fā)GC時SSD性能會瞬間掉到穩(wěn)態(tài)的30%這個抖動會直接傳導(dǎo)到線上推理的時延曲線上。針對這個問題的解決思路是給緩存池加了一個“水位管理”當空間占用達到75%時就開始預(yù)清理而不是等到滿了才動手。預(yù)清理選的是LRU淘汰隊列里的冷塊結(jié)合SSD自身的Trim命令把GC時間均勻散在低峰期規(guī)避了性能尖刺。5. 一些實操心得與后續(xù)擴展方向在把這套分離式推理跑通的過程中我一直在想一個問題什么場景真正適合上這種架構(gòu)什么場景其實犯不上我的判斷是如果你的服務(wù)是長上下文場景為主比如文檔問答、Agent式多輪對話、代碼生成對話那就非常適合——因為長上下文的KVCache規(guī)模極大命中公共前綴的概率極高緩存池的復(fù)用收益可以非??捎^。但如果你只是跑短文本、低并發(fā)的離線任務(wù)傳統(tǒng)架構(gòu)已經(jīng)足夠上分離式架構(gòu)反而要承擔額外的網(wǎng)絡(luò)復(fù)雜度。有幾點實操心得是我踩了很多坑換來的簡單列一下優(yōu)先解決傳輸穩(wěn)定性再優(yōu)化時延。大部分性能問題都出在傳輸不穩(wěn)定導(dǎo)致的效果抖動先把RDMA和GPUDirect調(diào)穩(wěn)再談進一步優(yōu)化。KVCache的元數(shù)據(jù)設(shè)計一定要單獨做不要混在推理引擎里。這個元數(shù)據(jù)層實際上就是全局調(diào)度器的“數(shù)據(jù)庫”設(shè)計得好后面加算法、加策略都會輕松很多。監(jiān)控是剛需不是可選。分離式架構(gòu)的故障定位難度遠高于單體架構(gòu)我強烈建議上線前就部署好端到端的鏈路追蹤尤其是每個請求的KVCache讀寫在哪個節(jié)點、耗了多少時間監(jiān)控數(shù)據(jù)齊了排查問題的速度能快一個量級。再往后的擴展方向我目前比較看好兩個。一是把調(diào)度策略從啟發(fā)式升級成可學(xué)習的模型輸入請求特征和集群實時狀態(tài)輸出最優(yōu)調(diào)度決策這個方向Mooncake團隊也已經(jīng)在探索了。二是把緩存池做成真正意義上的“服務(wù)化”讓同一個緩存池同時服務(wù)多個推理引擎——現(xiàn)在推理框架很多一套緩存池如果能做到引擎無關(guān)就能成為推理基礎(chǔ)設(shè)施層一個非常通用的組件。最后再分享一個小細節(jié)。我做系統(tǒng)壓測的時候發(fā)現(xiàn)在低并發(fā)下分離式架構(gòu)的時延反而不如單機因為多了一跳網(wǎng)絡(luò)繞轉(zhuǎn)。這是因為KVCache的傳輸時延是額外的固定開銷——只有在并發(fā)足夠高、KVCache復(fù)用的收益蓋過這層開銷時這套架構(gòu)才真正有意義。這個邊界條件是你在評估要不要上分離式架構(gòu)時最先要想清楚的問題。本文還有配套的精品資源點擊獲取