英伟达桌面超算,邪修玩法来了!
两台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 Things
EXO这种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的含金量似乎更高了。
参考链接:
[1]https://blog.exolabs.net/nvidia-dgx-spark/
[2]https://www.tomshardware.com/software/two-nvidia-dgx-spark-systems-combined-with-m3-ultra-mac-studio-to-create-blistering-llm-system-exo-labs-demonstrates-disaggregated-ai-inference-and-achieves-a-2-8-benchmark-boost
[3]https://x.com/awnihannun/status/1978465715121250801
文章来自于微信公众号“量子位”,作者是“克雷西”。