Repository navigation
qwen4exp: run MTP heads that share the target's token_embd and output - #240
danielhanchen wants to merge 2 commits into
Conversation
MTP heads exported with nextn_shared_target_tensors carry no token_embd or output, which load_arch_tensors required, so mainline refused them. For such a head created as an MTP draft, take both from the target the draft context is created against. ctx_other itself stays unset: for draft-mtp it means the draft shares the target's memory, which this head does not. Any other use of such a head fails with a clear error rather than a null dereference. Heads converted before upstream MTP support also mark the MTP block dense (compress ratio 0). graph_mtp built the k-pool input regardless, but the draft's indexer cache then holds no layer, so the input was never allocated and the first decode asserted in set_input. Build it only for a QSA MTP block; build_layer_attn already treats ratio 0 as dense.
|
You have reached your Codex usage limits for security reviews. Please try again later. |
![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: 36175a8325
ℹ️ 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".
| const llama_model & emb = model.tok_embd ? model : *cparams.model_tgt; | ||
| const llama_model & head = model.output ? model : *cparams.model_tgt; |
There was a problem hiding this comment.
When a shared head is run with --spec-draft-device excluding the device that holds the target LM head, this inserts a preallocated target tensor into a graph whose scheduler was created only from the draft model's model.devices. No registered backend can support that buffer, so scheduler assignment aborts instead of decoding. Include the target tensor devices in the draft context's backend set, or reject incompatible device configurations before building the graph.
Useful? React with 👍 / 👎.
| if (cparams.ctx_type != LLAMA_CONTEXT_TYPE_MTP || !tgt || !tgt->tok_embd || !tgt->output || | ||
| tgt->hparams.n_embd != model.hparams.n_embd || tgt->vocab.n_tokens() != model.vocab.n_tokens()) { | ||
| throw std::runtime_error(model.arch_name() + " head without token_embd/output only works as an MTP draft (--spec-type draft-mtp) of a matching target (this warning is normal during memory fitting)"); |
There was a problem hiding this comment.
With a separate shared MTP head and the server's default --fit on, the extra-model measurement creates this context without ctx_other, so this condition rejects it. common_params_fit_impl catches that failure and substitutes zero memory for the entire extra model, omitting the head weights, KV cache, and compute buffers; on memory-constrained systems the fitter can therefore retain an oversized context or offload plan and the real load can OOM. The no-allocation fitting path needs a way to validate/build this context without requiring a live target context.
Useful? React with 👍 / 👎.
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. |
…ed heads (unslothai#242) Drop ggml-org#27754, unslothai#144 and unslothai#137: upstream merged GLM-5-Next (ggml-org#27773), Qwen4Exp MTP (ggml-org#29761) and row prefetch (ggml-org#29599), and the old pins no longer merge. Old-GGUF loading moves to unslothai#239, shared MTP heads to unslothai#240, and what the dropped pins had beyond upstream to unslothai#241. Repin ggml-org#24423, ggml-org#25731, unslothai#61 and unslothai#176 to their conflict-fixed heads.
This branch has not been deployed
1 participant
Summary
Upstream merged Qwen4Exp MTP (ggml-org#29761), superseding our #144, but the MTP heads published in
unsloth/Qwen3.8-Flash-Next-GGUF/MTP/still do not run on mainline:mtp-*-shared-*.gguf(recommended, notoken_embd/output)check_tensor_dims: tensor 'token_embd.weight' not foundmtp-*-{BF16,Q8_0,Q4_K_M}.gguf(self-contained)GGML_ASSERT(buffer) failedThe shared heads were introduced in #142 (
nextn_shared_target_tensors), which borrowed the target's tensors in the loader; that went away with #144.Changes
token_embd/outputand is created as an MTP draft (LLAMA_CONTEXT_TYPE_MTP),llama_contextrecords the target model fromparams.ctx_otherin a newcparams.model_tgt, andgraph_mtptakes the embedding and LM head (with its scale) from it.cparams.ctx_otheris deliberately left unset: for draft-mtp,common/speculative.cppreads it as "the draft shares the target's memory" (gemma4-assistant), and it is also passed asmem_otherto the memory.load_arch_tensorsmakestoken_embdoptional only for MTP-only files and no longer duplicates a missingtoken_embdasoutput. Any other use of such a head (e.g.--spec-type draft-simple) fails with a clear error instead of a null dereference; the memory-fit pass fails gracefully as before.compress_ratios[n_layer] = 0(the MTP block attends densely), while upstream's converter writes the QSA ratio there.graph_mtpbuilt the k-pool input unconditionally, but the draft's indexer cache then holds no layer, sok_idxswas never allocated andset_inputasserted. The k-pool input is now built only for a QSA MTP block;build_layer_attnalready treats ratio 0 as dense. Heads converted by upstream are unaffected.Testing
Qwen3.8-Flash-Next UD-IQ1_S +
-md <head> --spec-type draft-mtp --spec-draft-n-max 2, llama-server, greedy, 3 prompts x 3 repeats, one exclusive B200:-np 2.test-llama-archs(133/133) passes.Related: #227 (fit: measure a draft head that borrows the target's embeddings) is complementary; without it the fit pass logs one "failed to measure the memory of the extra model" line for shared heads, as before.
Base is b11368 (
base/upstream-1fb7ef3e3). The full nightly pin set merges cleanly with this PR on top.