Feature Request: MTP (Multi-Token Prediction) Speculative Decoding for Gemma 4

[vduke]

Feature Request: MTP (Multi-Token Prediction) Speculative Decoding for Gemma 4

Summary

Google recently released official MTP (Multi-Token Prediction) draft models for Gemma 4, enabling 2x decode speedup with zero quality loss. Currently, oMLX strips MTP weights during quantization (_should_skip_tensor in oq.py) rather than using them for inference.

Context

Google's Gemma 4 architecture includes optional MTP heads that predict multiple tokens per forward pass. These are shipped as separate "assistant" models:

• google/gemma-4-26B-A4B-it-assistant — MTP drafter for 26B-A4B MoE
• google/gemma-4-31B-it-assistant — MTP drafter for 31B dense

vLLM already has a PR implementing MTP speculative decoding: vllm-project/vllm#41745 (submitted by a Google engineer).

Current Behavior

In omlx/oq.py, MTP tensors are explicitly excluded:

def _should_skip_tensor(name: str) -> bool:
    if ".mtp." in name or name.startswith("mtp."):
        return True
    return False

MTP weights are treated as bloat to be stripped, not a feature to be leveraged.

Proposed Behavior

1. Preserve MTP weights during model loading (don't strip them)
2. Add an mtp_enabled model setting (similar to specprefill_enabled)
3. When MTP is enabled, use the built-in MTP heads for speculative decoding during generate:
   • Forward pass produces N candidate tokens from MTP head
   • Full model verifies candidates in parallel
   • Accepted tokens are emitted, rejected tokens trigger single-token fallback
4. Support loading the separate assistant model as the MTP drafter

Expected Performance

• 2x decode speedup with zero quality loss (per Google's benchmarks)
• Particularly impactful for longer generation sequences
• Works alongside existing SpecPrefill (prefill speedup) for combined gains

Implementation Reference

• Google's MTP architecture: Gemma 4 technical report
• vLLM PR: [Spec Decode] Add Gemma4 MTP speculative decoding support vllm-project/vllm#41745
• The "assistant" models contain the MTP weights in the standard format

Environment

• oMLX: 0.3.8
• macOS: Darwin 25.4.0 (arm64)
• Hardware: Apple M3 Ultra, 256GB

[miabchdave]

Just want to add that the MLX Community Google Assistant model card on Hugging Face has the mlx_vlm instructions for running the Gemma 4 assistant model with recommended settings for single and batch decode:

https://huggingface.co/mlx-community/gemma-4-31B-it-assistant-bf16

EDIT:
I can confirm using the mlx_vlm Gemma 4 BF16 with assistant draft model (above instructions):

Gemma 4 31B BF16 Without MTP Draft: Generation: 75 tokens, 7.318 tokens-per-sec
Gemma 4 31B BF16 With MTP Draft: Generation: 75 tokens, 13.408 tokens-per-sec

M5 Max 128GB

[jeffhamons]

Benchmark: M1 Max 64GB, gemma-4-31b-it-8bit target + gemma-4-31B-it-assistant-bf16 drafter

Hardware: M1 Max, 64GB unified memory
Stack: mlx_vlm main (0.4.5-dev), --draft-kind mtp, --draft-block-size 6
Target: gemma-4-31b-it-8bit (31GB on disk)
Drafter: gemma-4-31B-it-assistant-bf16 (926MB on disk)

Results:

Config	tok/s	Accepted / 6 drafted	Acceptance rate
temp=1.0, analytical prompt	4.8	2.32	38.7%
temp=0 (greedy), plain prompt	7.6	4.15	69.2%
Baseline: 31B alone, mlx_lm.server (no MTP)	8.1	—	—
oMLX native (no MTP), warm, real workload	10–27	—	—

Key findings:

1. Quantization mismatch is a real but manageable problem. At greedy (temp=0), the 8-bit target and bf16 head agree 69% of the time — that's usable. The bf16 result above (7.3 → 13.4 tok/s on M5 Max) likely sees ~80-85% acceptance under similar conditions.

2. Temperature dominates acceptance rate. Going from temp=1.0 to temp=0 moved acceptance from 38.7% → 69.2%. For production workloads that run at temp=1.0 (analytical summarization, synthesis), the 8-bit+bf16 combination regresses vs baseline.

3. The mlx_lm.server serving path is the practical bottleneck. Even before MTP overhead, mlx_lm.server on this hardware runs the 31B at 8.1 tok/s vs oMLX native at 10-27 tok/s. MTP can't overcome a slower baseline.

The case for native oMLX MTP integration:

If oMLX integrated the MTP decode loop natively (at its current 10-27 tok/s baseline), the math changes significantly. At temp=1.0 with ~40% acceptance, a 10 tok/s floor becomes ~14 tok/s and a 27 tok/s peak becomes ~38 tok/s. That would meaningfully reduce the KV cache saturation watchdog firings that degrade performance on long overnight runs (21 flushes in our most recent 11-hour run).

Happy to share the bench scripts or run additional configs if useful.