MTP for Qwen3.8-Flash-Next - #144
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b8e89f2f36
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| for (const auto & [n, t] : llama_internal_get_tensor_map(model_shared)) { | ||
| if (n == name) { | ||
| src = t; |
There was a problem hiding this comment.
When --mtp-shared-embd is used with a target that ties its output head to the token embeddings, this exact-name scan cannot find output.weight: target loading represents the tied head by assigning model.output to the tensor named token_embd.weight, so tensors_by_name contains no output entry. The loader therefore throws before the architecture's normal tied-head fallback, making shared sidecars unusable with tied-output targets; resolve LLM_TENSOR_OUTPUT through the target model's output pointer or token-embedding fallback.
Useful? React with 👍 / 👎.
| if (model_shared == nullptr) { | ||
| throw std::runtime_error(format("%s: this model is a draft head without its own '%s'; " | ||
| "load it as a draft of its target model, not on its own", __func__, name.c_str())); |
There was a problem hiding this comment.
With the default --fit and a --mtp-shared-embd sidecar, the extra draft is measured before the target model is loaded, so its copied parameters have no model_shared and this branch throws on the first omitted tensor. common_fit_params catches that measurement failure and fits the main model alone, excluding the draft block and its context memory, so the chosen layer/context settings can exceed available memory during the subsequent real draft load; the no-allocation measurement path needs to model borrowed tensors without requiring a live target.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
I think this is exactly what the bot remarks.
$ llama-cli -m models/unsloth/Qwen3.8-Flash-Next-GGUF/Qwen3.8-Flash-Next-UD-Q4_K_XL-00001-of-00004.gguf --jinja -c 131072 -ngl 999 -ncmoe 40 --load-mode mmap -ub 4096 -b 8192 --reasoning-effort medium --model-draft models/unsloth/Qwen3.8-Flash-Next-GGUF/mtp-Qwen3.8-Flash-Next-shared-Q4_K_M.gguf --spec-type ngram-mod,draft-mtp --spec-draft-n-max 4
Loading model... |0.00.787.728 E llama_model_load: error loading model: borrow_shared_tensor: this model is a draft head without its own 'token_embd.weight'; load it as a draft of its target model, not on its own
0.00.787.734 E llama_model_load_from_file_impl: failed to load model
It works regardless. And with --fit off this warning is not printed.
b8e89f2 to
4818b47
Compare
Adds the MTP head's own hyper-connection mixer tensor names and lists the NextN tensors under the qwen4exp architecture.
Adds --spec-type draft-mtp support for Qwen3.8-Flash-Next. The MTP head folds the next token's embedding into the trunk's wide hyper-connection residual, runs one trunk-style block (dense attention + MoE) over it, and collapses the result with its own mixer before reusing the trunk's LM head. - read nextn_predict_layers so n_layer() excludes the MTP block - load the trailing block through the existing trunk path: is_recr() and is_ple() are already false past the trunk, so it needs no special casing - eh_proj fuses the checkpoint's fc_embedding and fc_hidden side by side, so one matmul computes fc_embedding@e + fc_hidden@h - the head carries its own hyper-connection mixer, mirroring the trunk's hc_head_*, which stands in for the output norm qwen4exp does not have - export the wide pre-collapse residual as t_h_nextn from both graphs, so the driver can feed it back for the next draft step - route MTP contexts to a plain KV cache filtered to the trailing layer The draft block attends densely for now: the trunk's QSA only prunes context past a 2048-token budget, so dense is a numerical superset and drafts are verified either way. Indexer tensors are still loaded.
The MTP block is one trunk-shaped block (dense attention + MoE wrapped in hyper-connections) plus a head-level combiner, so once _QwenMtpMixin renames mtp.layers.0.* to the trailing block index its tensors ride the existing qwen4exp mappings unchanged. Two head-level pieces need handling: - fc_embedding and fc_hidden fuse into the eh_proj the shared NextN code expects, since W_e@e + W_h@h == [W_e|W_h] @ concat(e, h) - mtp.hyper_connection_mixer.* is the head's own copy of the trunk's hc_head_* output mixer, unindexed in the checkpoint and per-block in the GGUF compress_ratios is read with length block_count, so it gains a trailing 0 for the MTP block, which attends densely. --no-nextn drops the head; --mtp exports it on its own.
4818b47 to
6fc8df1
Compare
* unsloth: pin ggml-org#27941, #152 and #154 ggml-org#27941 at 02eb201, no longer a draft and MERGEABLE upstream. It fixes four qwen4exp correctness defects and is the only account of the Flash-Next reporter who is on Vulkan only, where the allocator variable is never set. The previous pin PR #148 carried the stale 8161d11 and is closed. #152 at 258345e and #154 at 31e432e both merge onto b10708 on their own. #137 is already pinned at 4e1865e, which is still its head, so no change. #142 and #144 are NOT pinned here: both predate an upstream refactor that folded lazy_mode and model_shared into a lazy struct, so they conflict on src/llama.cpp and src/llama-model-loader.cpp against b10708. b10708 has 'ml.lazy.mode = params.lazy_mode' where those branches still write 'ml.lazy_mode' and 'ml.model_shared'. They need rebasing onto current upstream before they can be pinned; pinning them now would fail the resolve. * unsloth: repin ggml-org#27941 to 6b2b85c The PR moved on at 10:44Z, after 02eb201 was pinned. Verified to merge onto the current base tag: 6 files changed, 480 insertions, 138 deletions. * unsloth: pin #144, rebased, and leave #142 out as contained in it #144 rebased onto b10709 at 6fc8df1. Two of its nine commits were dropped rather than carried: 'qwen4exp: fix QSA correctness defects and harden metadata loading' and the comment tidy on top of it. That work is what ggml-org#27941 supersedes, and keeping both copies is what made the two pins conflict in llama-kv-cells.h and llama-memory-hybrid-idx.cpp. Listed after ggml-org#27941 so the upstream version of that work lands first. #142 is deliberately NOT pinned: its single commit is byte-identical to #144's c7bd6f2 apart from the lazy API spelling, so #144 already contains it and pinning both would apply the same change twice. Verified: b10709 + ggml-org#27941 + #144 merges clean, 24 files, +612/-94. * unsloth: repin ggml-org#27754 to 949f7ef The PR moved on at 10:39Z; 5796547 was pinned by #159 earlier today. Verified to merge onto b10709: 44 files changed, 2673 insertions, 38 deletions. Every entry in the set is now at its PR's current head. --------- Co-authored-by: danielhanchen <elliegouldingstuff@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6fc8df13a1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // LLM_GRAPH_TYPE_DECODER_MTP draft head for qwen4exp. | ||
| // | ||
| // The head folds the next token's embedding into the trunk's wide hyper-connection residual, | ||
| // runs one trunk-style block over it, and collapses the result with its own mixer before | ||
| // reusing the trunk's LM head. The wide post-block residual is exported as t_h_nextn so the |
There was a problem hiding this comment.
Condense this 11-line overview to the non-obvious dense-attention limitation and TODO. It hard-wraps several sentences across comment lines and substantially exceeds the repository's usual 1-2-line limit, adding review and maintenance overhead without clarifying the constructor beyond what the code already shows.
AGENTS.md reference: AGENTS.md:L74-L78
Useful? React with 👍 / 👎.
A NextN/MTP draft exported with --mtp carries the token embeddings, output norm and lm head so it can be loaded as a standalone model. For every current sidecar those three tensors are most of the file: ggml-org/Qwen3.8-27B-GGUF mtp-Qwen3.8-27B-Q4_0.gguf is 1.565 GiB, of which 1.332 GiB (85%) is the copy, against 0.223 GiB for the MTP block itself. Add an opt-in --mtp-shared-embd that leaves them out and marks the file with nextn_shared_target_tensors. The loader then resolves those names against the already loaded target model. The graph side needs no change: the nextn blocks of twelve archs already fall back to model.tok_embd and model.output. The borrow is gated on the new key, so a sidecar published before this change cannot reach it and keeps its current behaviour. Shapes are checked against the target and a mismatch is refused, as is loading such a file on its own.
qwen4exp was missing from llm_arch_supports_rs_rollback, so n_rs_seq was clamped to 0 and every speculative round took the full-checkpoint path: 158 saves at 8.86 ms and 98 restores at 13.31 ms of 112.57 MiB each, about 10.4 ms per round. Adding the arch on its own is not enough and breaks generation after four tokens, because build_conv_state_at always wrote slot 0. The shared delta-net helper writes n_rs_seq + 1 snapshots at (s_slot * mem_size + kv_head); the qwen4exp helper had no such branch, so state groups 1..n_rs_seq were never populated and a rollback read stale rows. The helper is used twice here, for the GDN conv state and for the PLE conv state, which is why the defect was not visible from either call site alone. At n_rs_seq == 0 the loop runs once at s_slot 0 and is byte for byte the previous single write, so the non-speculative path is unchanged by construction. Confirmed by md5 of greedy output against the parent commit. Measured on Qwen3.8-Flash-Next UD-IQ1_S, one B200, 512 tokens, temp 0, seed 1, three reps: MTP goes from 0.671x to 1.294x against the same model with the draft head disabled. Draft acceptance rises from 54.6 to 60.8 percent, which is the sensitive probe here, since speculative decoding cannot change the output and correctly restored state can only show up as acceptance.
The graph cache is keyed on cgraph->nodes[0] alone, so two evaluations that share a first node but differ in shape collide on one entry. Warmup needs two consecutive calls with unchanged node properties, so a workload whose batch shape varies resets warmup on nearly every call and falls back to eager launch. Speculative decoding is exactly that workload. The qwen4exp verify batch is distributed 2:13 percent, 3:11 percent, 4:75 percent as the accepted count varies, where qwen35 sits at 4:98 percent and is effectively constant. Host launch time for the qwen4exp target decode was 1.52 ms with the draft head disabled and 12.35 ms with it enabled, while GPU time was unchanged, so the regression was entirely host side. The key now mixes the first node, the last node and the node count. This is O(1) rather than a walk over every node: the existing uid early return fires on 127 of 128 decodes, so the hot path must not touch node data. An earlier all-nodes hash reintroduced exactly the per-node walk a CUDA graph exists to avoid. Measured overhead against the previous key is 0.2 to 0.6 percent, with both variants built into one binary to avoid comparing across runs. Capture churn on Qwen3.8-27B UD-Q2_K_XL drops from 52 captures and 50 destroys to 4 and 0, with identical output md5 and an unchanged speculative ratio. Across 14 distinct prefill shapes the cache instantiates 16 entries against 14 before, with no destroys and no growth, and is capped at 64 by LRU on top of the existing sweep. test-backend-ops passes 13646 of 13646 on CUDA0, and Llama-3.2-1B-Instruct Q8_0 is byte identical with no throughput change.
6fc8df1 to
586b15e
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 586b15ef8e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // not counted in n_created or size_data: the tensor is not in this file and is neither | ||
| // allocated nor freed here | ||
| return src; |
There was a problem hiding this comment.
When a shared sidecar is loaded with --device-draft on a different device from the target, this returns the target's preallocated tensor without making its backend available to the draft context. The CLI explicitly supports an independent draft device list (common/arg.cpp:4213-4217), while the context scheduler registers only model.devices of the draft (src/llama-context.cpp:330-357); if none supports the target tensor's buffer type, scheduler assignment aborts on the preallocated leaf. Register the target backend in the draft scheduler or create a draft-owned copy of each borrowed tensor.
Useful? React with 👍 / 👎.
|
I landed the same rollback fix upstream as ggml-org#28123, so c6e318e is now in master and you can drop it. I'm running some tests now Quick Opus 5 agentic loop result : Ran it on an RTX PRO 6000 (96 GB), target UD-Q4_K_XL, one slot, greedy, --ctx-checkpoints 0, 4 prompts of 256 tokens each. Better than I expected, and higher acceptance than your B200 numbers. One thing that does not match your conclusion though: the non-shared Q8_0 is 8 percent faster than the shared one here, 191 vs 176, with identical acceptance and mean accepted length. Same weights, so it looks like the borrow path costs something rather than the quant tier. The other thing, and this one worries me more: greedy output is not identical with the head on. 3 of my 4 prompts diverge, and all four sidecars diverge at the same byte, bf16 included, so it is not draft quality. Control passes, two MTP-off runs of the same binary are byte identical 4/4, so it is not my harness either. Text stays coherent on both sides, it just picks a different continuation. My guess is that turning on draft-mtp makes n_rs_seq non-zero, which switches build_conv_state_at from one write to K, so the target's own recurrent state differs from the non-speculative path. That is the area c6e318e and ggml-org#28123 touch. Happy to dig if you want, I have the setup warm. |
The tree is that of hetero-qwen38-mtp: ggml-org master + unslothai/llama.cpp#144 + ggml-org#28118 + the halo-hybrid patch set ported. main's history stays reachable; its tip is tagged main-pre-mtp. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0113rfHcRv4MecQaiiXKphq5
4 participants
Adds the NextN/MTP draft head for Qwen3.8-Flash-Next (
qwen4exp), plus the two fixes that turn it from a large regression into a speed win.Generation speed
Before this PR, enabling the draft head made generation slower at every setting tried. One B200, 512 tokens, temp 0, seed 1, 3 reps,
-fa on:By target model, with the Q8_0 shared sidecar:
The qwen35 control is what says this is now in the expected band rather than an outlier. Draft acceptance is 60.8 percent, mean accepted length 2.98.
Scope this to low-concurrency serving. Driven through the same client as vLLM, at concurrency 1 both stacks now extract the same speculative gain (llama.cpp 1.331x TPOT, vLLM 1.337x), and the decode-only gap to vLLM narrows from 3.45x to 1.54x. At concurrency 8 the sidecar is a net loss: 0.871x on UD-IQ1_S and 0.806x on UD-Q4_K_XL, where vLLM's gain only shrinks to 1.066x. That regression is not yet profiled. Note separately that llama.cpp's throughput collapse at concurrency is upstream behaviour rather than this architecture's: the qwen35 27B control collapses identically, 3.31x TPOT growth from 1 to 8 slots against qwen4exp's 3.29x.
Acceptance rate depends on the sampler, so any speculative figure needs its protocol stated. The same Q8_0 sidecar against UD-Q4_K_XL measures 1.667x greedy, 1.566x with
--temperature 0pinned through the benchmark client, and 1.183x at that client's default sampler, where logged acceptance falls to 0.387 to 0.504 with mean length 2.16 to 2.50 against 70.67 percent and 3.12 greedy. All three are correct for the same artifact.How it got there:
The two fixes
The regression was entirely host side: the target verify decode went from 1.52 ms of launch time to 12.35 ms while GPU time stayed flat. The draft itself was never the problem at 0.77 ms per call.
ggml-cuda: key the CUDA graph cache by shape. The cache keys oncgraph->nodes[0]alone, and warmup needs two consecutive calls with unchanged node properties. The qwen4exp verify batch varies (2:13 percent, 3:11 percent, 4:75 percent) where qwen35 is effectively constant at 4:98 percent, so warmup reset repeatedly and fell back to eager launch. The key now mixes first node, last node and node count. O(1) on purpose: the uid early return fires on 127 of 128 decodes, so the hot path must not walk node data. Overhead against the previous key is 0.2 to 0.6 percent.qwen4exp: cover the recurrent conv state for rollback.qwen4expwas missing fromllm_arch_supports_rs_rollback, son_rs_seqwas clamped to 0 and every round took the full-checkpoint path, about 10.4 ms per round. Adding the arch alone breaks generation after four tokens, becausebuild_conv_state_atalways wrote slot 0 where the shared delta-net helper writesn_rs_seq + 1snapshots. State groups1..n_rs_seqwere never populated and rollback read stale rows. The helper is used twice here (GDN conv state and PLE conv state), which is why it was not visible from either call site. Atn_rs_seq == 0the loop is byte for byte the previous single write, so the non-speculative path is unchanged by construction.Sidecar artifact
Q8_0 in the shared/borrow form, 2.595 GiB. It is the speed optimum rather than a compromise: bf16 is 1.88x larger and also slower (the LM head is cheaper at 8 bits), and every tier below Q8_0 loses speed. Acceptance across tiers, ~10,300 draft events per arm, plus or minus 0.91 percent: bf16 66.50, Q8_0 66.14, Q6_K 65.85, Q5_K_M 65.19, Q4_K_M 64.35, Q3_K_M 63.22, Q2_K 54.08.
An MTP-aware imatrix is not worth building: the bf16 to round-to-nearest gap bounds what one could recover, and at Q8_0 that bound is 0.36 pp against UD-IQ1_S and 0.00 pp against UD-Q4_K_XL.
Commits
Based on
qwen4exp/qsa-correctness-fixes(#143). The first three commits are Ryan Monsurate's upstream work from ggml-org#27836, cherry-picked with authorship intact. The borrow commit is #142, included because the draft-only load path depends on it; it drops out on rebase once #142 lands.Validation
test-backend-ops13646 of 13646 on CUDA0;test-llama-archs -a qwen4expCUDA 8.66e-08, CPU 0.00e+00, roundtrip OK.Open items
--kv-unifiedwith four or more concurrent sequences andn_kvabove 2051, a sequence holding fewer thancompress_ratiocells owns no block, its whole bias row is-inf, and the softmax produces NaN that reaches every other sequence through the dead block. Bisected to a one commit A/B: the parent5d6fc1862passes 4/4 in two reps,713156608fails 0/4 and 1/4, with non-unified and shallow controls passing in every arm. It is qwen4exp only, since only that arch has the QSA indexer. The same defect in mild form also drops each sequence's last one to three tokens in ordinary single-slot long-context use. Fix is 11 lines and validated; it belongs on the base, not here.--kv-unifiedwith MTP off, which is the base defect above.ggml-cudacommit affects graph keying for every model, not just this one, and is a reasonable candidate to split out for upstreaming.llama-serverhere needs--ctx-checkpoints 0; checkpoint recompute perturbs greedy output on its own, on qwen35 as well as qwen4exp.