mlx-mtp-graft

MTP graft — restoring Multi-Token-Prediction heads to MLX quantizations that dropped them

Most MLX quantizations of Qwen3.5/3.6/3.8 keep mtp_num_hidden_layers: 1 in config.json while silently dropping the MTP weights. The checkpoint loads, generates correctly, and gives you no reason to suspect anything — it just cannot do multi-token prediction, and nothing tells you why.

This repo contains the MTP head only (314 MB), plus a tool to graft it onto a recipient checkpoint. No base weights are redistributed.

Measured on an Apple M5 Max, 128 GB, macOS 26.4.1.

What's in here, and what it applies to

Not claimed to work on Qwen3.5-27B or Qwen3.6-27B — not tested. Nothing here was trained; the head is Qwen's own (see Provenance).

Result

Recipient: Qwen3.8-27B abliterated with Heretic (trohrbaugh/Qwen3.8-27B-heretic-ara, MLX affine 4-bit g64), which ships with 0 MTP tensors. Served through oMLX, greedy, thinking off, A/B/A interleaved with cooldowns, 3 rounds:

Acceptance at depth 1 is 93.7% on code, 53% on Chinese prose, from oMLX's own per-request log. tok/cycle 3.53 and 1.88 respectively.

The serving runtime matters more than any of this. See below.

The runtime dominates — read this before optimising anything else

The same MTPLX-packaged Qwen3.8-27B artifact, same machine, same day:

Both runtimes measure the model identically — depth-1 acceptance 0.935 (MTPLX) vs 0.937 (oMLX), tok/cycle 3.38 vs 3.39. The model behaves the same. The execution differs by more than 2×.

Two things follow, and both cost me a wrong conclusion first:

Draft-head precision. MTPLX ships its head as an FP16 sidecar (mtp_policy: keep_bf16, 849 MB vs 314 MB here). Loading that FP16 head into oMLX, the draft step took 252 ms against 27 ms for the 4-bit head — and acceptance was identical (93.7% both). I concluded FP16 was a bad trade. In MTPLX's own runtime the same FP16 head costs 0.78 s over 229 draft calls = 3.4 ms each. The 252 ms was oMLX's execution path for a sidecar it was not built around, not a property of FP16. That comparison is void. The 4-bit head here is smaller and works; whether it is better is untested.

Depth. oMLX runs an adaptive controller (cap 8) scoring (1 + p1 + p1·p2 + …) / t_est(d) from live measurements, and chose depth 3. MTPLX's tune measured depth 2 as best on this Mac by a wide margin (115.6 vs 81.6). The controller is well-designed — its t_est is just measured on its own slower path. Correct optimisation over a wrong cost model.

If you graft a head and care about throughput, benchmark the runtime before you tune the artifact.

Sparse MoE: no gain under oMLX

Grafted onto a qwen3_5_moe 35B-A3B recipient (donor: mlx-community/Qwen3.6-35B-A3B-MTP-4bit, 46 tensors), oMLX reports Lightning MTP (model_type=qwen3_5_moe, active) and 59% acceptance — and buys nothing: 102–117 tok/s ungrafted, 101–108 grafted. The controller detects it and stops:

MTP[1] finish=parked  tok/cycle=1.29  accept=25.8%  depth[d1=8/31, d2=0/0, d3=0/0]

oMLX's source names the case: "MoE expert loads scale with verify tokens … the whole depth menu can be worse than standard decoding." The MoE's single-token step is already 9.5 ms against the dense model's 33.8 ms; there is little to amortise.

Counter-evidence, and it is not weak: EryriLabs grafts a head onto the same architecture and reports +29% under llama.cpp on an RTX 3090, at comparable acceptance. Given the MTPLX result above, the likeliest explanation is again the runtime, not the architecture. Measure your stack; do not take the MoE result here as general.

Output fidelity

Verify-then-accept means the head cannot make the model say something the base model rejects. But batched K-token verification and single-token decoding take different floating-point paths, so near-ties flip at temperature 0.

Fixed seed, temperature=0, top_p=1, top_k=0:

before: …因為它假設近期被使用過的資料[[，未來被再次使用的機率較高。]]
after : …因為它假設近期被使用過的資料[[在未來被再次使用的機率較高。]]

Both fluent, neither better. Low-entropy text has no near-ties and reproduces exactly. If you need bit-reproducible output, do not enable MTP.

Known incompatibility: mlx_lm breaks on MTP tensors

A checkpoint containing MTP tensors produces garbage under mlx_lm.load, silently. Controlled test, same checkpoint, only the 46 MTP tensors differing:

BigBang-v1-heretic-MLX-4bit      (MTP = 0)   -> normal Chinese output
BigBang-v1-heretic-MLX-4bit-mtp  (MTP = 46)  -> ";av L WL two one lya"**开始中j位ode理ments谷雨谷…"

• oMLX / MTPLX — work, and are the only stacks that use the head.
• mlx_vlm — works (ignores it).
• mlx_lm — broken. Keep an ungrafted copy if you need this path.

