models: Qwen3.8-Flash-Next MTP - #28243
Conversation
github-actions
Bot
added
model
| // target for a draft head that declares nextn_shared_target_tensors; must outlive this model | ||
| const struct llama_model * model_shared; |
There was a problem hiding this comment.
The llama_context_params already support ctx_other. Reuse that logic instead of adding this.
| // a captured graph hard-codes its shapes, so with one key per split an alternating shape | ||
| // (a speculative verify batch) resets warmup forever. O(1) on purpose: walking nodes undoes the | ||
| // point of a cuda graph. A shape this fails to separate re-captures as before, so it cannot regress. | ||
| static uint64_t ggml_cuda_graph_get_key(ggml_cgraph * cgraph) { | ||
| // unlike the previous key this dereferences nodes[0], so an empty graph is not safe here | ||
| if (cgraph->n_nodes <= 0) { | ||
| return 0; | ||
| } | ||
|
|
||
| uint64_t key = (uint64_t) (uintptr_t) cgraph->nodes[0]; | ||
|
|
||
| auto mix = [&key](uint64_t v) { | ||
| key = (key ^ v) * 0x100000001b3ull; | ||
| }; | ||
|
|
||
| mix(cgraph->n_nodes); | ||
|
|
||
| for (int d = 0; d < GGML_MAX_DIMS; d++) { | ||
| mix(cgraph->nodes[0]->ne[d]); | ||
| mix(cgraph->nodes[cgraph->n_nodes - 1]->ne[d]); | ||
| } | ||
|
|
||
| return key; |
There was a problem hiding this comment.
Better keep the CUDA changes for a follow-up PR.
|
@ggerganov thanks will check when I wake :)) |
|
Works great on dual consumer GPUs (5090 + 5060 Ti, 32 GB RAM): +25-40%, acceptance 0.77 9950X, 32 GB DDR5, RTX 5090 + RTX 5060 Ti (gen5 x8/x8), Ubuntu, CUDA 13.3, PR @ 2857e51.
Acceptance 0.767 (16849/21960), mean len 2.53. One minor issue I encountered: the fitter can't measure the shared head ( Thanks for all of your hard work! |
Validation: 2× RTX A6000 48GB (PCIe, no NVLink), Windows Server 2019, CUDA 12.6 (sm_86) — both heads draft at ~0.37 acceptance, no config beats baselineBuilt this PR @ Target: Arm 1 — shared Q4_K_M sidecar, Arm 2 — same but without Arm 3 — standalone Q8_0 grafted head (the drluoto file that gave 82–95% acceptance on #27836 @ Same draft file, same target, same prompts: 0.82–0.95 acceptance on #27836 vs ~0.37 on this branch — so something between Side note: draft compute buffer (2253 MiB) initially OOMs on device 1 next to the 90 GB target, recovered by |
The CUDA graph cache keyed captured graphs by the raw memory address of their first node (cgraph->nodes[0]). A captured graph hard-codes its shapes, but speculative decoding constantly alternates between different batch shapes (draft steps, verify batches, catch-up decodes) on the same context - when a new shape happens to reuse the same first- node address as a stale cached graph for a different shape, capture either reuses the wrong graph or thrashes, permanently resetting warmup instead of ever converging to steady-state replay. Hash node count and both endpoint tensors' shapes into the key instead (O(1) - walking all nodes would defeat the point of a CUDA graph), add LRU eviction capped at 64 graphs so the map can't grow unbounded now that distinct shapes get distinct entries. A shape this still fails to separate re-captures exactly as before, so it can't regress anything. Cherry-picked from ggml-org#28243 (open, unmerged), which found this while working on qwen4exp MTP performance - the effect is generic to any speculative-decoding workload on this fork, not qwen4exp-specific, so pulling in just this piece rather than the rest of that PR (which also reworks qwen4exp trunk/draft tensor sharing and doesn't fix the mixer/export-timing bugs already fixed on this branch). Verified: qwen4exp MTP speculative decoding still produces correct, byte-identical (temp 0) output after this change. Speed effect is hard to isolate cleanly from the dominant MoE-cache warmup effect already documented on this branch, but the fix is justified on its own correctness merits regardless of measured delta. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
| 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.
This triggers when loading with --fit on as the draft model is loaded independently as part of that. Everything still works though: it's later loaded again together with the target model.
Metal data point on this exact head (2857e51): dn=2 is greedy-identical, dn=5 is slower and divergesAll validations above are CUDA, so here is Apple Silicon. M5 Pro 64 GB, Metal, built at
Three things worth separating:
Happy to rerun on any later head; the harness is three commands. |
|
I have been trying this and the baseline from master actually has faster token gen than with MTP on, both shared and standalone. Using 5090+A6000(Ampere) + 128GB DDR5. |
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.
_QwenMtpMixin is not a ModelBase subclass, so it re-declares the attributes it reads off cls for the type checker. filter_tensors reads cls.mtp_shared_embd without a matching declaration, which ty reports as unresolved-attribute. The declaration is a bare annotation, matching no_mtp and mtp_only above it. That creates no class attribute, so it cannot shadow ModelBase.mtp_shared_embd even though the mixin precedes the model class in the MRO; a default value here would have. Assisted-by: Claude
A draft-only export declares the full block count but ships the MTP block alone, so the trunk tensors load as null and only the MTP graph is buildable. Context reservation builds the trunk graph, which walked those nulls and segfaulted. A shared-embedding draft is caught earlier by the borrow check, since it has no token_embd of its own. A self-contained draft keeps one, so it passed that check and reached here. Assisted-by: Claude
Reuses the existing llama_context_params::ctx_other plumbing instead of adding a model_shared load parameter and a nextn_shared_target_tensors metadata key. A draft-only export now simply omits token_embd/output, and the qwen4exp MTP graph resolves them against the target context at graph build, following dflash and gemma4-assistant. Loading such a file on its own reports that it needs -md. Assisted-by: Claude
Assisted-by: Claude
2857e51 to
2c96729
Compare
This comment was marked as spam.
This comment was marked as spam.
|
Pulled the latest commits on this PR and started getting this error in the logs and it seems MTP wasn't activated |
A qwen4exp draft that borrows the target's embeddings sets ctx_other but keeps its own memory, so it must be caught up and rolled back like any other draft. Treating it as memory-shared skipped the catch-up decode and placed every draft token at the same position, which the M-RoPE position check rejects. Assisted-by: Claude
|
@gabrielfreire Apologies it should be fixed now |
9 participants
Overview
Enables 1.3 to 2x faster MTP support for Qwen3.8-Flash-Next + shared MTP modules (re-uses embed_tokens to save disk space and VRAM / RAM). Builds on top of #27836
Additional information
Requirements