feat(inference): support the qwen3_5_mtp drafter for Qwen3.5/3.6/3.8 MTP speculative decoding

Summary

Add support for the qwen3_5_mtp drafter so Qwen3.5 / 3.6 / 3.8 targets can run MTP speculative decoding. Today mlxcel's MTP round loop exists but only Gemma 4 implements MtpTarget, and qwen3_5_mtp is not a known drafter model_type. Both failure paths are named and fail closed, so this is a missing capability rather than a bug.

Where the MTP weights actually live

This was the first thing to establish, because the naming suggests the weights ride along inside the model.

They do not, in any MLX build. Verified against the safetensors index of every mlx-community Qwen3.8-27B variant, and against the downloaded shard headers directly:

repository	tensors	`mtp.*` tensors
`mlx-community/Qwen3.8-27B-4bit`	2180	0
`mlx-community/Qwen3.8-27B-8bit`	2180	0
`mlx-community/Qwen3.8-27B-bf16`	1184	0
`mlx-community/Qwen3.8-27B-mxfp8`	1682	0
`mlx-community/Qwen3.8-27B-nvfp4`	1682	0
`Qwen/Qwen3.8-27B` (upstream HF bf16)	1199	15

The upstream base checkpoint does ship MTP, as a sibling of model., not nested under it:

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.mlp.{gate,up,down}_proj.weight
mtp.layers.0.self_attn.{q,k,v,o}_proj.weight
mtp.layers.0.self_attn.{q,k}_norm.weight

The MLX conversion pipeline drops them, and mlx-community republishes them split out as a standalone drafter:

mlx-community/Qwen3.8-27B-MTP-bf16 — model_type: "qwen3_5_mtp", block_size: 3, single 849 MB model.safetensors, full text_config mirroring the target (hidden_size 5120, head_dim 256, mtp_num_hidden_layers: 1, mtp_use_dedicated_embeddings: false).

mlxcel independently drops the same keys at load: weights.retain(|k, _| !k.contains("mtp.")) (src/models/qwen3_5.rs:2337). That line is correct and should stay: upstream uses the presence of mtp. keys as a raw-HF-layout probe for the RMSNorm +1.0 shift, and #776 already hardened mlxcel against double-shifting because of it.

Current behavior (measured, `main` @ `6e781fcf`)

Target models/qwen3.8-27b-4bit, drafter models/qwen3.8-27b-mtp-bf16:

$ mlxcel generate ... --draft-model models/qwen3.8-27b-mtp-bf16
Error: Unsupported model type: qwen3_5_mtp

$ mlxcel generate ... --draft-model models/qwen3.8-27b-mtp-bf16 --draft-kind dflash
Error: Unsupported model type: qwen3_5_mtp

$ mlxcel generate ... --draft-model models/qwen3.8-27b-mtp-bf16 --draft-kind mtp
Error: --draft-kind mtp is only supported for Gemma 4 (text, VLM, or Unified) targets;
       the loaded target is not MTP-capable. Omit --draft-kind to use the classic
       SpeculativeGenerator with your --draft-model drafter.

Both messages are accurate and fail closed. No silent wrong-path behavior to fix.

Gap analysis

drafter_kind_by_model_type() (src/lib/mlxcel-core/src/drafter/mod.rs:208-220) maps only gemma4_assistant and gemma4_unified_assistant. A qwen3_5_mtp drafter therefore auto-resolves to DEFAULT_DRAFTER_KIND = Dflash (:198) and dies in the generic loader.
No MtpTarget impl for this family. The only impls are Gemma 4 (src/models/gemma4_mtp_target.rs, six variants). Qwen35Model and Qwen35VLModel implement dflash::SpeculativeTarget (qwen3_5.rs:1892, src/vision/qwen3_5_vl.rs:218) but not MtpTarget.
No drafter model. There is no Rust analogue of upstream's mlx_vlm/speculative/drafters/qwen3_5_mtp/.

Upstream contract to port