This cost me a wrong diagnosis: the same garbage appeared on a VL checkpoint that had both vision and MTP tensors, and I blamed vision. It is the MTP tensors.

Naming conventions differ between runtimes

oMLX expects language_model.mtp.*. MTPLX ships a sidecar named mtp.* and loads it through its own runtime spec (mtplx_runtime.json). Handing an MTPLX artifact to oMLX fails cleanly:

VLM load failed: Received 15 parameters not in model: mtp.*
LLM fallback also failed: Lightning MTP is enabled ... the converted weights are
missing the mtp.* tensors. Default mlx-lm converters strip them; you need a
converter that preserves MTP.

(Note oMLX's error text — its author documents this exact failure mode.) graft_mtp.py normalises mtp.*, bare standalone heads, and already-namespaced language_model.mtp.* to what oMLX expects.

Usage

python graft_mtp.py \
  --recipient /path/to/model-without-mtp \
  --donor     /path/to/donor-or-standalone-head \
  --out       /path/to/model-with-mtp

Recipient weights are hard-linked, not copied; the graft costs only the MTP shard. The original directory is never modified. The tool refuses mismatched hidden_size, head counts, vocab_size, num_experts, moe_intermediate_size, or a different base quantization scheme (--force overrides).

In oMLX: set mtp_enabled and unload/reload the model. oMLX reads model settings at load time only — editing model_settings.json while the model is resident does nothing (verified twice).

Verifying a graft

Three checks, cheapest first.

1. Are the tensors there? Check weights, not config — the config lies. Of eight Qwen3.5/3.6/3.8 MLX checkpoints on hand, four declared mtp_num_hidden_layers: 1 with zero MTP tensors present.

wm = json.load(open(f"{d}/model.safetensors.index.json"))["weight_map"]
print(sum(1 for k in wm if "mtp" in k.lower() or "nextn" in k.lower()))

2. Are they trained? Counting is not enough — a head can be correctly shaped random initialisation. EryriLabs reported this on 2026-08-20 for Ornith-1.5-35B-A3B: 785 mtp.* tensors at std = 0.0200 with Gaussian kurtosis 3.0 — literally initializer_range — accepting ~13% of drafts. Every tensor-count check passes. (Fixed upstream on 2026-08-23; sampling six of those tensors by HTTP range request now gives norm means 0.74–0.87 and projection kurtosis 5.7–10.4, and they match EryriLabs' grafted head to within 1 bf16 ULP — upstream adopted the same fix.)

Two tells, no inference needed: norm weights sit at mean ≈1 or above when trained, at the initializer range when not; projections are heavy-tailed when trained, Gaussian (≈3.0) with identical std across every layer when not.

The head in this repo: norm means 1.036 / 1.206 / 1.791 / 1.779 / 2.252; projection kurtosis 8.10 / 3.71 / 14.65, std 0.0159 / 0.0127 / 0.0198. Trained, and corroborated by the 93.7% acceptance above.

3. Is the wiring right? The head consumes concat(embedding, hidden), and there are choices: which hidden (pre- or post-norm), concat order, and local vs global position indexing. Get one wrong and acceptance collapses without an error. MTPLX exposes them as --base-hidden-variant, --concat-order, and its artifacts declare an explicit mtp_contract. This graft matched by defaults because donor and recipient share a base — that is luck, not method. If you graft across less related checkpoints, sweep them.

Provenance

The head is Qwen's own:

1. Qwen/Qwen3.8-27B — original MTP heads
2. tozp/Qwen3.8-27B-OBLITERATED-V2-mlx-4bit — MLX affine 4-bit g64; its abliteration_metadata.json records "Restore MTP + vision tensors from stock", so the heads it carries are unmodified
3. this repo — those 29 tensors, extracted unchanged

29 tensors, 314.3 MB, sha256[:16] = 505827ce17922d2d.

For qwen3_5_moe (hidden 2048, 256 experts) the donor is mlx-community/Qwen3.6-35B-A3B-MTP-4bit — same tool, 46 tensors, grafts cleanly. Read the MoE section first.

Limitations

• One machine (M5 Max). oMLX numbers only; the graft has not been run under MTPLX's runtime, which is where the 3.79× lives.
• No long-context, no batch >1.
• Grafting a stock head onto an abliterated model is not obviously sound in theory. Empirically: the recipient's KL to stock is 0.0535, while a checkpoint shipping this head natively sits at 0.9624 — the recipient is ~18× closer to stock than a model where the head demonstrably works. And the failure mode is bounded: a mismatched head lowers acceptance, it cannot corrupt output.
• Contract parameters (§Verifying, step 3) assumed rather than swept.
• Not tested with vision inputs on VL recipients.

Credits

• Qwen — the MTP heads
• tozp — the 4-bit quantization they came from
• Heretic / trohrbaugh — the recipient
• oMLX — Lightning MTP
• MTPLX — the runtime that showed how much of this is execution rather than weights, and the contract vocabulary
• EryriLabs — the random-init diagnostic and the llama.cpp MoE counter-result