Attach Qwen3.6 MTP head for speculative decoding + fix cache-rewind for hybrid recurrent models

[h9q2cyxvgm-ui]

Fixes #1462.

Two things bundled together because the second was needed to actually exercise/verify the first end to end:

1. Attach the MTP head instead of loading it as a peer draft model

The community-published qwen3_5_mtp drafter checkpoints (e.g. mlx-community/Qwen3.6-27B-MTP-bf16) are DeepSeek-V3-style MTP heads: no embedding table or output head of their own (they borrow the target model's), and their forward pass needs the target'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 — hence the Model type qwen3_5_mtp not supported error in #1462.

• models/qwen3_5.py: MTPHead module, lazily attached in sanitize() only when a checkpoint's weights actually bundle mtp.* keys — not eagerly from config alone, since the config can declare mtp_num_hidden_layers > 0 as an architecture capability while the specific file doesn't carry the tensors (every currently-published Qwen3.6 checkpoint does this).
• utils.py: load_mtp_head(model, mtp_path) attaches a separately published MTP checkpoint onto an already-loaded target model; is_mtp_head_checkpoint(config) detects one.
• generate.py: mtp_speculative_generate_step(), parallel to speculative_generate_step but chaining the attached MTP head against the target's own hidden states. stream_generate dispatches here when model.has_mtp and no separate draft_model was given.
• server.py: ModelProvider._load() detects an MTP-head checkpoint and attaches it instead of peer-loading it.

2. Give ArraysCache checkpoint/restore so speculative decoding can actually run

Once loading worked, a second and independent bug showed up: any speculative decoding on Qwen3.6 — including the original, unmodified peer-model mechanism, no MTP code involved — fails with "Speculative decoding requires a trimmable prompt cache (got {'ArraysCache'})". ArraysCache holds this model family's hybrid linear-attention layers' recurrent state as a fixed-size array, not an append-only buffer like KVCache, so is_trimmable() correctly returns False — there's nothing to slice back to an earlier position after a rejected draft token.

• models/cache.py: is_checkpointable() / checkpoint() / restore() added to _BaseCache (default: unsupported) and implemented on ArraysCache (a shallow list-copy of self.cache is a safe, cheap snapshot — verified across every hybrid-model file that uses ArraysCache that cache slots are always replaced wholesale via __setitem__, never mutated in place). New module functions can_rewind_prompt_cache(), checkpoint_prompt_cache(), rewind_prompt_cache().
• generate.py: both speculative_generate_step() and mtp_speculative_generate_step() gate on can_rewind_prompt_cache() instead of can_trim_prompt_cache(). When the cache mixes trimmable KV layers with checkpoint-only recurrent layers, a rejected round can't just trim (recurrent state has no notion of "state as of token k" other than recomputing it) — instead it checkpoints every non-trimmable layer before each round, and on partial rejection undoes the whole round (trims KV back to the pre-round offset for free, restores checkpointed layers) and replays exactly the tokens that turned out correct through the full model in one extra forward call. This is unavoidable — a monolithic forward touches every layer together, so cheaply trimming KV while separately patching recurrent layers isn't possible; they have to move in lockstep. This is generic infra shared by every hybrid linear-attention model in mlx_lm, not just Qwen3.6.
• Also fixes a tensor-rank bug in mtp_speculative_generate_step (only reachable once the cache fix let execution get this far for the first time) and adds a guard so an exception during draft generation doesn't crash the finally-block cleanup.

Verified

Temp=0 speculative output is token-for-token identical to plain generation, for both the peer-draft-model path (small Qwen3-0.6B draft) and the attached-MTP-head path, on unsloth/Qwen3.6-27B-UD-MLX-6bit. Full existing test suite has the same pre-existing failures before and after this change (confirmed via git stash on a clean checkout).

Honest caveat on speedup

Modest and inconsistent right now — roughly 1.2x on a 40-token generation, ~0.9x (slightly slower) on 200 tokens, at ~66% draft acceptance either way (with the MTP head; see the follow-up PR for how that acceptance rate was fixed). 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; whether this specific rewind strategy is the right performance tradeoff for hybrid-recurrent models in general is a fair thing to discuss — happy to iterate (e.g. tuning num_draft_tokens down, or a cheaper reject-recovery path) based on feedback.

Also open

A follow-up PR fixes a real accuracy bug in MTPHead itself (wrong concat order + wrong hidden-state norm stage) that took MTP draft acceptance from 0% to ~66% — separated out since it's a standalone correctness fix independent of the cache work above, and depends on the MTPHead class this PR introduces.

🤖 Generated with Claude Code

[h9q2cyxvgm-ui]

Follow-up data point on the speedup caveat above, since it's worth being precise about what's actually causing it.

I ran a controlled comparison: same prompt, same num_draft_tokens sweep, temp=0, 150 tokens — but against a dense Qwen3 model (mlx-community/Qwen3-32B-6bit, plain KVCache only, no hybrid recurrent layers) with a small peer draft model (Qwen3-0.6B), using the existing, unmodified speculative decoding path (this model was never affected by the ArraysCache issue in the first place).

num_draft_tokens	Dense Qwen3-32B (cheap reject via .trim())	Hybrid Qwen3.6-27B + MTP head (expensive reject via checkpoint+replay)
1	1.21x (43% accepted)	1.04x (48% accepted)
2	1.13x (54% accepted)	1.00x (60% accepted)
3	1.00x (61% accepted)	0.93x (67% accepted)

At every setting, the dense model comes out ~0.15–0.2x ahead of the hybrid model — despite the hybrid model's MTP head actually guessing right more often at each setting. That's consistent with the theory in the description above: it isn't that this rewind approach is poorly implemented, it's that any rejected round on a hybrid model pays a real, structural cost (a full extra forward pass to rebuild the recurrent state) that a pure-KV model simply doesn't pay. A better guesser doesn't fully compensate for a more expensive miss.

Practical takeaway for anyone picking this up: tuning num_draft_tokens down helps a bit (1 is best of the three tested) but doesn't get hybrid-recurrent models to a clear win on its own — closing that gap further would need a cheaper way to recover from a rejected round than a full-model replay, which is a separate, harder problem than what's fixed in this PR.

Also flagging, found while running this comparison and unrelated to anything in this PR: speculative_generate_step with num_draft_tokens=4 on the dense Qwen3-32B + Qwen3-0.6B pairing above produced output that didn't match plain generation at temp=0 (should be lossless). Reproduces on main too (checked before any of this branch's changes), so it's a separate, pre-existing bug — will file it as its own issue rather than mixing it in here.