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 也在它的支援範圍內，參考一些大神的文章，就決定是它了ＸＤ

兩個量化版本的取捨

第一次踩坑的版本用的是 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