Preserve in-checkpoint MTP head through convert and add a speculative-decoding path (Qwen3.5/3.6/3.8)

[kira-ariaki]

Summary

Qwen3.8-27B ships an in-checkpoint MTP head (15 tensors, mtp.*). base-convert drops them, and the engine has no speculative-decoding path, so decode stays strictly bandwidth-bound: one full sweep of the weights per token.

Measured on an M1 Max 64 GB, the same model, same quantization class, same machine:

BaseRT 0.2.0  base-q4  14.39 GB    12.5 tok/s     (no MTP)
oMLX  0.5.7   oQ4e     16.10 GB    14.2 tok/s     (Lightning MTP, accept 72-85%)

The interesting part is not that oMLX is 14% faster overall — it is that BaseRT's kernels are 1.7× more efficient and still lose, because the other runtime emits 2.0-2.5 tokens per backbone pass while BaseRT emits exactly one.

Environment

• BaseRT 0.2.0, official macOS arm64 release bundle
• macOS 26.4.1 (build 25E253), Apple M1 Max, 64 GB
• Model: Qwen/Qwen3.8-27B (bf16 original weights, 55.56 GB / 18 shards), converted locally with base-q4
• Comparison runtime: oMLX 0.5.7 with scottlowry/Qwen3.8-27B-oQ4e-mtp

Current behaviour

base-convert reports the drop explicitly:

  arch:    qwen3_5
  config:  hidden=5120, layers=64, heads=24/4, ffn=17408, vocab=248320
  mapped:  851 tensors kept, 333 mmproj, 15 dropped
  quantized 851 tensors
  wrote 851 tensors (14731 MB)

The 15 dropped tensors are exactly the MTP head:

mtp.fc.weight
mtp.norm.weight
mtp.pre_fc_norm_embedding.weight
mtp.pre_fc_norm_hidden.weight
mtp.layers.0.input_layernorm.weight
mtp.layers.0.post_attention_layernorm.weight
mtp.layers.0.self_attn.{q,k,v,o}_proj.weight
mtp.layers.0.self_attn.{q,k}_norm.weight
mtp.layers.0.mlp.{gate,up,down}_proj.weight

strings basert-serve | grep -ic mtp returns 0; the only speculation-adjacent symbol in the binary is a bare speculation string with no draft/verify machinery behind it. No release from 0.1.0 through 0.2.2 mentions MTP, multi-token prediction, or speculative decoding in its changelog.

Measurement

Kernel efficiency, isolated from speculation

The fair comparison is how long one full sweep of the weights takes, since that is pure kernel work and independent of how many tokens the sweep yields. oMLX logs backbone time and cycle count per request, which makes this directly computable:

                  weights    per sweep    ideal @316 GB/s    efficiency
BaseRT base-q4    14.39 GB     80.0 ms        45.5 ms           57%
oMLX oQ4e         16.10 GB    148.5 ms        50.9 ms           34%

(316 GB/s is the measured large-buffer bandwidth on this machine, not the 400 GB/s spec figure.)

BaseRT's kernels are 1.7× more efficient per byte moved. That advantage is real and it is BaseRT's.

What speculation does to the result

From the same oMLX logs, per-request MTP statistics:

MTP[2] tokens=65 cycles=32 tok/cycle=2.03 accept=31/43 (72.1%)
       depth[d1=20/27,d2=8/12,d3=3/4]
       timing[backbone=3898.8ms mtp=121.3ms sample=2.7ms cache=15.5ms]

MTP[3] tokens=60 cycles=24 tok/cycle=2.50 accept=34/40 (85.0%)
       depth[d1=17/21,d2=14/16,d3=3/3]
       timing[backbone=3562.9ms mtp=99.3ms sample=2.1ms cache=9.2ms]

Dividing backbone time by cycle count gives what that runtime would do without speculation, and the arithmetic closes cleanly:

oMLX, speculation disabled (1 token/cycle)     6.7 - 8.2 tok/s
oMLX, speculation enabled  (2.0-2.5 tok/cycle)     14.2 tok/s
BaseRT, no speculation available                   12.5 tok/s

