為了找到一些在地端也能讓 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 也在它的支援範圍內，參考一些大神的文章，就決定是它了ＸＤ
兩個量化版本的取捨
第一次踩坑的版本用的是 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，怎麼這麽省記憶體ＸＤ 不過後來就發現，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 底層不熟ＸＤ，只是有觀察到但是不知道是什麼原因。
| 併發數 | 總吞吐量 | 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