qwen4exp : port NextN/MTP speculative decoding to master (+50% decode at 70k) - #28104
Closed
JayToltTech wants to merge 6 commits into
Closed
qwen4exp : port NextN/MTP speculative decoding to master (+50% decode at 70k)#28104JayToltTech wants to merge 6 commits into
JayToltTech wants to merge 6 commits into
Conversation
The qwen4exp converter drops the MTP head, but community sidecar drafters now
exist (agentionai's Qwen3.8-Flash-Next-MTP-Q8_0: arch qwen4exp, block_count 49,
nextn_predict_layers 1, upstream tensor naming, hc_norm gammas verified
(1+w)-folded against the raw HF checkpoint by range read - exact +1.0 offset).
Follows the deepseek4 pattern: optional NEXTN_PREDICT_LAYERS KV with a
tensor-presence downgrade, mtp_only detection when the trunk is absent, trunk
tensors TENSOR_NOT_REQUIRED in that case, NextN tensors TENSOR_SKIP unless the
context requests them.
qwen4exp-specific differences from deepseek4, both dictated by the file format:
- the NextN block is always a full-attention QSA layer (gated attention, own
nh=4 indexer, full 512-expert MoE, both HC pairs); is_recr() is derived from
full_attention_interval and would misclassify blk.48 as linear attention, so
the MTP layer forces the full-attention branch.
- nextn.hnorm is hc-space ({hc_dim}, 10240), not {n_embd}: the draft head
consumes the target's 4-stream hyper-connection state.
A standalone load of a sidecar now aborts with a clear message instead of
segfaulting in the first hc_mix (the trunk tensors are null by design; the
mainline graph cannot be built from a drafter file).
Draft graph and --spec-type draft-mtp wiring are the follow-up commit; this one
is loader-only and mainline GGUFs are unaffected (KV absent -> n_layer_nextn 0,
identical load path).
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
(cherry picked from commit 23d8b2c)
(cherry picked from commit 3543908)
…o end Second half of the MTP wiring (loader was the previous commit). Verified end to end on the 64 GB box: Q4_K_M target + agentionai's community MTP sidecar, --spec-type draft-mtp, coherent temp-0 generation through the Vulkan lane. - graph_mtp: [enorm(embd(tok)) repeated across streams ; grouped-RMS hnorm(h)] -> eh_proj per stream -> one HC-wrapped full-attention block -> the shared head mixer (which doubles as the output norm). h arrives hc-space: the mainline graph exports the 4-stream residual under cparams.embeddings_nextn. - the MTP context holds a PLAIN attention cache over the nextn layer(s), dense, no recurrent state - the deepseek32 pattern. Discovered the hard way: a hybrid memory with an empty recurrent layer set fails its buffer allocation. - exports hand the scheduler REAL nodes: a naked reshape view gets no backend assignment and the h_nextn extraction asserts. - the draft runs dense attention (QSA gate is null-safe now); the reference precedent is deepseek32, whose MTP head also runs dense. Correctness only at this point: no acceptance-rate or throughput claims - the smoke ran under heavy memory pressure and the interactive-trap log spin. Measured cleanly it may or may not pay; that is the next session's A/B. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> (cherry picked from commit b5b7c8c) (cherry picked from commit 39817c4)
This reverts commit a17e843 (the v0.6.9 revert). The stall that motivated the revert was not in the state-save path: it was a replay re-verification livelock in the server's full-checkpoint rollback, fixed in the previous commit. With that in place, full checkpoints restore token-exact MTP rollback on recurrent targets again. Verified on gfx1151 Vulkan (Qwen3.6-35B-A3B Native-MTP Q6_K, greedy): the v0.6.8 repro completes 800 tokens in 12 s, a 3000-token request runs through the former stall horizon to EOS at 62 tok/s, and repeat runs are bit-identical. CPU spec-vs-nospec is token-exact through the canonical 192-token matrix. Assisted-by: Claude Fable 5 (cherry picked from commit f25eefe)
…store With full-checkpoint rollback, a partial draft acceptance restores the pre-round state and re-decodes the accepted tokens to rebuild it. The replay went through the same verification as a fresh draft. On backends where logits change with batch shape or memory layout (Vulkan), that re-verification can reject a token the original verification accepted; the rejection restores the same checkpoint and replays again, and the slot loops on one position without emitting anything. qwen35moe with --spec-type draft-mtp stalled this way a few hundred tokens into long generations (the v0.6.8 MTP hang): the loop repeated "accepted 2/3, restore at pos 995" every 27 ms with the GPU at 90 percent. Accept the replayed tokens without re-verifying and sample only the continuation from the final position. The replayed prefix was accepted by the verification that triggered the restore; the replay exists to rebuild state. On backends with batch-shape invariant logits the re-verification always agreed, so behavior there is unchanged (verified bit-identical on CPU with and without this change, 800-token greedy pair, 118 restore rounds). Assisted-by: Claude Fable 5 (cherry picked from commit 9c5d899)
The MTP drafter consumes one hyper-connection row per prompt token: its prefill path memcpys n_tokens-1 rows straight out of the nextn buffer. The mainline graph dropped the non-output rows inside the last layer, before the export, so the exported residual carried only n_outputs rows and the reader walked off the end -- GGML_ASSERT(offset + size <= ggml_nbytes(tensor)) on the first decode. Skip the early gather while a drafter is attached and re-apply it immediately after the export, so the drafter sees every row and the output head still sees only the output rows. Costs one extra layer of compute on dropped rows in the last layer, which is why the gather sits where it does when no drafter is attached. Setting the target context masked instead does not work: the flag also gates a second gather in the export, and the drafter has no output-indexed path.
qwen4exp is a recurrent hybrid, so the target context cannot partially seq_rm and the server classifies it SEQ_RM_TYPE_FULL. Every speculative round then takes a full recurrent-state checkpoint, and every rejected draft restores one. Through the host path that means serializing each GDN layer conv+state row, the 4-stream hyper-connection residual and the PLE history into a host vector with one synchronous backend read per tensor, then pushing it all back. The cost is flat in context and swamps everything else: with the MTP drafter attached, decode ran 201 ms/token against 29 ms/token on a tree that keeps the state on device, and a 120-token generation spent roughly 600 ms of each 825 ms round in checkpoint traffic. It also inflates the reported prompt eval time, which runs to first token and so absorbs the first save: 827 ms for an 11-token prompt. Request ON_DEVICE at the six speculative checkpoint sites so the state stays in device buffers. The library already implements this path; only the server never asked for it. Measured on gfx1151 Vulkan, Qwen3.8-Flash-Next UD-Q4_K_XL with a Q4_K_M MTP head: 4.77 -> 25.83 tok/s at short context, 4.33 -> 16.08 at 70k, draft acceptance unchanged at 70-80 percent. The prompt-cache checkpoints are deliberately left host-resident: they retain several historical states rather than one live round. Diagnosed by Claude Fable 5. The ON_DEVICE flag and the mechanism come from Gaetan Puleo (c8b681b6f), carried in Nathanw1014/llama.cpp as 08a3255.
Author
|
Converted to draft, and holding pending @Nathanw1014's permission. A maintainer rightly pointed out on #28061 that I should not submit someone else's work without asking first. That applies here too: four of the six commits are @Nathanw1014's, and I filed this without asking. My mistake, and I have closed the other PR. @Nathanw1014 — this stays a draft until you say otherwise. Three options, all fine by me: you take it over and submit it yourself (I will close this), you give the go-ahead and I take it out of draft, or you would rather it not be upstreamed at all and I close it. No obligation either way. Two of the commits here are mine and are not affected by that — I am noting them separately in case they are useful to anyone independently of the rest:
@llbzow — your dual-A6000 report of "acceptance matches, throughput regresses ~2.2x" on CUDA may be the same cause. The checkpoint traffic I measured is a per-round host round-trip of the full recurrent state, which is not Vulkan-specific; I only measured it on Vulkan because that is the hardware I have. If you are able to test |
|
Hi @JayToltTech, thanks for your contribution! Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:
Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below. |
ggml-gh-bot
Bot
added
the
draft
sodre90
added a commit
to sodre90/llama.cpp
that referenced
this pull request
Aug 31, 2026
Superseded by PR ggml-org#28104's port, which additionally fixes the checkpoint round-trip and replay-livelock issues this graph-only version was blocked on (get_idx() segfault under the reworked hybrid-memory graph). This reverts commit 482228e.
Author
|
@Nathanw1014 closing this out rather than leave it hanging on you. I should have asked before filing it, not after. The MTP work is yours. Take it upstream whenever suits, or don't. My two commits on the branch (the h_nextn gather fix, and the on-device checkpoint flag) are free for you to use or drop as you see fit, no credit needed. Branch is still on my fork if it's any use: |
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
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.
Ports qwen4exp NextN/MTP speculative decoding onto current master, working and fast.
Related: #27836 (an earlier MTP attempt, different implementation), #28060 / #28061 (a Vulkan livelock found on the way, submitted standalone).
Result
AMD Strix Halo (gfx1151), Vulkan/RADV, Qwen3.8-Flash-Next UD-Q4_K_XL, Q4_K_M MTP head, temp 0:
+50% decode at 70k, draft acceptance 70-80%. The gain grows with context, which is the point: qwen4exp runs full attention on 12 of 48 layers at ~24 KiB/token, so verifying several drafted tokens per forward pass amortises the per-step cost that otherwise dominates at depth.
At temp 1.0 / top-p 0.95 / top-k 20 on the reference tree the same approach holds up across generation lengths (50/100/200/400 tokens: 18.32 / 15.01 / 15.92 / 15.61 tok/s), so this is not a temp-0-only artifact.
Commits
Four from @Nathanw1014's
strix-halo-vulkan, authorship preserved — the MTP implementation is entirely theirs:354390810load NextN/MTP draft tensors (sidecar and in-file)39817c476NextN/MTP draft graph and contextf25eefeaffull checkpoints for MTP rollback9c5d899ffreplay livelock fix (standalone in server : do not re-verify replayed draft tokens after a checkpoint restore #28061)Two of mine, both required to make it work on master:
qwen4exp : defer the output-row gather past the h_nextn export— the drafter consumes one hyper-connection row per prompt token (common_speculative_impl_draft_mtp::process()memcpysn_tokens-1rows out of the nextn buffer), but master's graph drops non-output rows inside the last layer, before the export. The reader then walks off the end:I skip the early gather while a drafter is attached and re-apply it right after the export. Marking the target context masked instead does not work — the flag also gates a second gather in the export and the drafter has no output-indexed path; that route segfaults in the memcpy.
server : keep speculative checkpoints on device— the substantive one. qwen4exp is a recurrent hybrid, so the target cannot partiallyseq_rmand the server takes a full recurrent-state checkpoint every speculative round, restoring it on every rejection. Through the host path that serialises every GDN layer conv+state row, the 4-stream residual and the PLE history with one synchronous backend read per tensor. Cost is flat in context and swamps everything: 201 ms/token decode, roughly 600 ms of each 825 ms round. The six speculative call sites now requestLLAMA_STATE_SEQ_FLAGS_ON_DEVICE; the library already implemented that path, the server just never asked. 4.77 -> 25.83 tok/s short, 4.33 -> 16.08 at 70k.Prompt-cache checkpoints are deliberately left host-resident — they retain several historical states rather than one live round.
The ON_DEVICE flag and mechanism are Gaetan Puleo's (
c8b681b6f, carried as08a3255). Diagnosis of which call sites mattered: Claude Fable 5.Note for #27836
That PR reports ~4.7 tok/s at 90% acceptance — the same profile this had before the checkpoint fix. If it is hitting the same host-path checkpointing, the fix may be independent of its graph work.
Caveats
strix-halo-vulkanis likely the PLE host-side gather (an in-graphggml_get_rowsover a CPU-resident table, splitting the graph GPU->CPU->GPU per ubatch); not addressed here.dzannotti/Qwen3.8-Flash-Next-MTP-GGUF, Q4_K_M); the converter does not emit one.