Support Qwen3.5-MoE MoE MTP heads

[janfeddersen-wq]

Summary

Adds support for Qwen3.5-MoE MTP heads to the native Qwen MTP path.

The MTP head on Qwen3.5-MoE checkpoints (e.g. Qwen/Qwen3.5-122B-A10B) is itself an MoE block — router gate + per-expert MLPs + a shared_expert and shared_expert_gate — whereas the existing path assumed a dense single-MLP head. As a result these models were recognized as qwen3-next-mtp but rejected at the tensor gate with invalid-mtp-tensor-layout.

No new backend is required: mlx-lm's qwen3_5 DecoderLayer already instantiates SparseMoeBlock when num_experts > 0, so the MTP block builds correctly once the weights are stacked. The change is two small, dense-safe additions:

• artifacts.py — for MoE heads, derive the expected MTP key set from num_experts (mtp.layers.*.mlp.{gate, experts.{i}.*, shared_expert.*, shared_expert_gate}) instead of the fixed dense 15, so inspect's tensor gate passes. No-op for dense heads.
• mtp_patch.py — _stack_mtp_moe_experts stacks per-expert mtp.layers.*.mlp.experts.{i}.{proj} tensors into the switch_mlp layout SwitchGLU expects, mirroring the stacking mlx-lm performs for the main decoder layers. No-op for dense heads.
• Tests for the MoE tensor gate and the expert stacking.

Verification

Tested on Qwopus3.5-122B-A10B (qwen3_5_moe, 256 experts / 8 active), bf16 MTP sidecar grafted onto a 4-bit base:

• mtplx inspect → can_run: true, tensor gate 785/785, 0 missing / 0 extra.
• 120-token greedy run: accepted_by_depth = [40, 19, 3] of [57, 57, 56] drafted → ~70% depth-1 acceptance, 120 tokens in 57 verify passes (~2.1 tokens/verify). A mis-loaded MoE head would accept ~0%.
• pytest tests/test_artifacts.py tests/test_mtp_patch.py green.

Known limitation — MoE exactness

At temperature 0, MTP vs non-MTP greedy decode is ~98% identical and re-converges immediately, but occasionally flips a single token. This is the MoE router hitting a near-tie that resolves differently under batched verification vs single-token autoregressive decode (a known MoE/FP effect), not a drafting error. The max_diff = 0.0 exactness guarantee was established on a dense model; strict bit-exactness for MoE heads likely needs separate handling (e.g. fp32 router logits during verify). Flagging for discussion — happy to follow up.

🤖 Generated with Claude Code

[youssofal]

Closing this one out with apologies for the long silence — and with a correct diagnosis on your part. The MoE-MTP-head gap you identified (router gate + per-expert MLPs + shared_expert/shared_expert_gate in the head itself, rejected at the tensor gate) was closed in-house the same week this PR was opened: artifacts.py derives the expected MoE key set from num_experts, and the loader stacks numbered experts into the switch_mlp layout — functionally the same two changes you proposed. That support has shipped in every release since v1.0.4.

This branch has since drifted a long way behind main (the diff now spans 277 files of old history), so rather than ask for a rebase of equivalent functionality I'm closing it. If your Qwen3.5-MoE checkpoints hit any gap on current v2.0.2 — inspect gate, expert stacking, or the MoE-router near-tie exactness behavior you documented — please open an issue with the artifact details; that report would be genuinely valuable. Thank you for the careful work and the acceptance-rate receipts.