KaedeTai/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

file	applies to
qwen3.8-27b-mtp-4bit.safetensors	Qwen3.8-27B dense only — qwen3_5, hidden 5120, MLX affine 4-bit g64. Verified on one recipient.
graft_mtp.py	Any MLX checkpoint with an MTP subtree. qwen3_5 dense and qwen3_5_moe verified here; other families untested.

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:

prompt	before (no MTP)	after graft	speedup
Python LRU cache	29.0 / 29.9 tok/s	61.1 (spread 18.5%)	2.11×
Chinese prose	29.4 tok/s	40.9	1.39×

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:

runtime	depth	decode tok/s
MTPLX 2.9.1 (own runtime)	AR (no MTP)	30.5
MTPLX 2.9.1	D1	104.6 (3.43×)
MTPLX 2.9.1	D2	115.6 (3.79×)
MTPLX 2.9.1	D3	81.6 (2.68×)
oMLX	3 (its adaptive choice)	54.6

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:

content	tokens (before / after)	identical	similarity
Python code	165 / 165	yes	1.000
Arithmetic	3 / 3	yes	1.000
Chinese explanation	70 / 70	no — 1 token	0.992
Chinese creative prose	192 / 217	no — diverges at ~50%	0.663

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