From mlx_vlm/speculative/drafters/qwen3_5_mtp/ (unchanged since 72aa649f, 2026-07-10, so the pinned reference tree is current):

Module shape, all sized on the target's hidden_size:

self.fc = nn.Linear(2 * hidden_size, hidden_size, bias=False)
self.pre_fc_norm_embedding = nn.RMSNorm(hidden_size, eps=rms_norm_eps)
self.pre_fc_norm_hidden    = nn.RMSNorm(hidden_size, eps=rms_norm_eps)
self.layers = [layer_cls(args=layer_config, layer_idx=0) for _ in range(mtp_layers)]
self.norm = nn.RMSNorm(hidden_size, eps=rms_norm_eps)

Every MTP layer is built with replace(text_config, num_hidden_layers=mtp_layers, full_attention_interval=1). full_attention_interval=1 forces is_linear = False, so MTP layers are always full attention, never gated-delta. That matches the 15 published tensors, which carry self_attn.* and no linear_attn.*.

Forward, the only thing that crosses from target to drafter besides offsets:

h = mx.concatenate([self.pre_fc_norm_embedding(token_embed),
                    self.pre_fc_norm_hidden(hidden)], axis=-1)
h = self.fc(h)

Borrowed, not owned. bind(target_model) takes the target's embed_tokens (walking target, target.model, target.language_model.model) and its lm_head, falling back to embed_tokens.as_linear. The drafter has neither. draft_block raises if bind was skipped.

KV caches are separate, not shared. set_shared_kv(...) explicitly discards the passed states and consumes only the offsets, to align positions. Each MTP layer gets its own KVCache / BatchKVCache.

block_size is the total round budget including the bonus token, not the draft count: flat.setdefault("block_size", int(mtp_depth) + 2) and while len(tokens) < block_size - 1. For Qwen3.8-27B, mtp_num_hidden_layers: 1 gives block_size: 3, i.e. 2 drafted tokens per round. The drafter sets prefer_requested_block_size = True, bypassing the adaptive block-size heuristic.

Sanitize strips the mtp. prefix and applies the +1.0 RMSNorm offset only for raw-HF layout (key.startswith("mtp.")), over the suffixes .input_layernorm.weight, .post_attention_layernorm.weight, .q_norm.weight, .k_norm.weight, norm.weight, pre_fc_norm_embedding.weight, pre_fc_norm_hidden.weight.

Target-side surface the round loop requires (mlx_vlm/speculative/mtp.py), all of which exist on upstream qwen3_5.LanguageModel:

rollback_speculative_cache — hard requirement, checked before the loop starts
speculative_verify_hidden(inputs, cache) — must return hidden states, not logits
speculative_logits_from_hidden(hidden)
speculative_argmax_from_hidden(hidden) — optional fast path
speculative_verify_logits(inputs, cache, sampler) — non-hidden fallback
make_cache() returning the mixed list [ArraysCache(size=2) if is_linear else KVCache()]

Advertised drafter flags: supports_greedy_draft_argmax, prefer_requested_block_size, supports_ragged_batch_acceptance all true; requires_uniform_batch_acceptance false.

Implementation plan

Add a Qwen35Mtp drafter model and loader keyed on model_type: "qwen3_5_mtp"; register it in drafter_kind_by_model_type() so auto-detect resolves to DrafterKind::Mtp instead of falling through to DFlash.
Implement MtpTarget for Qwen35Model and Qwen35VLModel, including the mixed recurrent/KV rollback. The rollback is the hard part: partial acceptance must trim the KV of the 16 full-attention layers and restore the gated-delta recurrent state of the 48 linear layers to the accepted prefix. A recurrent state cannot be trimmed; it must be snapshotted before the draft block and restored on partial accept, or recomputed. Land B=1 first and gate B>1 behind a passing ragged-acceptance test.
Wire the drafter into the server round loop (src/server/speculative_dispatch.rs) and the offline CLI path.
Weight-inventory test against the pinned mlx-community/Qwen3.8-27B-MTP-bf16 index: every one of the 15 source tensors maps exactly once, no unknown required tensor remains, and the +1.0 shift is applied to exactly the seven norm suffixes and nothing else.
Exactness gate: at temperature == 0, MTP output must be byte-identical to classic decode over a fixed prompt set. This is the same contract src/server/batch/mtp_policy.rs:51-58 states for Gemma 4, and issue fix(speculative): Qwen3.5 MTP verification drift breaks temperature-0 parity #77 is the precedent for how MTP verification drift breaks it.
Extend docs/supported-models.md's speculative table (currently | DFlash | Qwen 3.5 text/VLM paths | at :523) with the MTP row and the required drafter repository.

