mlx_lm.server --draft-model fails to load any Qwen3.6 MTP checkpoint (model_type qwen3_5_mtp not registered)

[h9q2cyxvgm-ui]

Summary

mlx_lm.server --draft-model fails immediately at load time for every Qwen3.6 MTP drafter checkpoint I've tried, regardless of which target model it's paired with. The drafter's config.json declares model_type: "qwen3_5_mtp", which has no corresponding module in mlx_lm.models and no entry in MODEL_REMAPPING, so _get_classes() raises immediately, before any inference is attempted.

This is a different failure mode from #1292 (which is about truncated completions on an already-loaded MTP model); this one never gets the model loaded at all.

Affected versions

• mlx-lm 0.31.3 (pip, current latest at time of writing)
• mlx 0.31.2
• macOS 26.4.1, Apple Silicon (M5 Max, 128GB)
• Python 3.13

Affected models

• Target: unsloth/Qwen3.6-27B-UD-MLX-6bit (loads fine standalone — model_type: qwen3_5)
• Draft: mlx-community/Qwen3.6-27B-MTP-bf16 (model_type: qwen3_5_mtp — this is what fails)

Minimal reproduction

pip install mlx-lm==0.31.3
mlx_lm.server \
    --model unsloth/Qwen3.6-27B-UD-MLX-6bit \
    --host 127.0.0.1 --port 8082 \
    --draft-model mlx-community/Qwen3.6-27B-MTP-bf16

Observed output

Traceback (most recent call last):
  ...
  File ".../mlx_lm/server.py", line 695, in _generate
    self.model_provider.load_default()
  File ".../mlx_lm/server.py", line 385, in load_default
    self.load("default_model", None, "default_model")
  File ".../mlx_lm/server.py", line 363, in _load
    draft_model, draft_tokenizer = load(draft_model_path)
  File ".../mlx_lm/utils.py", line 491, in load
    model, config = load_model(model_path, lazy, model_config=model_config)
  File ".../mlx_lm/utils.py", line 334, in load_model
    model_class, model_args_class = get_model_classes(config=config)
  File ".../mlx_lm/utils.py", line 191, in _get_classes
    raise ValueError(msg)
ValueError: Model type qwen3_5_mtp not supported.

Root cause

mlx_lm/utils.py's _get_classes() does:

model_type = config["model_type"]
model_type = MODEL_REMAPPING.get(model_type, model_type)
arch = importlib.import_module(f"mlx_lm.models.{model_type}")

MODEL_REMAPPING has no qwen3_5_mtp entry, and there's no mlx_lm/models/qwen3_5_mtp.py. Checked the full registry directly (pkgutil.iter_modules(mlx_lm.models.__path__), 118 entries): nothing containing mtp exists at all, for any model family. So this isn't specific to this one checkpoint, it affects every Qwen3.6 MTP drafter on the Hub today, since none of them have a registered loader.

Expected output

The draft model should load and speculative decoding should proceed (or fail with something actionable at generation time, not at import time for a documented CLI flag).

Workaround

None found. mlx_vlm (a separate package) does have gemma4-family multimodal classes registered, but nothing for qwen3_5_mtp either. The MTPLX project (github.com/youssofal/MTPLX) reportedly works around this by implementing its own inference path for MTP heads instead of going through mlx_lm's standard model registry, rather than fixing the registry gap itself.

[h9q2cyxvgm-ui]

Update: fixed on a fork branch, with the fix and full findings pushed here:

https://github.com/h9q2cyxvgm-ui/mlx-lm/tree/feature/qwen3-5-mtp-speculative-decoding

What it took

The model_type: qwen3_5_mtp registration gap turned out to be the first of several layered issues, not the whole problem:

1. Loading: these MTP checkpoints are DeepSeek-V3-style heads — no embedding table or output head of their own, and their forward pass needs the target model's hidden state as input, not just token ids. That's architecturally incompatible with mlx_lm's generic --draft-model mechanism (which only knows how to call an independent peer model with token ids). Fix: attach the head directly onto the target model instead of loading it as a peer (utils.load_mtp_head(), generate.mtp_speculative_generate_step(), wired into server.py's ModelProvider). See commit 0c67fee.

2. A second, independent blocker: once loading worked, any speculative decoding on Qwen3.6 — including the original unmodified peer-model mechanism, no MTP code involved — failed with "Speculative decoding requires a trimmable prompt cache (got {'ArraysCache'})". ArraysCache holds this model's hybrid linear-attention layers' recurrent state as a fixed-size array, not an append-only buffer like KVCache, so it can't be trimmed back to an earlier position after a rejected draft. Fix: gave ArraysCache checkpoint/restore support, and had the speculative-decoding rewind logic checkpoint before each round and replay the accepted prefix on rejection instead of trimming. This is generic infra affecting every hybrid linear-attention model in mlx_lm, not just Qwen3.6. See commit 4847a49.

3. A real accuracy bug in the MTP head itself: even after 1 and 2 worked mechanically, the MTP head's drafted tokens had ~0% agreement with the target model's actual output — not occasionally wrong, essentially unrelated. Two bugs in MTPHead: (a) it concatenated [hidden, embedding] before its projection, but the correct DeepSeek-V3-style order is [embedding, hidden]; (b) it was being fed the post-final-norm hidden state instead of the pre-final-norm residual stream the MTP head is trained on. Both bugs silently produce a structurally valid but semantically scrambled vocabulary projection rather than an error, which is why they were easy to miss. Confirmed by manually swapping the concat order and comparing draft vs. actual target output — fixing both took draft acceptance from 0% to ~66%. See commit 06d99ad. This one is worth a standalone PR independent of everything else — it's a plain correctness bug in the MTP head, not specific to speculative decoding.

Verified

Temp=0 speculative output is token-for-token identical to plain generation, for both the peer-draft-model path and the attached-MTP-head path, on unsloth/Qwen3.6-27B-UD-MLX-6bit. No regressions in the existing test suite.

Honest caveat on speedup

It's modest and inconsistent right now — 1.23x on a 40-token generation, 0.92x (slightly slower) on 200 tokens, ~66% draft acceptance either way. The reason: full_attention_interval=4 means 3 of every 4 layers are the non-trimmable recurrent kind, so any rejected round costs a full extra replay forward through most of the network. Correctness is solid; the performance case for hybrid-recurrent models needing this cache-rewind approach specifically may need more work (e.g. tuning num_draft_tokens down, or a cheaper reject-recovery path) before it's a clear win.

Happy to open this as a PR (or split into separate PRs per the three issues above) if that's useful — let me know which shape would be easiest to review.