為了找到一些在地端也能讓 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
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)