Support for Gemma 4 assistant draft models

[cryi]

Is your feature request related to a problem? Please describe.
omlx currently lacks support for the new Gemma 4 assistant draft models, preventing the use of speculative decoding to accelerate text generation.

Describe the solution you'd like
Implement support for Gemma 4 assistant draft models, specifically models like gemma-4-26B-A4B-it-assistant-bf16, to enable faster inference speeds natively within omlx.

Describe alternatives you've considered

Additional context
Here is the link to the specific assistant model on Hugging Face: https://huggingface.co/mlx-community/gemma-4-26B-A4B-it-assistant-bf16

Can we expect to see this ported or supported in omlx anytime soon?

[miabchdave]

Speculative decode is actually most effective for dense models. So if oMLX was to support it, that's where the best possible chance of a real speedup would be - with Gemma 4 31B:

https://huggingface.co/google/gemma-4-31B-it-assistant

[cryi]

I think I read somewhere that even MoE can gain something like extra 50%. For me any speedup is good. But I assume the implementation for both will be shared. 🙂

[crazyi]

It is the same to DFLASH？

[cryi]

It is the same to DFLASH？

Same goal but likely slightly different. See the mlx-vlm docs — MTP requires the --draft-kind mtp and --draft-block-size <size> flags.

Notes: Google Gemma MTP docs

[panwudi]

+1 — this is becoming a real bottleneck for Mac-local Gemma 4 deployments.

A few data points from running gemma4-moe-26b-a4b on M5 Max 128GB via oMLX 0.3.9.dev1:

• BF16 inference at ~50 tok/s single-stream is fine for chat, but the official Google MTP drafter promises ~3× speedup, which would put Gemma 4 26B-A4B firmly in the same throughput tier as Qwen 27B Dense + DFlash.
• Blaizzy's mlx-vlm already shipped Gemma 4 MTP support, so the algorithm + Apple MLX kernels are proven to work; the remaining gap is dispatching from oMLX's engine into that path.
• Per-model specprefill_draft_model already exists in the admin UI for Gemma 4, so the user-facing config surface is mostly there — the work seems concentrated in omlx/engine/ and omlx/adapter/.

A few questions for the maintainer:

1. Is there a tentative timeline or is this awaiting community contribution?
2. If we (community) were to draft a PR porting the mlx-vlm Gemma 4 MTP dispatch into oMLX's engine, would you accept it, and is there a preferred architecture for the integration (e.g., reuse mtp_enabled flag and bump the
   mtp_compatible check, vs. introducing a parallel gemma_mtp_* field set)?
3. Any reason this couldn't reuse the gemma4_assistant config_model_type detection that the admin endpoint already reports?

Happy to test branches and contribute benchmark data on M5 Max if it helps.

[panwudi]

Update after reading mlx-vlm gemma4/language.py more carefully: this isn't
actually a MTP integration — Gemma 4 uses standard drafter-based speculative
decoding (separate small drafter, e.g., e2b → 31b). The mtp_enabled / mlx_lm_mtp
path won't fit. A new path (e.g., vlm_drafter_enabled) wiring to
LanguageModel.rollback_speculative_cache would be more accurate.

[cryi]

Available in 0.3.9rc1