qwen3_next: load in-checkpoint MTP head for self-speculation#1532
Open
pierre427 wants to merge 8 commits into
Open
qwen3_next: load in-checkpoint MTP head for self-speculation#1532pierre427 wants to merge 8 commits into
pierre427 wants to merge 8 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>
…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>
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>
Stacks on ml-explore#1456 (qwen3_5 in-checkpoint MTP head + self-speculation machinery + qwen3_next speculative-rollback support). The base already provides qwen3_next's GDN rollback; this adds only the MTP (nextn) head. Mirrors the qwen3_5 MTP head from ml-explore#1456: adds ModelArgs.mtp_num_hidden_layers, a Qwen3NextMTP module (fc over [norm(embed(t_{p+1})); norm(hidden_p)] -> one full-attention decoder layer -> norm -> trunk lm_head), and exposes the exact same self-spec interface (mtp attribute, logits, make_mtp_cache, mtp_step) so qwen3_next self-speculation runs through the identical machinery. sanitize keeps mtp.* tensors only when both the checkpoint and the built module have them, dropping the module otherwise for strict loading. No-op for the common MLX repacks that strip the head. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
qwen3_next: load in-checkpoint MTP head for self-speculation
Stacked on #1456 (qwen3_5 in-checkpoint MTP head + self-speculation
machinery + qwen3_next speculative-rollback support). Please review/merge
#1456 first.
What
Ports the Qwen3-Next multi-token-prediction (nextn) head onto the model so a
trained or grafted
-mtpqwen3_next checkpoint can drive self-speculativedecoding with no external draft model — the same capability #1456 already
gives qwen3_5.
The base (#1456) already provides qwen3_next's speculative rollback support
(
record_rollbackon the GDNArraysCache,supports_speculative_rollback = True), which is what makes the hybrid (recurrent + full-attention) cachetrimmable during speculation. What it lacked was the MTP head itself. This PR
adds only that.
How it mirrors the qwen3_5 MTP head (#1456)
The new
Qwen3NextMTPmodule and theModelinterface are a direct analog ofMTPModule/TextModelinqwen3_5.py:ModelArgs.mtp_num_hidden_layers: int = 0— 0 = no head (default; communityMLX repacks strip it);
> 0builds the head.Qwen3NextMTP(nn.Module):fc([norm(embed(t_{p+1})); norm(hidden_p)])->one full-attention decoder layer (
layer_idx = full_attention_interval - 1sois_linear=False, own KV cache) ->norm-> the trunk'slm_head.Predicts token p+2 from the trunk hidden at p and the committed token p+1.
Model.__init__buildsself.mtponly whenmtp_num_hidden_layers > 0.mtpattribute,logits(hidden),make_mtp_cache(), andmtp_step(hidden, tokens, mtp_cache) -> (logits, post_norm_hidden). The self-speculation machineryconsumes qwen3_next's head through the identical attribute/method surface it
uses for qwen3_5 — nothing model-specific.
sanitize: keepmtp.*tensors only if the checkpoint carries them andthe module was built (
mtp_num_hidden_layers > 0); otherwise drop both thetensors and the module so strict loading stays consistent. No-op for the
common MLX repacks that ship no MTP head. (Runs before the existing
expert-stacking / norm-shift path, which is untouched.)
Unrelated changes on our downstream fork (e.g.
RotatingKVCache-based windowedKV in
make_cache) are intentionally not included here — this PR is onlythe MTP head.
Test
tests/test_qwen3_next_mtp.pymirrorstests/test_qwen3_5_mtp.py: builds atiny qwen3_next config with
mtp_num_hidden_layers=1, runs the trunk to getpost-final-norm hiddens, then exercises the MTP head forward — asserts
next-token logits
[1, 7, 64], chained (recursive) MTP step, KV-cacheoffset/trim bookkeeping, and that the module is dropped when a checkpoint has
no
mtp.*tensors (and kept when it does).Verification
pytest tests/test_qwen3_next_mtp.py -q— pass (3)pytest tests/test_qwen3_5_mtp.py -q— pass (3), no regressionpy_compileclean;black --check+isort --profile black --checkclean