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.jsonwhile 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:
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.
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_enabledand 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.
python
1wm = json.load(open(f"{d}/model.safetensors.index.json"))["weight_map"]2print(sum(1for 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:
Qwen/Qwen3.8-27B — original MTP heads
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
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.
