英偉達(dá)桌面超算，邪修玩法來(lái)了！

兩臺(tái)DGX Spark串聯(lián)一臺(tái)蘋(píng)果Mac Studio，就能讓大模型推理速度提升至2.77倍。

這是GitHub三萬(wàn)星大模型框架作者EXO Lab團(tuán)隊(duì)發(fā)布的最新成果。

這個(gè)EXO Labs，專門(mén)研究把大模型放到各種家用設(shè)備上運(yùn)行。

之前讓MacBook、iPad、手機(jī)等設(shè)備組成集群跑405B大模型的分布式推理框架，還有兩臺(tái)Mac跑滿血DeepSeek-R1，都是這家工作室的手筆。

這次他們又把DGX Spark和M3 Ultra結(jié)合，利用它們各自的優(yōu)勢(shì)，在大模型部署上整出了新活。

那么，這套邪修組合具體是如何實(shí)現(xiàn)的呢？

PD分離+流式傳輸，讓設(shè)備各司其職

要想理解如何結(jié)合兩種設(shè)備的優(yōu)勢(shì)，需要先了解大模型推理的工作方式。

大模型的推理，主要分為Prefill和Decode兩個(gè)階段。

Prefill處理提示并為每個(gè)Transformer層構(gòu)建一個(gè)KV緩存，Decode階段則是根據(jù)構(gòu)建好的整個(gè)KV緩存生成token。

兩個(gè)階段任務(wù)不同，對(duì)硬件性能的側(cè)重也不一樣，整體上來(lái)說(shuō)Prefill更吃算力，而Decode吃內(nèi)存帶寬。

具體來(lái)說(shuō)，Prefill階段計(jì)算量隨提示長(zhǎng)度呈二次增長(zhǎng)，利用Flash Attention等技術(shù)，可以優(yōu)化為線性增長(zhǎng)，但計(jì)算量依然龐大，因此主要受制于計(jì)算能力；

到了Decode階段，KV緩存已經(jīng)計(jì)算完畢，不需要再重新運(yùn)算，矩陣-矩陣乘法變成了運(yùn)算量更低的向量-矩陣乘法比，對(duì)算力的需求降低，主要受制于內(nèi)存帶寬。

再看EXO Labs手里的兩種設(shè)備，DGX Spark算力強(qiáng)但是帶寬不行，Mac Studio搭載的M3 Ultra則剛好相反，內(nèi)存帶寬高但算力不如DGX Spark。

具體來(lái)說(shuō)，DGX Spark有100TFLOPS的fp16算力，M3 Ultra只有26TFLOPS；而M3 Ultra有256GB@819GB/s的內(nèi)存，DGX Spark卻只有128GB@273GB/s。

所以，EXO Labs的思路就是把Prefill和Decode階段分開(kāi)，分別分配給擅長(zhǎng)的設(shè)備，DGX Spark負(fù)責(zé)Prefill，Mac則負(fù)責(zé)Decode，這也就是AI Infra業(yè)界常說(shuō)的PD分離。

最簡(jiǎn)單的PD方式就是先把Prefill做完，然后再傳輸給Decode設(shè)備進(jìn)行Decode。

但這就增加了兩個(gè)階段之間的通信成本，如果傳輸時(shí)間過(guò)長(zhǎng)，效果可能適得其反。

所以，進(jìn)行PD分離運(yùn)算需要解決的關(guān)鍵問(wèn)題是，就是KV緩存?zhèn)鬏敗?/span>

這里EXO Labs運(yùn)用了流式傳輸?shù)乃枷搿?/span>

我們?cè)诰W(wǎng)上看電影、刷B站時(shí)，并不需要把整個(gè)視頻文件加載完才能開(kāi)始播放，而是將一小段加載到內(nèi)存之后就可以觀看，后面的內(nèi)容邊看邊加載，這就是流式傳輸。

音視頻可以邊傳邊看，KV緩存也可以邊算邊傳，因?yàn)榇竽Ｐ彤?dāng)中包含了多個(gè)Transformer層，使得KV緩存不一定非要以一個(gè)Blob的形式到達(dá)Decode設(shè)備，而是可以逐層到達(dá)。

第1層的Prefill完成后，其KV緩存就開(kāi)始傳輸?shù)浇oM3 Ultra去Decode，同時(shí)第2層的Prefill則在DGX Spark上開(kāi)始，每一層的通信都與后續(xù)層的計(jì)算重疊。

實(shí)際上， EXO還會(huì)在處理層的過(guò)程中傳輸該層的KV向量，因?yàn)镵V向量的計(jì)算是在最繁重的計(jì)算步驟之前進(jìn)行的。

利用EXO框架，PD分離、逐層KV流以及硬件感知都可以自動(dòng)完成。

啟動(dòng)EXO時(shí)，它會(huì)自動(dòng)發(fā)現(xiàn)連接的所有設(shè)備，并針對(duì)計(jì)算吞吐量、內(nèi)存帶寬、內(nèi)存容量和網(wǎng)絡(luò)特性對(duì)每個(gè)設(shè)備進(jìn)行分析。

給定一個(gè)模型和拓?fù)浣Y(jié)構(gòu)， EXO就會(huì)規(guī)劃哪個(gè)設(shè)備應(yīng)該處理Prefill，哪個(gè)設(shè)備應(yīng)該處理Decode，是否需要跨層流水線，何時(shí)傳輸KV對(duì)，以及如何在網(wǎng)絡(luò)條件發(fā)生變化時(shí)進(jìn)行調(diào)整。

最終，在DGX Spark和Mac Studio的組合下，Llama-3.1 8B在Prefill階段的速度提升至了Mac的3.79倍，Decode速度提升至DGX Spark的3.37倍，整體提升至Mac Studio的2.77倍。

Three More Things

EXO這種PD分離的做法，英偉達(dá)自己也在進(jìn)行嘗試，其即將推出的Rubin CPX平臺(tái)將使用計(jì)算密集型Rubin CPX處理器進(jìn)行Prefill，配備巨大HBM3e內(nèi)存帶寬的標(biāo)準(zhǔn)Rubin芯片則負(fù)責(zé)Decode。

再說(shuō)EXO團(tuán)隊(duì)這次用的DGX Spark，最近正在進(jìn)行配送，馬斯克、奧特曼還有LeCun都收到了，其中還有老黃親自送貨上門(mén)。

另外蘋(píng)果最新發(fā)布的M5，AI性能也有一定提升。

在M5芯片的MacBook Pro上，首個(gè)Token生成速度（主要受Prefill影響）提升到了M1的6.4倍、M4的3.55倍。

另外，也有更快的圖像/視頻生成、更快的微調(diào)，以及更高的吞吐量。

不過(guò)宣傳上說(shuō)的是性能的又一次躍升，但仔細(xì)一看，M5甚至不如M4 Max，M4 Max又甚至不如M3 Ultra……

而EXO的這波操作下，M3 Ultra的含金量似乎更高了。