qwen3_5: load in-checkpoint MTP head + speculative rollback for hybrid (GDN) caches#1456
qwen3_5: load in-checkpoint MTP head + speculative rollback for hybrid (GDN) caches#1456pierre427 wants to merge 7 commits into
Conversation
Speculative decoding currently raises for models with recurrent
ArraysCache layers ("requires a trimmable prompt cache", ml-explore#980): the
recurrent state summarizes the whole past and cannot be trimmed.
This adds an exact, cheap rollback instead of forbidding the trim:
- ArraysCache: start/stop_speculation + record_rollback(T, fn, snapshot);
while speculating the cache is trimmable, and trim(n) restores the
entries via the recorded fn (or the snapshot for a full rollback).
- GDN layers (qwen3_5, qwen3_next): while speculating, record a rollback
closure over the exact per-token tensors the kernel consumed. Replaying
the recurrence from the pre-forward state over the first m stashed
inputs reproduces the post-m state bit-for-bit (causality: rows of a
chunk do not depend on later rows), and the conv state is a slice of
conv_input. No deepcopy, no re-forward, no extra weight streaming —
the rollback costs one tiny kernel launch per GDN layer.
- speculative_generate_step: allow models declaring
supports_speculative_rollback, and start/stop recording around the
decode loop (recording is off during prefill so no prompt-sized
tensors are retained).
Tests: tail-independence bit-exactness (same m-prefix, different tails
=> identical caches and logits after rollback), greedy speculative ==
greedy vanilla end-to-end with a mismatched draft (rollback every
round), and the unsupported-model error path.
Addresses findings from an adversarial audit of the first version: - ArraysCache keeps a small stack of rollback records (bounded by a 64-token window) instead of a single record, so multi-forward windows can be rewound — enabling HYBRID DRAFT models: a Qwen3.5 draft records one entry per T=1 draft step and its cache is now rewound exactly (the setup reported in ml-explore#1446 uses a hybrid draft). Partial rollbacks re-record the still-valid prefix. - trim() raises when asked to trim beyond the recorded window instead of silently clamping (a clamp would desync ArraysCache layers from KVCache layers, which trim the full amount). - speculative_generate_step: per-cache capability check (trimmable OR rollback-capable with model support) instead of a model-wide bypass; start/stop speculation on the draft cache too; n is initialized to num_draft at the top of each round so a mid-round exception or generator close makes the finally-block rewind a no-op rather than trimming with stale values — previously the rollback RuntimeError could mask the original exception. Tests: hybrid draft == vanilla greedy (rollback across stacked T=1 records), mid-round draft exception surfaces unmasked, trim beyond the window raises. 8/8 in this file; full suite failure set identical to main (pre-existing failures only). Real-weights e2e: dense draft 1.37x identical; hybrid draft pair 1.33x identical.
…d caches Fix: converted '-mtp' checkpoints (which keep mtp.* tensors with already-shifted norms) were double-shifted at load because mtp presence was treated as a raw-checkpoint signal; the signal is now the unsanitized conv1d layout only. Feature: mtp_num_hidden_layers builds an MTPModule (fc/pre_fc_norms/ full-attn decoder layer/norm, shared lm_head), self-dropped when a checkpoint has no mtp tensors. TextModel gains mtp_step() and rollback_speculative_cache(): KV caches trim; GatedDeltaNet states are rebuilt by replaying gdn_sink-captured verify inputs on the accepted prefix in one batched gated_delta_update (single-sequence scope). Replay pattern after mlx-vlm qwen3_5 / MTPLX (Apache-2.0). Measured on Qwen3.5-122B-A10B-oQ4-mtp (M5 Max): MTP chain k=4 = 1.36-2.11x vs plain greedy, output lossless modulo sub-ULP kernel ties. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Nice work — and a funny coincidence: we independently built the same capture-and-replay rollback this week, which I'd read as two votes for the approach being right. Your ULP-level near-tie observation matches our measurements exactly (top-2 gaps of 0.000–0.125 at divergence points, same class the unmodified batched-vs-sequential pipeline already shows). Since your PR explicitly defers generation-loop integration ("can follow separately if there's interest") — that part exists, tested: https://github.com/lBroth/mlx-lm/tree/gdn-exact-replay What it adds, complementary to this PR:
Your single-kernel batched replay (stacking layers along the head axis) is nicer than our per-layer replay loop — we'd adopt that with credit. If you're open to it, I think your MTP head + this integration layer make one complete story (#872 + #1446 + unblocking #1297's strategies on hybrids): either fold the branch under this PR, or keep them stacked — whatever you prefer. Happy to coordinate. |
|
Sure, go ahead and roll it all up into one.
…On Fri, Jul 3, 2026, at 5:56 PM, Broth wrote:
*lBroth* left a comment (ml-explore/mlx-lm#1456) <#1456 (comment)>
Nice work — and a funny coincidence: we independently built the same capture-and-replay rollback this week, which I'd read as two votes for the approach being right. Your ULP-level near-tie observation matches our measurements exactly (top-2 gaps of 0.000–0.125 at divergence points, same class the unmodified batched-vs-sequential pipeline already shows).
Since your PR explicitly defers generation-loop integration ("can follow separately if there's interest") — that part exists, tested: https://github.com/lBroth/mlx-lm/tree/gdn-exact-replay
What it adds, complementary to this PR:
• an `ArraysCache`-level protocol (`start/stop_speculation`, `record_rollback`, `trim` applies the recorded replay) instead of model-file-local plumbing — the layer records a closure over the exact kernel inputs, cache.py stays model-agnostic;
• wired into the existing `speculative_generate_step`, so `--draft-model` speculation works today on hybrids (including *hybrid draft models* — rollback records stack across the draft's T=1 steps; the Qwen3.6-target + Qwen3.5-draft pair from #1446 <#1446> runs end-to-end);
• covers `qwen3_next` as well as `qwen3_5*`;
• 8 tests, the key one being bit-exact tail-independence: two verify chunks sharing an m-token prefix with different tails must leave IDENTICAL caches and next-token logits after rollback-to-m — any off-by-one in the replay/slice leaks the tail. Validated on M2/M3/M5 Metal + CPU ops backend with real weights (greedy spec ≡ vanilla).
Your single-kernel batched replay (stacking layers along the head axis) is nicer than our per-layer replay loop — we'd adopt that with credit.
If you're open to it, I think your MTP head + this integration layer make one complete story (#872 <#872> + #1446 <#1446> + unblocking #1297 <#1297>'s strategies on hybrids): either fold the branch under this PR, or keep them stacked — whatever you prefer. Happy to coordinate.
—
Reply to this email directly, view it on GitHub <#1456?email_source=notifications&email_token=ABD3A353GGIAHRVBHXUPP3D5DATXRA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBXHE2TGNJVGIY2M4TFMFZW63VGMF2XI2DPOKSWK5TFNZ2KYZTPN52GK4S7MNWGSY3L#issuecomment-4879535521>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ABD3A36QY3YATA3WUNXTXLT5DATXRAVCNFSNUABFKJSXA33TNF2G64TZHM4TINRXGY3TQNZYHNEXG43VMU5TIOBQGQYTENZQHEY2C5QC>.
Triage notifications, keep track of coding agent tasks and review pull requests on the go with GitHub Mobile for iOS <https://github.com/notifications/mobile/ios/ABD3A36REYZBWSV2ZMFOBJD5DATXRA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBXHE2TGNJVGIY2M4TFMFZW63VGMF2XI2DPOKSWK5TFNZ2KUZTPN52GK4S7NFXXG> and Android <https://github.com/notifications/mobile/android/ABD3A35JO4BVSQFPQJAVYAD5DATXRA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTIOBXHE2TGNJVGIY2M4TFMFZW63VGMF2XI2DPOKSWK5TFNZ2K4ZTPN52GK4S7MFXGI4TPNFSA>. Download it today!
You are receiving this because you authored the thread.Message ID: ***@***.***>
|
…oth) Combines two independently-built halves into one story: - MTP head loading + norm-shift fix (our PR ml-explore#1456): un-strips mtp.* weights, builds MTPModule, mtp_step/make_mtp_cache; converted -mtp checkpoints load. - GDN exact-replay speculative rollback (lBroth/gdn-exact-replay): model- agnostic ArraysCache protocol (start/stop_speculation, record_rollback, trim applies the recorded replay) wired into speculative_generate_step, so --draft-model speculation works on hybrid (GatedDeltaNet) caches today, incl. hybrid draft models; covers qwen3_5* and qwen3_next. Our model-file-local gdn_sink + rollback_speculative_cache is DROPPED in favor of lBroth's cache-level protocol (cleaner: model files stay agnostic, the generation loop's existing trim_prompt_cache just works). Our single- kernel batched replay remains a credited follow-up optimization for the per-layer trim path. Tests: 8 exact-replay (incl. bit-exact tail-independence) + 3 MTP loading + norm-shift contract — all pass. Co-Authored-By: lBroth <noreply@github.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pierre Lamy <pierre@userid.org>
Extends the exact-rollback protocol to the OTHER non-trimmable cache: sliding-window RotatingKVCache. Once it wraps (offset >= max_size) the tokens a verify step pushes out of the window are gone, so stock trim can't rewind. Same protocol as ArraysCache: start/stop_speculation + record_rollback; update_and_fetch stashes the pre-forward window + new K/V while speculating, trim restores + re-appends the accepted prefix. No model-file hook (KV is not recurrent). is_trimmable() is True while speculating; gpt_oss declares supports_speculative_rollback=True, so speculative_generate_step accepts it unchanged -> --draft-model spec now works on gpt-oss-family sliding-window models, not just GDN hybrids. 4 tests incl. bit-exact tail-independence (two verify chunks sharing an m-prefix with different tails leave identical windows after rollback-to-m) + stacked draft records. 37/37 across prompt_cache + all rollback + MTP. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pierre Lamy <pierre@userid.org>
|
@lBroth — thanks, that's a clean design and the independent convergence is reassuring. I took you up on "fold it into one" and pushed a consolidation to this PR's head (force-update): What's now on the branch
All merged green: 8 exact-replay + 3 MTP loading + norm-shift contract, plus the item below. New: the same protocol for the other non-trimmable cache — Honest caveat we measured: external-draft speculation loses on a fast MoE target (gpt-oss-puzzle-88B: 0.42x, ~45% accept with a 20B draft) — the rollback is correct and general, it just doesn't pay at that operating point. Other adjacent things we have, mention-only (not proposing to include here):
Totally happy for you to drive the final shape — take any/all of this, or I'm glad to rebase onto your branch if you'd rather own it. Whatever lands it cleanest. |
|
Consolidation validated on our side (M5 Pro, macOS, mlx 0.31.2):
Two small notes for anyone testing with a fresh environment:
The |
|
As promised — cross-hardware validation of THIS consolidated head (
¹ Qwen3.5-9B-MLX-4bit target + Qwen3.5-0.8B-4bit draft — both hybrid GDN (the exact pair class from #1446), greedy, No regressions found anywhere. From our side this PR is validated and ready for review. |
… PR head, runs tests, optional real-weights e2e)
|
Following up on landing the speculative-rollback work — I think this PR is carrying too much for one review: it bundles the model-agnostic rollback cache-API (@lBroth's Proposal: split the rollback cache-API into its own focused PR and land it first, with the MTP head as a follow-up on top. Rationale:
Since the protocol is substantially @lBroth's work, I'd suggest you drive the rollback PR; I'll contribute the I'll then re-target this PR to just the MTP head on top of the landed API. Does that work for you? |
|
Quick follow-up to keep the split moving. On my side the MTP-head changes ( Since you've already validated the consolidated head across M2/M3/M5, easiest is probably to cut the rollback PR straight from there. If a pre-separated rollback-only branch off |
|
Opened the split as #1486 — the rollback cache-API alone (my three GDN/ArraysCache commits + your One note on your piece: I brought it in via cherry-pick, so the author stays you ( Then you re-target #1456 to just the MTP head on top of the landed API, as you suggested. Thanks for splitting it out — makes both much easier to review. |
|
Heads up for reviewers on the current state of the split: this PR is now scoped to the MTP-head loading only (the |
The MTP (nextn) head lives on the inner language_model/TextModel, but self_mtp_generate_step reads mtp/logits/mtp_step/make_mtp_cache from the top-level model. For wrapped checkpoints (qwen3_5_moe) that left the head unreachable. Delegate the four members to language_model so self-spec works on the top-level model. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
Pushed a small follow-up ( |
|
Requesting review of the rollback semantics and validation in this PR. Speculative decoding normally discards rejected work by truncating the KV cache. Qwen3.5 also carries recurrent GDN state, which cannot be restored by slicing it at a token boundary. This implementation reconstructs the retained state through replay, while loading the model’s MTP head and correcting an extra normalization that otherwise changes the proposal path. The current measurements show 83–100% acceptance and a 1.36–2.11× decoding speedup. Comparisons against ordinary target decoding remain lossless apart from equivalent low-precision ties. The points that seem most important to check are:
The broader question is how speculative decoding should represent rollback for architectures whose decoding state is recurrent rather than token-indexed. |
2 participants
What
Two related changes to
mlx_lm/models/qwen3_5.py:1. Bugfix: mlx-community
-mtpcheckpoints load as garbage today.sanitize()treatsmtp.*presence as a raw-HF-checkpoint signal andapplies the +1.0 zero-centered-norm shift. Converted checkpoints that
keep the MTP tensors (e.g.
mlx-community/Qwen3.5-122B-A10B-oQ4-mtp,converted with mlx-vlm) already have shifted norms — so every RMSNorm
gets double-shifted and the model emits fluent-speed gibberish. The
raw-checkpoint signal is now the unsanitized conv1d layout only, which
every real Qwen3.5 checkpoint has (kernel size 4). The existing
test_qwen3_5_family_convert_then_load_norm_not_shift_twicetest isupdated to use a kernel-4 config so the conv signal is exercised, and
extended with the kept-mtp regression case.
2. Feature: load the vendor-trained MTP head + support speculative
rollback on hybrid caches.
mtp_num_hidden_layersconfig field builds anMTPModule(
fc([norm(embed); norm(hidden)])→ full-attention decoder layer →norm → shared lm_head — the Qwen3-Next MTP structure). The module is
dropped in
sanitize()when a checkpoint has nomtp.*tensors, sonon-mtp conversions keep loading strictly. Conversion now preserves
and quantizes the MTP tensors instead of silently discarding them.
TextModel.mtp_step(hidden, tokens, mtp_cache)runs one MTP forward(batched catch-up or single-token recursive drafting).
TextModel.rollback_speculative_cache(caches, gdn_states, keep, block_size)rewinds after a verify forward: KV caches trim; the 36GatedDeltaNet recurrent states — which cannot trim — are rebuilt by
replaying the captured layer inputs (
gdn_sinkkwarg on the forwardpath) on the accepted prefix, all linear layers batched into a single
gated_delta_updatecall. Single-sequence scope; batched caches raisea clear error. (Replay pattern after mlx-vlm's qwen3_5 implementation
and MTPLX, both Apache-2.0.)
Why
Qwen3.5's hybrid linear-attention layers use
ArraysCache, which is nottrimmable — stock speculative decoding cannot rewind rejected tokens, so
these models are currently locked out of speculation in mlx-lm entirely,
and the shipped MTP head (the natural, draft-model-free way to speculate)
is discarded at load. This PR provides the loading + rollback primitives;
generation-loop integration can follow separately if there's interest.
Measured (M5 Max 128GB, Qwen3.5-122B-A10B-oQ4-mtp, greedy)
MTP self-speculative chain built on these primitives, vs interleaved
same-harness plain decode:
Depth-1 acceptance 83–100%. Losslessness: output ≡ plain greedy on
structured + prose probes, with divergences only at positions where
single-row recomputation shows a top-2 logit gap ≤ 1 bf16 ULP (0.000 /
0.125 observed) — the same batched-vs-sequential kernel tie behavior the
unmodified pipeline already exhibits. A forced-reject stress test
(every cycle rejects all drafts, exercising GDN replay on every cycle)
matches plain greedy under the same criterion.
Tests
tests/test_qwen3_5_mtp.py(new): rollback-replay ≡ sequential groundtruth for keep = 0..block; mtp_step shapes + recursive chaining +
trim; module dropped when checkpoint lacks mtp tensors; norm-shift
fires for raw layout and not for converted layout (both directions).
tests/test_models.py: norm-shift contract test updated as above.🤖 Generated with Claude Code