輝達桌面超算，邪修玩法來了！兩台DGX Spark串聯一台蘋果Mac Studio，就能讓大模型推理速度提升至2.77倍。這是GitHub三萬星大模型框架作者EXO Lab團隊發佈的最新成果。這個EXO Labs，專門研究把大模型放到各種家用裝置上運行。之前讓MacBook、iPad、手機等裝置組成叢集跑405B大模型的分佈式推理框架，還有兩台Mac跑滿血DeepSeek-R1，都是這家工作室的手筆。這次他們又把DGX Spark和M3 Ultra結合，利用它們各自的優勢，在大模型部署上整出了新活。那麼，這套邪修組合具體是如何實現的呢？PD分離+流式傳輸，讓裝置各司其職要想理解如何結合兩種裝置的優勢，需要先瞭解大模型推理的工作方式。大模型的推理，主要分為Prefill和Decode兩個階段。Prefill處理提示並為每個Transformer層建構一個KV快取，Decode階段則是根據建構好的整個KV快取生成token。兩個階段任務不同，對硬體性能的側重也不一樣，整體上來說Prefill更吃算力，而Decode吃記憶體頻寬。具體來說，Prefill階段計算量隨提示長度呈二次增長，利用Flash Attention等技術，可以最佳化為線性增長，但計算量依然龐大，因此主要受制於計算能力；到了Decode階段，KV快取已經計算完畢，不需要再重新運算，矩陣-矩陣乘法變成了運算量更低的向量-矩陣乘法比，對算力的需求降低，主要受制於記憶體頻寬。再看EXO Labs手裡的兩種裝置，DGX Spark算力強但是頻寬不行，Mac Studio搭載的M3 Ultra則剛好相反，記憶體頻寬高但算力不如DGX Spark。具體來說，DGX Spark有100TFLOPS的fp16算力，M3 Ultra只有26TFLOPS；而M3 Ultra有256GB@819GB/s的記憶體，DGX Spark卻只有128GB@273GB/s。所以，EXO Labs的思路就是把Prefill和Decode階段分開，分別分配給擅長的裝置，DGX Spark負責Prefill，Mac則負責Decode，這也就是AI Infra業界常說的PD分離。最簡單的PD方式就是先把Prefill做完，然後再傳輸給Decode裝置進行Decode。但這就增加了兩個階段之間的通訊成本，如果傳輸時間過長，效果可能適得其反。所以，進行PD分離運算需要解決的關鍵問題是，就是KV快取傳輸。這裡EXO Labs運用了流式傳輸的思想。我們在網上看電影、刷B站時，並不需要把整個視訊檔案載入完才能開始播放，而是將一小段載入到記憶體之後就可以觀看，後面的內容邊看邊載入，這就是流式傳輸。音視訊可以邊傳邊看，KV快取也可以邊算邊傳，因為大模型當中包含了多個Transformer層，使得KV快取不一定非要以一個Blob的形式到達Decode裝置，而是可以逐層到達。第1層的Prefill完成後，其KV快取就開始傳輸到給M3 Ultra去Decode，同時第2層的Prefill則在DGX Spark上開始，每一層的通訊都與後續層的計算重疊。實際上， EXO還會在處理層的過程中傳輸該層的KV向量，因為KV向量的計算是在最繁重的計算步驟之前進行的。利用EXO框架，PD分離、逐層KV流以及硬體感知都可以自動完成。啟動EXO時，它會自動發現連接的所有裝置，並針對計算吞吐量、記憶體頻寬、記憶體容量和網路特性對每個裝置進行分析。給定一個模型和拓撲結構， EXO就會規劃那個裝置應該處理Prefill，那個裝置應該處理Decode，是否需要跨層流水線，何時傳輸KV對，以及如何在網路條件發生變化時進行調整。最終，在DGX Spark和Mac Studio的組合下，Llama-3.1 8B在Prefill階段的速度提升至了Mac的3.79倍，Decode速度提升至DGX Spark的3.37倍，整體提升至Mac Studio的2.77倍。Three More ThingsEXO這種PD分離的做法，輝達自己也在進行嘗試，其即將推出的Rubin CPX平台將使用計算密集型Rubin CPX處理器進行Prefill，配備巨大HBM3e記憶體頻寬的標準Rubin晶片則負責Decode。再說EXO團隊這次用的DGX Spark，最近正在進行配送，馬斯克、奧特曼還有LeCun都收到了，其中還有老黃親自送貨上門。另外蘋果最新發佈的M5，AI性能也有一定提升。在M5晶片的MacBook Pro上，首個Token生成速度（主要受Prefill影響）提升到了M1的6.4倍、M4的3.55倍。另外，也有更快的圖像/視訊生成、更快的微調，以及更高的吞吐量。不過宣傳上說的是性能的又一次躍升，但仔細一看，M5甚至不如M4 Max，M4 Max又甚至不如M3 Ultra……而EXO的這波操作下，M3 Ultra的含金量似乎更高了。 (量子位)