DEV Community

JH5
JH5

Posted on

DiffusionGemma 26B 登陸 M2 Max:MLX 吞吐量實測與 Context 極限挑戰

為了找到一些在地端也能讓 Agent 有無限 token 自由的毒駕的方法,原本用手邊的M4 24GB Mac 上嘗試執行 DiffusionGemma 26B,卻悲慘的連 1,000 tokens 的 Context 都撐不住,直接迎來 OOM(記憶體不足)的悲劇。

換到 M2 Max 96GB 後,終於可以展現出它應有的實力? 我改用MLX(mlx-vlm 0.6.3),過程中雖然踩了 MXFP4 的量化 Bug 並手動處理了 Patch,但最後成功在4-bit 格式下跑完整套 Benchmark。

本文記錄這幾天 DiffusionGemma 26B 在 Apple Silicon 上的吞吐量極限、Prompt 載入成本、以及 Context 長度與對記憶體的代價,同時,我們也會拿這些實測數據來作為後續 GH200 與 GB10 跨平台效能對比的 Baseline 系列首篇文章。

兩個量化版本

第一次 deploy 踩到 MXFP4 的 dequantize bug,後來換 4-bit 才穩定跑完整個 benchmark:

項目 MXFP4(初版) standard 4-bit(最終)
硬體 Apple M2 Max,96 GB 統一記憶體(38-core GPU)
框架 mlx-vlm 0.6.3 + mlx 0.31.2
模型 mlx-community/...-mxfp4 mlx-community/...-4bit
量化格式 MXFP4(4-bit group) 標準 4-bit
峰值記憶體 ~19 GB ~45.7 GB
部署方式 Python API + OpenAI-compatible server(mlx_vlm.server)

mlx-vlm 是 MLX 社群專門給 VLM 用的推理框架,DiffusionGemma 的 block diffusion decoder 也在它的支援範圍內,參考一些大神的文章,就決定是它了XD

兩個量化版本的取捨

第一次踩坑的版本用的是 MXFP4(mlx-community/diffusiongemma-26B-A4B-it-mxfp4),雖然載入好像成功但第一次 generation 馬上噴錯:

ValueError: [dequantize] Biases must be provided for affine quantization
Enter fullscreen mode Exit fullscreen mode

mlx-vlm 的 _diffusion_soft_embedding_weight 在 dequantize embed_tokens 時預設用 affine mode,但 DiffusionGemma 的 MXFP4 格式沒有 bias 參數,目前試出來的解法是 detect 到 biases is None 時改用 mode="mxfp4"

Patch 完 MXFP4 就能跑了,一開始還想說短 context(~8K)下 peak 只有 19 GB,怎麼這麽省記憶體XD 不過後來就發現,context 一超過 8K 速度就線性往下掉,16K 時幾乎動不了。

所以後來改測 standard 4-bit(mlx-community/diffusiongemma-26B-A4B-it-4bit),雖然是記憶體從 19 GB 跳到 45.7 GB,但短context速度快了一倍以上,穩定性也好很多。

項目 MXFP4 standard 4-bit
模型大小 14.8 GB 16.18 GB
Peak 記憶體 ~19 GB ~45.7 GB
Short gen 峰值 14.7 tok/s 31.6 tok/s
Context 1K ~13 tok/s 0.61 tok/s
相容性 需手動 patch 直接可用
結論 記憶體省但慢 快但不適合長 context

MXFP4 省記憶體、長 context 比較穩,而 standard 4-bit 的生成速度快一倍但記憶體吃好吃滿,最後跟其他平台的 baseline 我是以 standard 4-bit 為主來比較速度,如果你真的需要較長的 context,可以考慮換回 MXFP4。

Prompt Encoding

Prompt encoding 的速度曲線很有趣:

Prompt 長度 Encoding 速度
14 tokens 198 tok/s
269 tokens 459 tok/s
525 tokens 646 tok/s
1,037 tokens 687 tok/s
2,061 tokens 694 tok/s
4,109 tokens 646 tok/s

短 prompt 的 encoding 很慢(198 tok/s),但過了 500 tokens 以後穩定在 650-700 tok/s 左右,這應該是因為 MLX 在短序列的時候沒辦法充分利用 Metal GPU 的平行機制,overhead 相對就比較明顯, 前 1K tokens 的冷啟動成本對實際使用來說沒什麼影響,反正 encoding phase 本來就比 generation 快兩個數量級。

Generation Throughput