Follow-ups deliberately not blocking

MTP for the 4-bit target. The published drafter is bf16 while the practical target is 4-bit. Confirm the bf16 drafter pairs with a quantized target before assuming it; if not, a quantized drafter must be converted or split from Qwen/Qwen3.8-27B.
Splitting MTP out of a raw HF checkpoint. Upstream ships split.py for this. A mlxcel-surgery subcommand would let users produce a drafter from Qwen/Qwen3.8-27B without a Python detour. Separate issue.
Multimodal. The upstream drafter README states image/video prefill still runs on the target and MTP only accelerates the text decode tail. Match that; do not advertise speculative multimodal prefill.
Qwen3.5 / 3.6 drafters. The same model_type covers them (upstream's README is titled "Qwen3.5 / Qwen3.6 MTP Drafter"). Validate at least one non-3.8 pairing before closing.

Acceptance criteria

--draft-model <qwen3_5_mtp checkpoint> auto-resolves to DrafterKind::Mtp with no --draft-kind
mlxcel generate and mlxcel-server both run MTP speculative decode against mlx-community/Qwen3.8-27B-4bit + mlx-community/Qwen3.8-27B-MTP-bf16
Temperature-0 output byte-identical to classic decode across a fixed prompt set
Partial-accept rollback restores both the full-attention KV and the gated-delta recurrent state; a mutation test fails if the recurrent state is left un-restored
Weight-inventory test against the pinned drafter index (15 tensors, exact +1.0 suffix set)
Measured acceptance rate and end-to-end tok/s vs classic decode, medians over repeated runs with the machine load recorded, written to docs/benchmark_results/
B>1 either passes ragged-acceptance tests or is rejected with a named error
docs/supported-models.md speculative table updated
Loader never infers MTP availability from mtp_num_hidden_layers: the drafter is keyed on its own model_type, and a strict weight-inventory gate rejects any directory that is not a standalone qwen3_5_mtp head (comment 2026-08-15)

Split out of #1163 (Qwen3.8-27B qualification), where this was found. See that issue for the full research baseline and measurement context.

Additional evidence on where the MTP weights live, because the public framing and the MLX reality disagree and the disagreement is easy to miss.

Every official/reference Qwen3.8 checkpoint ships MTP inside the model repository. The MLX conversions are the outlier, not the norm:

repository	tensors	`mtp.*`
`Qwen/Qwen3.8-27B`	1199	15
`Qwen/Qwen3.8-27B-FP8`	1606	22
`unsloth/Qwen3.8-27B`	1199	15
`Qwen/Qwen3.8-2.4T-A95B`	1609	19
`mlx-community/Qwen3.8-27B-{4bit,8bit,bf16,mxfp8,nvfp4}`	2180 / 2180 / 1184 / 1682 / 1682	0

The Qwen3.8-27B model card lists MTP as a model component in Model Overview, alongside the FFN and the LM output ("MTP (Multi-Token Prediction): trained with multiple steps"), and the card directs production users to vLLM / SGLang / TokenSpeed, which all read the MTP block out of the same directory. So "Qwen3.8 has built-in MTP" is a correct reading of the published material; it just stops being true after the mlx-vlm conversion.

The MLX config still advertises MTP after the weights are gone. text_config.mtp_num_hidden_layers: 1 and mtp_use_dedicated_embeddings: false are byte-identical between Qwen/Qwen3.8-27B and mlx-community/Qwen3.8-27B-4bit, even though the latter has zero mtp.* tensors. Nothing in the MLX checkpoint signals the loss except counting the tensor list.

Two consequences for this issue:

A loader must not infer MTP availability from mtp_num_hidden_layers. Presence of the actual mtp.* tensors (or of a separate qwen3_5_mtp drafter directory) is the only reliable signal. Worth an explicit test.
The "split MTP out of a raw HF checkpoint" follow-up listed above is more valuable than it first appeared: since the reference weights carry MTP and only the MLX conversion drops it, users arriving from the vLLM/SGLang documentation will reasonably expect MTP to work and will need a supported way to recover it from Qwen/Qwen3.8-27B rather than depending on mlx-community/Qwen3.8-27B-MTP-bf16 continuing to exist.

The config/weights inconsistency is an upstream mlx-vlm converter issue rather than an mlxcel one, but it is the reason this is confusing and it is worth reporting there.

feat(models): qualify Qwen3.8-27B on the qwen3_5 path (#1174)

Premise correction before anyone starts this, because the issue body overstates the difficulty in one specific and consequential way.

The body says:

The rollback is the hard part: partial acceptance must trim the KV of the 16 full-attention layers and restore the gated-delta recurrent state of the 48 linear layers to the accepted prefix. A recurrent state cannot be trimmed; it must be snapshotted before the draft block and restored on partial accept, or recomputed.

That is already implemented. Qwen35Model::rollback_speculative_cache (src/models/qwen3_5.rs:1932) does exactly this today, for the DFlash path:

Attention layers: trims by block_size - (max(accepted) + 1), with per-row tail zeroing for B > 1 via zero_per_row_kv_tail.
Gated-delta layers: replays the captured block over [:, :n] from GdnRollbackSnapshot, so the post-block state matches the accepted-prefix position. The snapshots carry the original q/k/v/a/b and init_state tensors at their captured dtype, and the doc explicitly warns against promoting them to f32.

The verify pass that produces those snapshots already returns VerifyOutput { logits, hidden_states, gdn_states }, so the hidden states an MTP drafter needs are also already on hand rather than needing a new forward variant.

So the genuinely hard, family-specific piece of this issue does not need to be invented. What remains is real work, but it is wiring and a new drafter rather than novel cache surgery:

Register qwen3_5_mtp in drafter_kind_by_model_type() (src/lib/mlxcel-core/src/drafter/mod.rs:208), which currently maps only the two Gemma 4 assistant types. Without it the drafter auto-resolves to DEFAULT_DRAFTER_KIND = Dflash and dies in the generic loader.
Add a qwen3_5_mtp drafter model and loader for the 15-tensor MTP head, structurally mirroring the existing Gemma 4 assistant drafter.
Implement mlxcel's own MtpTarget trait for Qwen35Model and Qwen35VLModel, wrapping the existing verify and rollback machinery. Note this is mlxcel's trait shape (prefill_and_seed / embed_token / verify_forward / verify_finalize / num_layers / eos_token_ids plus batched variants), not upstream's Python method names that the issue body quotes.
Server round-loop wiring, and lifting the Gemma-4-only gate behind the --draft-kind mtp is only supported for Gemma 4 (text, VLM, or Unified) targets error.
The exactness gate: temperature-0 output byte-identical to classic decode.

The rest of the issue body stands, including the observation that mlx-community/Qwen3.8-27B-4bit ships zero mtp.* tensors so the drafter must come from mlx-community/Qwen3.8-27B-MTP-bf16 or be split out of the bf16 base.

Recording this rather than silently scoping it down, since "the recurrent rollback must be invented" is the sentence most likely to make someone size this issue wrong in either direction.

Reproduction instructions for other Apple Silicon generations

The verdict in this PR is hardware-specific: MTP measures 0.70x on M1 Ultra, and the cause is that this GPU generation does not amortize a multi-token forward at the block sizes MTP uses. This codebase already records the same split for Gemma 4 (mtp_b1_default, issue #165): M1 Ultra 0.75-0.96x, M5 Max 1.2-1.4x. The M5 cell for Qwen is open, and so is M3/M4.

If you have another machine, here is how to fill those cells. Results on any generation are welcome on this PR or on #1165, including negative ones.

Requirements

Unified memory: 32 GB minimum, 48 GB comfortable. The target is 15 GB of 4-bit weights, the drafter adds about 1.7 GB (bf16 on disk, loaded as f16), plus KV and activations. A 16 GB or 24 GB machine will not hold this pairing. Base M3 and base M4 are out; M3 Max / M4 Max / M4 Pro at 36 GB or more are in.
Xcode with the Metal toolchain, not just Command Line Tools. If xcrun -f metal fails, set DEVELOPER_DIR to a full Xcode, for example /Applications/Xcode.app/Contents/Developer. Every command below assumes it is exported.
Disk: about 17 GB for the two checkpoints.

Setup

git fetch origin && git checkout feature/issue-1165-qwen3-5-mtp-drafter

export DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer   # or your Xcode path
cargo build --release --features metal,accelerate

hf download mlx-community/Qwen3.8-27B-4bit    --local-dir models/qwen3.8-27b-4bit
hf download mlx-community/Qwen3.8-27B-MTP-bf16 --local-dir models/qwen3.8-27b-mtp-bf16

Test 1: the decisive one, and it is cheap

This is the measurement that actually decides the verdict, and it does not need the MTP path at all. It times the target's own forward at increasing block sizes, on both the verify path and the plain prefill path, so it isolates whether your GPU gives you K positions for less than K times the price.

export DEVELOPER_DIR=/Applications/Xcode.app/Contents/Developer
cargo test --release --features metal,accelerate \
  --test qwen38_mtp_chain_parity verify_forward_cost_scaling -- --ignored --nocapture

M1 Ultra baseline to compare against:

T	`forward_speculative`	plain batched prefill
1	43.6 ms	43.6 ms
2	75.7 ms	75.0 ms
3	106.2 ms	104.4 ms
4	135.5 ms	133.2 ms
8	244.0 ms	238.7 ms

How to read your numbers. Take (T=4 cost) / (T=1 cost). On M1 Ultra it is 3.1, meaning four positions cost more than three single-token steps and MTP cannot win. If your machine lands near 1.5 or below, MTP should be a clear win there and the static gate for your hardware should be revisited. Anything between is the interesting middle.

The two columns should stay within a couple of percent of each other on any hardware. If they diverge on yours, that is a different finding and worth reporting on its own, because it would mean the verify machinery, not the forward, is costing something.

Test 2: end-to-end, the number people will quote

PROMPT="Explain in detail how a transformer neural network works, covering attention, feed-forward layers, normalization, and training."

# classic decode
./target/release/mlxcel generate -m models/qwen3.8-27b-4bit -p "$PROMPT" -n 300 --temp 0

# MTP, drafter-configured block size (3)
./target/release/mlxcel generate -m models/qwen3.8-27b-4bit \
  --draft-model models/qwen3.8-27b-mtp-bf16 --draft-kind mtp --draft-block-size 3 \
  -p "$PROMPT" -n 300 --temp 0

# MTP, block 4
./target/release/mlxcel generate -m models/qwen3.8-27b-4bit \
  --draft-model models/qwen3.8-27b-mtp-bf16 --draft-kind mtp --draft-block-size 4 \
  -p "$PROMPT" -n 300 --temp 0

Run each arm 5 times and take the median; report min-max too. Record uptime load before and after, since this matters more on a loaded machine than the arms differing.

M1 Ultra baseline: classic 23.45 tok/s, MTP block 3 16.48 (0.70x, acceptance 0.591), MTP block 4 13.81 (0.59x, acceptance 0.465).

Acceptance is deterministic at temperature 0, so if your acceptance differs materially from 0.591 at block 3, something is wrong with the pairing rather than with the hardware, and that is worth reporting immediately rather than continuing the sweep.

Test 3: exactness, please run this on every generation

The chain-parity kernel is Metal-only and hardware-specific, so its bitwise guarantee needs re-confirming per GPU generation. These three should pass on any Apple Silicon machine that can hold the model:

cargo test --release --features metal,accelerate --test qwen38_mtp_chain_parity \
  parity_kernel_block_is_bitwise_equal_to_single_token_chain -- --ignored --nocapture
cargo test --release --features metal,accelerate --test qwen38_mtp_chain_parity \
  block_verify_chain_matches_single_token_chain -- --ignored --nocapture
cargo test --release --features metal,accelerate --test qwen38_mtp_chain_parity \
  adapter_chain_with_rejected_drafts_matches_classic_chain -- --ignored --nocapture

A failure here is more important than any speed number. It would mean temperature-0 MTP output is not byte-identical to classic decode on your hardware, which is the contract the feature rests on.

Optional, to see what byte-exactness costs on your GPU:

cargo test --release --features metal,accelerate --test qwen38_mtp_chain_parity \
  parity_kernel_cost_vs_standard -- --ignored --nocapture

M1 Ultra measured parity/standard ratios of 1.089 (T=1), 1.117 (T=3), 0.834 (T=4), 1.008 (T=64), i.e. inside dispatch noise, and about 3.4% end to end.

What to report back

Whatever you run, please include: the machine (chip, memory), macOS version, the uptime load during the run, and the raw numbers rather than a summary. A single line of "M4 Max 64 GB, T=4/T=1 ratio 2.2" is already useful, and a negative result is as useful as a positive one.

If a generation lands clearly positive, the follow-up is to revisit the static per-hardware gate rather than to change any of the code in this PR: the gate is what currently keeps MTP from engaging where it does not pay.

test: pin MTP dispatch-site coverage and fix the chain-parity test lint

fix(security): bound the qwen3_5_mtp drafter config scalars at load

fix(speculative): harden Qwen 3.5 MTP exactness gate and defaults

fix(speculative): correct stale MTP docs and trim a held-too-long buffer

docs: add PR #1182 technical reports and CHANGELOG entry

feat(speculative): add qwen3_5_mtp drafter and Qwen 3.5 MTP target su…

docs(comments): cut churn from the speculative-decoding comments

Navigation Menu

Repository navigation

feat(inference): support the qwen3_5_mtp drafter for Qwen3.5/3.6/3.8 MTP speculative decoding #1165

Summary

Where the MTP weights actually live

Current behavior (measured, `main` @ `6e781fcf`)

Gap analysis

Upstream contract to port

Implementation plan

Follow-ups deliberately not blocking

Acceptance criteria

Reproduction instructions for other Apple Silicon generations

Requirements

Setup

Test 1: the decisive one, and it is cheap

Test 2: end-to-end, the number people will quote

Test 3: exactness, please run this on every generation

What to report back

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

feat(inference): support the qwen3_5_mtp drafter for Qwen3.5/3.6/3.8 MTP speculative decoding #1165

Description

Summary

Where the MTP weights actually live

Current behavior (measured, main @ 6e781fcf)

Gap analysis

Upstream contract to port

Implementation plan

Follow-ups deliberately not blocking

Acceptance criteria

Activity

inureyes commented on Aug 15, 2026

inureyes commented on Aug 15, 2026

inureyes commented on Aug 16, 2026

inureyes commented on Aug 16, 2026

Reproduction instructions for other Apple Silicon generations

Requirements

Setup

Test 1: the decisive one, and it is cheap

Test 2: end-to-end, the number people will quote

Test 3: exactness, please run this on every generation

What to report back

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Current behavior (measured, `main` @ `6e781fcf`)