The MTP head costs 2.8% of backbone time (99.3 ms against 3562.9 ms) and returns 2.0-2.5× the tokens. Acceptance holds at 72-85% on ordinary Chinese and English prose, with depth-1 and depth-2 both landing consistently.

Applying the same 2.5 tokens/cycle to BaseRT's measured 80.0 ms sweep projects ~31 tok/s — roughly 2.5× the current figure, and above what either runtime achieves today.

Prior art

• SGLang reports 38.28 tok/s decode for Qwen3.8-27B on DGX Spark. That machine has 273 GB/s of memory bandwidth; against a ~13.5 GB NVFP4 checkpoint the pure-autoregressive ceiling is 273/13.5 ≈ 20.2 tok/s. The published number is 1.9× that ceiling, which is only reachable with speculation. Their Qwen3.8-27B cookbook page lists in-checkpoint MTP as a supported feature.
• vLLM supports Qwen3.x MTP speculative decoding out of the box (--speculative-config with "method": "qwen3_next_mtp").
• oMLX ships Lightning MTP for the qwen3_5 family on Metal, with the acceptance statistics quoted above; the model-side requirement is only that the mtp.* tensors survive conversion.

All three consume the same in-checkpoint head that Qwen publishes — no separate draft model, no extra download.

Why this looks tractable

The MTP head is one transformer layer plus a fusion projection: fc (concatenates the previous hidden state with the next token embedding), two pre-fusion norms, a standard attention block, an MLP, and a final norm. It reuses the main model's embedding and LM head, so the only new weights are those 15 tensors — 0.44 GB in bf16, well under 3% of the checkpoint.

Conversion looks like the smaller half of the work: the tensors are already enumerated and classified well enough to be deliberately dropped, and their names follow the same mtp.layers.N.* shape as the main stack, so the existing to_canonical_name prefix chain should extend to them without new naming rules. config.json carries mtp_num_hidden_layers: 1 and mtp_use_dedicated_embeddings: false for the runtime to key on.

Runtime needs the usual draft-verify loop: run the head to propose k tokens, batch-verify them in one backbone pass, accept the longest matching prefix. On the hybrid GDN architecture the verify step needs care — the recurrent state must roll back on rejection, unlike a pure-attention KV cache — but oMLX demonstrates it working on this exact architecture on Metal, and baseRT::GDNStateCache::{create,destroy} already exist as independently constructed objects in the shipped binary, which suggests the state boundary is where it needs to be.

Two adjacent notes:

• Support NVIDIA Nemotron-H architecture (Mamba-2 + MoE + Attention hybrid, e.g. Nemotron-3.5-Lightning-30B-A3B) #46 asks for Nemotron-H support and observes in passing that MTP is dropped at convert. This issue is narrower: the head is already in a supported architecture's checkpoints today, so the conversion side needs no new architecture work.
• Hybrid (Qwen3.5/3.6 GDN) models serialize all requests: support N independent recurrent states #31 (hybrid GDN request serialization) touches the same recurrent-state machinery. Speculative decode and multi-sequence lanes both need per-sequence GDN state that can be snapshotted and restored, so the two may share an implementation.

Suggested scope

An incremental path that yields something useful at each step:

1. Preserve mtp.* through base-convert behind a flag (--keep-mtp), quantized with the same profile as the main stack. Bundles grow ~0.2 GB at q4. Even with no runtime support this makes the weights available and turns the drop into a decision rather than a silent loss.
2. Single-draft speculation (k=1) in the engine: propose one token, verify in the next backbone pass. Roughly 1.5-1.7× at the acceptance rates measured above, and it exercises the full rollback path with the simplest possible state handling.
3. Multi-depth drafting (k=2-3), matching what oMLX gets 2.0-2.5 tokens/cycle from.

Step 1 alone would unblock experimentation; the current behaviour makes it impossible to try speculation on BaseRT even for someone willing to write the decode loop against the C API.

I have the converted bundles, the raw conversion logs, and the measurement harness for both runtimes, and I am happy to re-run anything on this M1 Max — including an A/B against a --keep-mtp build if that would help size the work.