Standard 4-bit 版的生成速度跟 MXFP4 版的差異很明顯

Output 長度 生成速度 延遲
32 tokens 7.1 tok/s 4.5s
64 tokens 15.5 tok/s 4.1s
128 tokens 25.8 tok/s 5.0s
256 tokens 31.6 tok/s 8.1s
512 tokens 29.1 tok/s 17.6s
1024 tokens 26.8 tok/s 38.2s

峰值在 256 tokens(31.6 tok/s),剛好 fit 一個 diffusion canvas。比 MXFP4 版的 14.7 tok/s 快了 115%。512 tokens 需要跨 canvas,降回 29.1 tok/s。

如果你想要更高吞吐量,可以試 max_denoising_steps=16(預設 48),品質會降但速度翻倍。

Context Length

Standard 4-bit 雖然有優點,不過也產生了一些悲劇,記憶體消耗跳到 45.7 GB,造成 KV cache 的空間反而比 MXFP4 少了不少

Context 長度 生成速度 延遲
~1.8K tokens 0.61 tok/s 52.1s
~9.3K tokens 1.38 tok/s 23.3s
~18.6K tokens 0.57 tok/s 56.5s

從數據上看起來花的滿多時間的,但是這些數字不是生成慢,主要是 prompt encoding 就吃掉了大部分時間,這顆將近45 GB 的模型佔用讓 KV cache 只能從剩下的 50 GB 擠,但 mlx_vlm.server 的 memory management 似乎沒有針對這種大模型做最佳化(可能快來了?),也造成較長 prompt 的 encoding phase 幾乎是線性時間的飆升。

併發:standard 4-bit 的 scaling

Standard 4-bit 版在併發測試的表現比 MXFP4 好一些,這部分我跟 MAC 底層不熟XD,只是有觀察到但是不知道是什麼原因。

併發數 總吞吐量 Wall time
Sequential
Concurrent 2 31.2 tok/s 16.4s
Concurrent 4 26.9 tok/s 38.1s

Concurrent 2 併發的總吞吐量跟單請求峰值差不多(31.2 vs 31.6 tok/s),代表在排程機制上的 overhead 不大,而 Concurrent 4 大約掉到 26.9 tok/s,scale 效率大概 85%。

另外,Concurrent 4 的 wall time 從 16.4s 跳到 38.1s ,最後一個 Request 等了快 22s 才開始處理,這不是 DiffusionGemma 的問題,而是 MLX server 的 design limitation,Metal backend 看起來是沒有 CUDA 那套 concurrent kernel execution,所有的 Request都是要乖乖排隊的,建議 Mac還是先不要當 Production endpoint 。

跟 M4 24GB 的比較

前面有提到我用 M4 24GB 上測同一顆模型,結果是就是一場悲劇。

項目 M4 24GB M2 Max 96GB(standard 4-bit)
模型 footprint 16.18 GB 16.18 GB
Peak 記憶體 接近 OOM 45.7 GB
可用 context < 1K tokens ~1-2K tokens(慢但可用)
最大生成速度 12.6 tok/s 31.6 tok/s
多輪對話 OOM 勉強可
部署方式 oMLX mlx-vlm

最大的 bottleneck 還是記憶體,M4 24GB 連模型都快裝不下,完全沒空間留給 KV cache,而 M2 Max 96GB 雖然 standard 4-bit 吃掉快 45.7 GB,但至少還有空間跑 inference。


雖然 M2 Max 96GB 看起來可以在本地端流暢執行 DiffusionGemma 26B (Standard 4-bit 峰值可達 31.6 tok/s),但是記憶體與跟後端的排程機制仍限制了它的在長 Context 與併發表現。

實際上接到CLI 或是開發環境的場景,體感上還是跟現在線上服務提供的使用者經驗差滿多的,後續第二篇將移師到 GH200 透過 vLLM 轟出 1180 tok/s 的極致速度,而第三篇則會在 GB10 上挑戰 32K Context 的極限。

如果你也對大模型在不同硬體架構上的極限感興趣,歡迎持續關注後續的跨平台綜合評測!


最終採用模型: mlx-community/diffusiongemma-26B-A4B-it-4bit (Standard 4-bit)

初版測試模型: mlx-community/diffusiongemma-26B-A4B-it-mxfp4 (MXFP4,需手動修正 dequantize bug)

測試環境: mlx 0.31.2 + mlx-vlm 0.6.3

Top comments (0)