Gemma-4-31b-it 4bit Performance worse than MLX on M1 MAX

[Ahdw]

Hardware: M1 MAX, 64GB unified Memory, 24c GPU

Model: gemma-4-31b-it
MLX inferer: oMLX
MLX quantization: oQ4-fp16, no vision

Prefill 1024 tokens at 94.1 tokens/s, Generate 128 tokens at 14.0 tokens/s

see: https://omlx.ai/benchmarks/y9evupla

baseRT
convert:

basert convert \
  path/to/google/gemma-4-31B-it \
  -o path/to/gemma-4-31B-it-q4-f16scale.base \
  --target base-q4 \
  --profile path/to/BaseRT/base-convert/profiles/default-q4-f16scale.json

benchmark:

basert bench path/to/gemma-4-31B-it-q4-f16scale.base -p 1024 -n 128

| model | size | test | t/s |
| --- | ---: | ---: | ---: |
| gemma-4-31B-it-q4-f16scale.base | 32458 MiB | pp1024 | 74.25 ± 0.06 |
| gemma-4-31B-it-q4-f16scale.base | 32458 MiB | tg128 | 11.74 ± 0.01 |

[prabod]

@Ahdw Thanks for testing this model on your device and sharing the benchmarks. We have only optimized for sub-30B dense models at the moment, which explains the performance gap with the 31B Gemma model. However, optimizations for these larger models are next on our roadmap. Thanks for flagging this! (I'll keep this issue open to track our progress on this)

[Ahdw]

@prabod any update on this issue? I've done some math:

─── HARDWARE ──────────────────────────────────────────────────────────────
FP16 Compute : 7.83 TFLOPS (M1 Max 24c, FP16:FP32 = 1:1)
Bandwidth    : 400 GB/s (0.4 TB/s)

─── MODEL PROFILE ─────────────────────────────────────────────────────────
Model        : gemma-4-31b-it-UD-MLX-4bit-fp16
Architecture : Gemma4ForConditionalGeneration (Dense)
Parameters   : ~31B (All Active)
Layers       : 60
Attention    : 32 Q-Heads / 16 KV-Heads (GQA 2:1)
Dimensions   : hidden=5376, intermediate=21504, head=256
Attn Type    : 50L Sliding Window (1024) + 10L Full Attention
Quantization : UD-MLX Mixed 4/5/6/8-bit (Dominant 4-bit, group_size=64)
Disk Size    : ~15.5 GB (31B × 0.5 Byte)

─── THEORETICAL LIMITS ────────────────────────────────────────────────────
Prefill (Compute-Bound)
  • FLOPs/Token : 2 × N ≈ 62 GFLOPS
  • Max Speed   : 7.83 TFLOPS ÷ 62 GFLOPS ≈ 126.3 tok/s

Decode (Memory-Bound)
  • Memory Read : ~15.5 GB/token (Model Weights)
  • Max Speed   : 400 GB/s ÷ 15.5 GB ≈ 25.8 tok/s
───────────────────────────────────────────────────────────────────────────

And based on my benchmarking:

┌──────────────────┬──────────┬────────┬────────┬───────────┬────────────────┐
│    Test Item     │   TTFT   │ pp TPS │ TG TPS │ Peak Mem  │ Bandwidth Util.│
├──────────────────┼──────────┼────────┼────────┼───────────┼────────────────┤
│ pp1024/tg128     │ 10,442ms │  98.1  │  12.1  │ 22.09 GB  │     66.82%     │
│ pp4096/tg128     │ 43,134ms │  95.0  │  11.6  │ 23.82 GB  │     69.08%     │
│ pp8192/tg128     │ 88,246ms │  92.8  │  11.0  │ 26.11 GB  │     71.80%     │
│ pp16384/tg128    │ 183,031ms│  89.5  │  10.1  │ 27.83 GB  │     70.27%     │
│ pp32768/tg128    │ 390,740ms│  83.9  │   8.6  │ 32.09 GB  │     68.99%     │
└──────────────────┴──────────┴────────┴────────┴───────────┴────────────────┘
Comparison Baselines:
- gemma-4-31B-it-oQ4-fp16 (Official oMLX oQ quantization, ~93 t/s pp1024)
- gemma-4-31B-it-oQ8-fp16 (Official oMLX oQ quantization, ~90 t/s pp1024)

* Note: Memory Bandwidth Utilization is estimated using (Peak Mem × TG TPS) / 400 GB/s, rounded to 2 decimal places (e.g., 80.72%).

Fun Fact is that, the Unsloth Dynamic 2.0 quantization gives better prefill speed. I suspect that it's due to the alignment of data.

Do you think with baseRT we can manage to have a 100+ tps prefill speed and 15+ tps decode on this machine particularly ?