[eauchs]

Summary

Adds Multi-Token-Prediction (MTP) self-speculative decoding for Tencent Hunyuan 3 (hy_v3).

This is stacked on #1211 (kernelpool), which adds the base hy_v3 model + tool parsers + tokenizer fixes. That PR intentionally strips the MTP layer at load. This branch keeps kernelpool's commits as-is and adds three commits on top that keep and use the MTP head:

• models/hy_v3.py: load the MTP block instead of dropping it in sanitize(); add Model.predict_next_tokens(...) and a return_hidden_states path. sanitize() also accepts already-stacked (MLX-converted) expert weights, not just the per-expert Tencent layout.
• generate.py: mtp_generate_step(...) + auto-routing in stream_generate when num_nextn_predict_layers > 0 — no external draft model; up to two tokens per target pass. Verification is distribution-preserving and honors sampler / logits_processors / temperature / top-p / top-k exactly.
• Verification uses the residual (Leviathan) rule when possible (accept x~q with prob min(1, p(x)/q(x)), resample from relu(p-q) on reject) for strictly higher acceptance, falling back to sampler-agnostic match verification (emit the target's own sample) for temp-0, custom samplers, XTC, or active logits processors. make_sampler now tags its callable with its params so the MTP path can reconstruct the processed distributions.

Only the last three commits are mine to review (the rest is #1211). Happy to rebase to an MTP-only diff once #1211 merges.

Test checkpoint (2/3-bit experts + MTP, MLX): ox-ox/Hy3-295B-Instruct-w2q3exp-AProjQ8-SExpQ8-OutQ8-MTP-mlx.

Testing

• sanitize remapping validated in isolation for both the MLX-converted (pre-stacked experts + mtp.*) and the original Tencent (per-expert + model.layers.N) layouts.
• End-to-end on the quantized MLX checkpoint (M3 Max): loads, generates coherent text, MTP acceptance ~35-40%.

Notes

• Two distribution-preserving verification rules, both respecting sampler / logits_processors / temperature / top-p / top-k: residual (Leviathan) — accept x~q with prob min(1, p(x)/q(x)), resample from the normalized relu(p-q) on reject — used when the sampler exposes its params (any make_sampler output), temp>0, XTC off, no logits processors; and match (emit the target's own independent sample, accept when it equals the draft), the sampler-agnostic fallback used otherwise and for temp-0 (where it reduces to greedy). Residual accepts Σ min(p,q) ≥ Σ p·q (match), i.e. strictly more often, for the same output distribution. Validated by Monte-Carlo: the reconstructed distribution and the residual-emitted distribution match the sampler's to within sampling noise, and measured acceptance matches the Σ min(p,q) theory.
• Acceptance is checkpoint-bound here: the test checkpoint's MTP drafter runs 2–3-bit experts (same precision as the target), so it reflects quant precision more than the method's ceiling. Measured on M3 Max at temp=0.9: residual ~25% vs match ~18% accept on this 2–3-bit drafter; higher-precision drafters measure ~64–66% at greedy (see thread).
• No changes to switch_layers.py or the default (non-MTP) generation path.

[gilby]

Nice work — we independently built the same MTP path for Hy3 on the oMLX serving stack (same block wiring: eh_proj(cat([enorm(embed), hnorm(hidden)])) → decoder layer → final_layernorm, pre-norm hidden into the drafter, k=1 draft + bonus, persistent MTP KV cache with trim-on-reject), and have two findings that may be useful for this PR:

1. Your 35–40% acceptance is likely the checkpoint, not the algorithm. We measured drafter-precision sensitivity directly on the full 295B release (target: 4-bit imatrix quant, eh_proj held bf16), same prompts across variants:

MTP block quant	accept @ greedy	accept @ temp 0.9 (official sampling)
4-bit	64.5%	43.1%
8-bit	65.8%	49.0%

Even 8→4-bit on the drafter costs ~6 points at temp 0.9. With w2q3exp (2–3-bit experts in both target and drafter), 35–40% at greedy is roughly where that curve lands — worth re-measuring on a ≥4-bit checkpoint with the MTP block at 8-bit+ before judging the method's ceiling. Holding eh_proj unquantized also matters; aggressive quantization of the fusion projection collapses acceptance in our experience with other MTP models.

2. Consider distribution-preserving verification instead of greedy-only. The current mtp_generate_step argmaxes draft and verify and ignores sampler/logits_processors, so any request with temperature > 0 silently gets temperature-0 output. Hy3's recommended sampling is temperature=0.9, top_p=1.0 (greedy degrades its CoT), so most real usage can't take this path as-is. Standard speculative sampling (accept with prob min(1, p_target/p_draft) under the processed distributions, resample from the residual on rejection) preserves the requested distribution exactly and still pays: at temp 0.9 we see the 49% acceptance above, which nets out to a measurable end-to-end decode win once fixed per-forward overheads amortize (25→29 tok/s on an M3 Ultra; +18% at greedy).

One honest caveat from profiling the same configuration: because Hy3's active parameters are dominated by routed experts, the 2-token verify pass activates nearly disjoint expert sets (~1.6× the weight traffic of a 1-token forward), so the k=1 speedup ceiling is lower than on dense models — acceptance quality is what decides whether it nets positive.

[eauchs]

Thanks @gilby — great cross-check, and the wiring matches exactly (eh_proj(cat([enorm(embed), hnorm(hidden)])) → single decoder layer → final_layernorm, pre-norm hidden into the drafter, k=1 + bonus, persistent MTP KV cache trimmed on reject).

On (1) — you diagnosed it correctly. I pulled the actual quant config for the test checkpoint (ox-ox/…w2q3exp…MTP) and the drafter is running about as hard as the target:

• mtp.layer.mlp.switch_mlp (routed experts): gate/up 2-bit, down 3-bit
• mtp.eh_proj: 8-bit (not bf16)
• MTP attn / shared_mlp / router: 8-bit

So the drafter's routed experts sit at the same 2–3-bit as the target and the fusion projection isn't held unquantized — 35–40% @ greedy lands right on the curve you measured rather than a method ceiling.

I can't re-measure the higher-precision case myself, though: I'm on a 128 GB box, and a ≥4-bit 295B checkpoint doesn't fit — squeezing it under that ceiling is the whole reason for the w2q3exp build, so 3-bit experts is a hard cap on my hardware. If you're able to run the accept-rate sweep on your ≥4-bit build with eh_proj held bf16 and the MTP block at 8-bit+, that'd be the definitive read on the method's ceiling — your 64–66% greedy already points that way. Happy to line up prompts/config so it's apples-to-apples.

On (2) — agreed, and it's the more important fix. mtp_generate_step is greedy-only and silently ignores sampler/logits_processors, so any temp>0 request gets temp-0 output — and since Hy3's recommended sampling is temperature=0.9, top_p=1.0 (greedy degrades its CoT), most real usage can't take the path as written. I flagged it as a known limitation in the PR notes, but "known" isn't "fine." I'll switch verification to distribution-preserving speculative sampling — accept with min(1, p_target/p_draft) under the processed distributions, resample from the residual on reject — so the requested distribution is exact and temp-0.9 requests actually benefit (your 49% accept / +18%-ish end-to-end is about what I'd expect).

And thanks for the honest expert-traffic caveat — it matches what I profile: the 2-token verify pass activates nearly disjoint expert sets (~1.6× the weight traffic of a 1-token forward), so on this MoE the k=1 speedup ceiling is lower than on a dense target and acceptance quality is what decides whether it nets positive.

I'll push the sampling fix to this branch.

[eauchs]

Update on (2), @gilby — implemented. mtp_generate_step now carries both rules and picks per sampler:

• residual (Leviathan) when the sampler exposes its params (make_sampler now tags temp / top-p / top-k / min-p on the returned callable), temp>0, XTC off, no logits processors: accept x~q with min(1, p(x)/q(x)), resample from the normalized relu(p − q) on reject.
• match (emit the target's own independent sample, accept on equality) as the sampler-agnostic fallback and temp-0 greedy.

Measured on the aggressive q2-lite build (MTP routed experts 2-bit, eh_proj 4-bit), M3 Max, temp=0.9 / top_p=1.0, 80-token runs, model RAM-resident:

rule	accept	tok/s
match	17.9%	5.62
residual	25.4%	5.93

So +7.5 pts (~+42% relative) from the rule alone — directionally exactly your point. Small sample (~65 steps/rule), so I'd treat the absolute gap as directional rather than tight, but it's consistent with the provable Σ min(p,q) ≥ Σ p·q, and I verified in isolation that (a) the reconstructed processed distribution matches what make_sampler actually draws from and (b) the residual-emitted distribution matches the target p, both to within sampling noise, with the accept rate landing on the Σ min(p,q) theory.

Absolute acceptance is still checkpoint-bound — 2-bit experts in the drafter is about as hard as it gets — and I can't go past 3-bit on 128 GB, so the ≥4-bit sweep with eh_proj at bf16 is still yours to run. But the residual path is in now (eauchs:hy_v3-mtp, PR body updated), so if you point your ≥4-bit build at it we should be comparing the same code. Thanks for the nudge — this was the right call.

[PhilipJohnBasile]

Real-world data point on the exact test checkpoint this PR references (ox-ox/Hy3-295B-Instruct-w2q3exp-AProjQ8-SExpQ8-OutQ8-MTP-mlx), M5 Max 128GB — we've been running this stack in production for a Hy3 MLX build (https://github.com/PhilipJohnBasile/hy3-demolition-mlx).

Per our RESTORE.md we were pinned to a7cc305 (this PR's first self-spec commit, before 532d9d5/2ef070e added distribution-preserving + residual verification), so this may already be resolved upstream in this PR's later commits — flagging in case it's still useful data:

path	decode tok/s	vs AR
AR (no MTP)	7.43	1.0x
MTP self-spec (a7cc305)	1.59	0.21x — 4.7x slower

Outputs matched byte-for-byte both paths (correctness was never in question, only throughput). Given @gilby's note above about drafter precision — this checkpoint's MTP block runs at the same 2–3-bit as the target experts — this lines up with the "checkpoint-bound, not algorithm-bound" read. Curious whether the residual verification commit changes this number on the same checkpoint; happy to re-benchmark against current hy_v3-mtp HEAD if useful. Rooting for this landing — it and #1211 are both direct dependencies for us getting Hy3 running on mainline mlx-lm / LM Studio.

[PhilipJohnBasile]

Haven't re-benchmarked against the fixed HEAD (2ef070e) yet — still on the old a7cc305 pin, will follow up with real numbers when I do.

In the meantime, linking a related data point in case it's useful context: a separate speculative-decoding implementation for the same Hy3 checkpoint (drafter-contract based, not this PR's native self-spec path) hit a different but adjacent failure mode — youssofal/MTPLX#142 (comment) — MTP-drafted depths pass quality only at the slowest setting (27% slower than AR), fail quality (repetition loop) at the depth that's competitive on speed, and the contract calibration reports no_agreement_signal at every depth on this same w2q3exp checkpoint. Different mechanism (contract calibration vs greedy/distribution sampling), same aggressively-quantized checkpoint — plausible the quantization is the common thread, but flagging as adjacent evidence rather than a direct answer to the acceptance-rate question here.

[pierre427]

Third independent implementation datapoint, on a different Hy3 variant — supports the "checkpoint-bound, not method-bound" read, with one addition: a cheap diagnostic that separates head quality from loop plumbing.

Setup: Hy3 152B REAP-pruned to 50% experts, 4-bit/g64 (router 8-bit). That quant also strips layer 80, so we streamed the nextn head back from tencent/Hy3 (5 of 99 shards carry model.layers.80.*) and grafted it — keeping the head's own un-pruned 192-expert MoE intact rather than slicing to match the trunk. The graft is sound because the head consumes the trunk's post-final-norm hidden, which expert pruning doesn't touch. Head quantized to the same 4-bit trunk recipe. Same wiring as this PR (eh_proj([enorm(emb); hnorm(h)]) → one decoder layer → final_layernorm, k-draft + bonus). M5, 128 GB.

The diagnostic: before touching the loop, teacher-force the drafter at full-context real positions and measure the acceptance ceiling:

probe	ceiling
code prompt	83.0%
prose prompt	44.2%

Then the loop, greedy, exact-match verified:

loop	acceptance	vs AR
fresh drafter context each step, k=1	36.8%	0.88× (a loss)
persistent MTP KV (trim-on-reject), k=1	84.4%	1.20×
persistent, k=2 (copy-heavy prompt)	84.9%	1.30×

So on our checkpoint the head was fine all along — a starved drafter context masqueraded as a 37% "method ceiling," and the persistent cache (which this PR and @gilby's implementation already have) recovers the teacher-forced ceiling almost exactly. The suggestion for the w2q3exp acceptance question: run the teacher-forced ceiling probe on that checkpoint. If the ceiling itself sits at ~35–40%, that's @gilby's precision curve and no loop change will fix it; if the ceiling is much higher than the loop's acceptance, something is still starving the drafter. It decomposes the two hypotheses in one cheap measurement, no ≥4-bit rebuild needed.

Two smaller notes:

• REAP datapoint: 50% expert pruning of the trunk did not degrade the head (83% code ceiling post-graft) — consistent with the head reading only the final-norm hidden. Grafting a stripped head back from the base checkpoint is viable at ~9.6 GB of downloads if others hit the stripped-quant problem Add --preserve-mtp to keep MTP weights during conversion #1543 is addressing.
• Throughput economics (re @PhilipJohnBasile): even at 84% acceptance, k=1 only yields 1.20× for us because the 2-token MoE verify touches ~2× the routed experts of an AR step (cycle delivers 1.84 tok at ~1.53× step cost). On MoE targets the acceptance→speedup conversion is much weaker than dense intuition suggests, so sub-1× results at modest acceptance are expected, independent of the drafter-precision issue. Prose is additionally capped by the head itself (our 44% prose ceiling ≈ the 40-49% range measured up-thread).

[PhilipJohnBasile]

Follow-up from our streamed-MoE work: the expert-traffic caveat is showing up as a first-order runtime cost, not just a theoretical ceiling.

On a separate 744B GLM-5.2 runtime, we instrumented consumer-blocked expert-load stall separately from asynchronous read service. In the baseline fixed-trace decode, exposed stall was 62.3s of 130.6s (47.7%). Holding expert bytes and read count constant, deterministic post-read completion delay increased read-completion p50 and decode time together:

Added completion delay Read-completion p50 Decode time
0 µs 2.30 ms 130.6 s
4,000 µs 7.21 ms 166.9 s
10,000 µs 14.75 ms 280.5 s
16,000 µs ~22.23 ms 390–412 s

The claim is intentionally narrow: token latency is causally sensitive to added expert-read completion delay under this treatment. This is not a physical SSD-bandwidth estimate, and we are not deriving compression-speedup projections from it.

Our first full-context canary also caught a setup confound: applying the treatment from process start throttled the 8K prefill. The experiment now runs prefill with delay off, activates treatment only at the decode boundary, warms for 32 tokens, resets measurement counters at loader-idle, and measures exactly the following 256 tokens.

That v3 path is currently 620/620 green, with numerical parity and single-variable trace proofs passing and a new binary epoch frozen. The 16 ms canary is running now.

Practical implication for this PR: on streamed MoEs, MTP acceptance should be reported together with extra expert bytes/reads and consumer-exposed stall. A k=1 verification pass can still lose to added expert traffic even when acceptance improves. I’ll post the randomized-block result when it completes.

[pierre427]

Follow-up to the datapoint above — we've now reviewed the branch code (HEAD 2ef070e) against the implementation we've been running at 152B, and have a set of concrete findings, roughly ordered by severity. Happy to send any of these as patches on top of your branch.

1. prompt_cache slicing silently drops the last decoder layer.
mtp_generate_step does model_cache = prompt_cache[:-1]; mtp_cache = prompt_cache[-1] (generate.py L734-735), but make_prompt_cache(model) returns num_hidden_layers caches — hy_v3 defines no make_cache. mlx_lm.chat passes exactly that, so the trunk runs with n_layers−1 caches, zip truncation never executes the final decoder layer, and its cache slot is silently repurposed as the MTP cache. No crash, corrupted logits. (Your speculative_generate_step uses the opposite convention, prompt_cache[: len(model.layers)].) Cleanest fix: keep the MTP cache out of prompt_cache (separate make_mtp_cache()), or give hy_v3 a make_cache() returning n_layers+1 entries so the layout is real.

2. Headless (MTP-stripped) repacks no longer load.
Model always builds the MTP block when config.num_nextn_predict_layers > 0, and sanitize never drops it when the weights aren't there — but #1211's conversion strips the MTP weights while the saved config still says num_nextn_predict_layers: 1. Existing community quants ship exactly this shape (e.g. pipenetwork/Hy3-REAP50-MLX-4bit), so this branch breaks them at load_weights(strict=True). Fix is a few lines in sanitize: if no mtp.* keys survive the remap, drop the module. (Companion suggestion for #1211: zero the config field when stripping at conversion.)

3. The MTP KV cache is never synced to the committed sequence — this is the measured low-acceptance regime.
The prefill loop never feeds the MTP head, so its first draft runs at RoPE offset 0 while the trunk sits at prompt_len; on accept, the accepted token's pair is never fed (one hole per acceptance); and the reject-path mtp_cache.trim(1) (L847) removes the committed pair that was just fed — the rejected draft never entered the cache — so the trim is exactly backwards. Net: the head drafts from a gap-ridden, position-compressed, promptless KV. That's the family we measured at 36.8% acceptance / 0.88× before switching to a persistent MTP cache (teacher-forced chunk-wise at prefill; after each verify, rewind the k speculative entries and re-feed the committed span with trunk hiddens from the verify forward — one MTP call per cycle either way): 84.4% / 1.20–1.30× on the same checkpoint, numbers in the comment above. Same lift on a second vendor head (Qwen3.5-122B: 72–83% → 91–94% at k=1). One caveat: heads distilled from trunk hiddens don't need this; Hy3's vendor head does. This slots into your loop as a change to the prefill loop + the accept/reject cache handling; glad to send it.

4. Question — which hidden does the head expect?
__call__ returns the pre-final-norm hidden to predict_next_tokens (hy_v3.py L293-296). Our port feeds the post-final-norm hidden, matching our read of the reference implementation, and the 84% figure was measured that way. Worth double-checking against the reference — if the head was trained on post-norm hiddens, this costs acceptance silently.

5. Smaller items:

• num_nextn_predict_layers > 1 configs leave orphan model.layers.{n+i} keys for i≥1 (sanitize remaps only layer n; Add Hy3 #1211 stripped all depths) → strict-load error. Latent — Hy3 ships depth 1 — but cheap to handle.
• mtp_generate_step accepts max_kv_size / kv_bits / kv_group_size / quantized_kv_start / prompt_progress_callback and ignores them all.
• The stream_generate routing is unconditional for any model with the attr > 0 — worth a flag (and a check that the module actually exists, see item 2), and every token is labeled from_draft=True, which corrupts acceptance stats downstream.
• Nits: the yielded logprobs come from unprocessed logits while tokens are sampled from processed ones; the accept test does a per-token .item() sync that could batch.

If a REAP-pruned trunk carrying the base model's un-pruned MTP head is a case you want to support (it's how we run — head grafted back from tencent/Hy3), that additionally needs an mtp_num_experts config so the head's expert count can differ from the trunk's; ~10 lines in ModelArgs/MTPBlock/fix_moe. As-is the branch fails on both layouts of that checkpoint family (per-expert: orphan keys; pre-stacked: shape mismatch).

[samithaj]

Fourth datapoint on the expert-traffic caveat, from a different model family — and measured from the router rather than inferred from timing.

Setup: GLM-4.7-Flash (glm4_moe_lite, 4-bit, 64 routed experts, top-4, 46 MoE layers) running its in-checkpoint nextn head as a k=1 drafter on a paged runtime (persistent per-request head KV, rewind on reject, greedy exact-match verify). M3 Max.

Expert divergence, measured directly. I instrumented the router across 96 greedy decode steps, all MoE layers, real hidden states — 4,370 layer-position pairs:

consecutive decode positions share 25.7% of their top-4 experts (chance would be 4/64 = 6.25%).

So a k=1 verify step reads ~1.74x the expert weights of a 1-token step — between @gilby's ~1.6x and @pierre427's ~2x, same mechanism you both describe. End-to-end it lands where the arithmetic predicts: acceptance 0.770, baseline 16.8 ms/step vs 26.7 ms/step with the drafter, i.e. 1.77 tokens per cycle at ~1.59x step cost = 1.11x. Same shape as @pierre427's 1.84 tok at 1.53x = 1.20x.

A teacher-forced ceiling from a loop that isn't starving the head. @pierre427's diagnostic is the right first move, and it seems worth having a reference point for what a healthy loop looks like. Same head, teacher-forced at full-context real positions, 191 slots per prompt:

prompt	ceiling
code	85.3%
chat	84.8%
prose	75.9%
overall	82.0%

Live acceptance in the serving loop is 77.0% — within ~5 points of the ceiling, and the prompt mixes differ, so some of that is prompt variance (live per-prompt acceptance ranges 0.71-0.83 on this head). So here the persistent cache is already capturing essentially all of the head's quality; the gap that turned into 37% -> 84% on hy_v3 is absent. That makes the diagnostic doubly useful — it tells you which of the two problems you have, and ~5 points is what "the loop is fine" looks like.

The prose/code split is consistent with your 44% prose vs 83% code ceiling: this head is much stronger on structured text than prose, on a similar spread.

On the quantization thread: this head is 4-bit against a 4-bit target, so by @gilby's curve there should be a few points available from an 8-bit drafter. I have a bf16 split of the same head published and will measure that pair next.

(#1543 is mine — glad it is useful. It exists because of exactly the stripped-quant problem @pierre427 hit.)

[PhilipJohnBasile]

---

Useful fourth point, and the router-measured framing is the right upgrade over
inferring from timing — 25.7% top-4 overlap against a 6.25% chance floor is a
clean number to anchor the arithmetic on.

One thing worth flagging before anyone stacks these across models: the
overlap measured here (consecutive decode positions, one model, one scheme)
isn't the same object as the scheme-vs-scheme overlap we measured elsewhere
(same position, two quantization schemes, one engine) — different axis
entirely, and they shouldn't be read as comparable percentages even though
both are "expert overlap under a perturbation."

On your last line — "this head is 4-bit against a 4-bit target, so there
should be a few points available from an 8-bit drafter" — we have a
same-engine, two-scheme measurement that bears on it, on a different
aux-loss-free-router MoE (not this one): quantization scheme changes perturb
routing but don't reshuffle it (~78-84% of top-k retained against a matched
null well below that), and a sigmoid+correction-bias router wasn't
particularly fragile to it. So the "few points from precision" intuition looks
right in direction; whether it holds numerically for a drafter head
specifically is untested by that work and worth being explicit about the gap.

We started a matrix on Hy3's MTP head that overlaps with your setup (teacher-
forced acceptance ceiling vs live, pre/post-norm hidden, bf16/fp32 router) but
it stalled after the model load with no arms collected yet -- flagging the
overlap now rather than promising results I don't have. If we get it running
cleanly, happy to compare notes, since a fifth model family on the
expert-traffic caveat would make the pattern harder to dismiss as
architecture-specific.

[nastya236]

Thanks for your contribution! We recently introduced a new policy regarding the use of AI in pull request descriptions and discussions #1770.
Unfortunately, this means that we need to close this pull request. However, you’re very welcome to open a new one that follows the updated policy. Thanks for your understanding!