Mtp/qwen4exp nextn - #59
Merged
InfernalDread merged 6 commits intoSep 2, 2026
Merged
Conversation
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.
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.
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.
InfernalDread
merged commit Sep 2, 2026
0d08860
into
InfernalDread:qwen_38_flash_unsloth_mtp
10 of 32 checks passed
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
3 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.
No description provided.