qwen4exp: support draft-head-only GGUFs (unsloth layout) + fix draft-load regression

[TheArchitectit]

Overview

This patch lets --spec-type draft-mtp work with the draft-head-only GGUFs
that unsloth ships for Qwen3.8-Flash-Next, and fixes a one-line bug where the
draft loader used the target model's path instead of the -md path. It is a
companion to #27836 (qwen4exp architecture + NextN/MTP speculative decoding)
and depends on that PR landing.

Both problems showed up while trying to run speculative decoding on a
147 GB MoE model on a pure-CPU box: the shipped draft pack does not fit
what #27836's loader demands, and the draft-load bug made it impossible to
use a separate draft file at all.

Fix 1: draft-head-only GGUFs

#27836's qwen4exp loader requires hc_head_norm/hc_head_down/hc_head_up
and the PLE block unconditionally. The unsloth GGUFs ship draft-head-only
files that lack these tensors — the target carries output_hc_* in its
shards; the draft pack has only nextn.shared_head_norm + the block's own
ffn projections.

The patch adds an mtp_only probe (n_layer_nextn > 0 && blk.0.hc_attn_norm absent) and, when it fires:

• trunk hc_head_* tensors marked TENSOR_NOT_REQUIRED when mtp_only
• PLE block guarded behind !mtp_only
• trunk tensor loop skipped when mtp_only (layers 0..n_layer-1 absent)
• MTP graph falls back to shared_head_norm + hc_ffn_down/hc_ffn_up
  when the head's own mixer is absent (mirrors the fallback path in the
  reference implementation)
• nextn.hc_head_* marked TENSOR_NOT_REQUIRED; new nextn.shared_head_norm
  tensor loaded when present

Fix 2: draft-load regression

common/speculative.cpp:2585 loaded params.model.path (the target
model path) instead of model_path (the draft model path) — the variable
holding the user-specified -md path. This caused the 147 GB target model
to be loaded again as the draft, consuming all RAM and failing. One-line fix.

Additional information

Keeping this PR in draft until #28389 (my other open PR, cuda: fix CUB argsort
corruption) is reviewed, since new contributors are limited to one open PR.
It is independent of #28389 code-wise and waits on #27836 regardless.

Testing

Measured on a 48-core Haswell Xeon (E5-2660 v3) running Qwen3.8-Flash-Next
(UD-Q4_K_XL target, 4-shard GGUF, ~147 GB mlock), with the native MTP head
requantized to Q8_0 as the draft model.

# upstream binary + PR #27836 + this patch
--spec-type draft-mtp --spec-draft-n-max 4 -md mtp-qwen4exp-Q8_0draft.gguf

# 8-prompt diverse suite (ab_8prompt.py)
ON  (draft-mtp Q8_0): mean 4.28 t/s, acceptance 0.51–0.75, mean_len 3.0–4.0
OFF (stock twin):     mean 3.33 t/s
Δ: +29% mean, 2.0× peak (best prompt 6.62 t/s vs 3.33)

Comparing the per-stream combiner from #27836 against a mean-pooled variant
(same head, same flags): acceptance was 0.51–0.75 vs 0.18–0.20. That lines
up with the note in #27836 that the combiner has to run per hc stream.

Drafter quant ladder (same upstream binary, dn=4)

Each rung is a requantization of the native BF16 MTP head pack (bit-identical
lineage); I arm-tested every one with a journal-verified load.

Quant	Size	Suite t/s	Mean acc_len	vs stock
Q8_0	3.9 GB	4.28	2.25	+29%
Q4_K_M	2.8 GB	3.97	2.34	+19%
Q5_0	2.8 GB	3.95	2.40	+19%
Q3_K_M	2.2 GB	3.44	2.15	+3%
Q2_K	1.8 GB	3.44	2.07	+3%

Acceptance is only the limiting factor below ~5 bits; every rung beats stock
(3.33 t/s). Q8_0 is the practical pick (identical speed to F16 at half the
RAM; F16 ties at 4.29 t/s with highest acceptance 2.51 but costs 7.3 GB).

Draft-n-max sweep (Q8_0)

dn	Suite t/s	Mean acc_len	vs stock
2	3.72	1.88	+12%
4	4.28	2.25	+29%
6	3.52	2.47	+6%
8	3.38	2.60	+2%

dn=4 is the sweet spot — past that, the extra verify work costs more than
the acceptance it buys back.

Files changed

This branch carries three commits from #27836, so the GitHub diff against master shows 10 files. The changes authored here are two:

• src/models/qwen4exp.cpp — mtp_only detection, optional tensors, graph fallback (+25 lines)
• common/speculative.cpp — draft-load path fix (1 line)

The remaining files in the diff (conversion/qwen4exp.py, gguf-py/gguf/constants.py, gguf-py/gguf/tensor_mapping.py, src/llama-arch.{h,cpp}, src/llama-model.{h,cpp}, src/models/models.h) come from #27836 and are shown for merge context only.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES — an AI assistant (Claude Code) helped polish the wording of this description and worked alongside me on the validation runs; the design, measurements, and every number above are my own, and I take responsibility for all submitted changes.

[ggml-gh-bot]

Hi @TheArchitectit, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• PR Template not respected: Please respect the template when creating a new pull request. Make sure to fill out all required sections.

• AI-generated content: While code is allowed to be generated by AI, please write the PR description and commit messages on your own without the help of AI.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.

[TheArchitectit]

Status/plan note: per the automated review on #28389 (new contributors are limited to 1 open PR at a time), I am keeping this PR in draft until #28389 (cuda: fix CUB argsort corruption) is reviewed/approved. It is independent of #28389 code-wise; this one is a companion to #27836 and depends on it landing regardless. I will also update the description to follow the PR template. No action needed from maintainers on this PR in the meantime.

[pwilkin]

/bot review

[ggml-gh-bot]

Automated code review

Review of PR #28097 (commit 3fb9b98). Note: this branch carries three commits from companion PR 27836 (the qwen4exp MTP head itself); this review covers the whole diff but concentrates on the changes authored here: the mtp_only loader/graph fallback in src/models/qwen4exp.cpp and the one-line draft-path fix in common/speculative.cpp.

Verified correct

• (point 1) common/speculative.cpp:2585: confirmed against master that the draft model was being loaded from params.model.path (the target) instead of model_path (the -md path, assigned two lines above). The fix is correct and is actually a fix of a live master bug, independent of PR 27836. Worth noting it could land standalone if this PR waits on 27836.
• (point 2) The mtp_only probe (n_layer_nextn > 0 && blk.0.hc_attn_norm absent) matches the established pattern in deepseek2.cpp, deepseek32.cpp, bailingmoe3.cpp, cohere2moe.cpp. ml.get_weight() returns nullptr safely, both values are uint32_t so the hparams guard is underflow-safe, and layers is sized n_layer_all so the extended loop indexes in range.
• (point 3) The trunk graph's deferred gather for unmasked t_h_nextn mirrors deepseek32.cpp (gather only when masked or no export), and the driver sets masked=true on the draft context, so the unconditional gather in graph_mtp is consistent.
• (point 4) Commit messages use Assisted-by: for the AI contribution - correct per the contribution guidelines.

Will slow the review

• (point 5) gguf-py/gguf/constants.py (QWEN4EXP tensor list, ~line 2920): MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM is missing, yet this PR's loader now reads blk.{bid}.nextn.shared_head_norm and _QwenMtpMixin.filter_tensors remaps a checkpoint mtp.norm to model.layers.{N}.shared_head.norm. TensorNameMap.__init__ skips any tensor not in MODEL_TENSORS[arch], so converting a checkpoint that carries mtp.norm will die with Can not map tensor. The unsloth draft files evidently contain this tensor, so the checkpoint likely has it. Add MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM next to the other NEXTN_* entries.
• (point 6) src/models/qwen4exp.cpp:494-496: the head-mixer fallback decides each weight independently - head_norm on hc_head_norm presence, head_down/head_up on their own. A file with a partial hc_head_* set silently mixes the head's norm with the block's hc_ffn_down/hc_ffn_up matrices. Gate on the trio as a whole (hc_head_norm && hc_head_down && hc_head_up ? ... : fallback) so a partial set either falls back cleanly or errors out.
• (point 7) src/models/qwen4exp.cpp:478: GGML_ASSERT(n_layer_nextn == 1) aborts the whole process at graph-build time on a GGUF-declared value. Since the model loader already throws for malformed files, validate this in load_arch_hparams with a throw std::runtime_error so a nextn_predict_layers > 1 qwen4exp GGUF produces a clean load error instead of a crash. (The n_layer_nextn < n_layer_all assert matches qwen35.cpp's existing pattern, so that one is consistent as-is.)

Nits

• (point 8) graph_mtp ignores cparams.nextn_layer_offset and always uses hparams.n_layer(), while every sibling MTP graph (deepseek4, deepseek32, bailingmoe3) uses n_layer() + offset plus a range assert. Harmless while n_layer_nextn == 1, but the speculative driver does call llama_set_nextn_layer_offset() on the draft context; adding the same range assert would catch driver misuse early.
• (point 9) Loading a draft-head-only GGUF as a plain model (llama-cli -m draft.gguf without --spec-type draft-mtp) leaves the trunk hc_head_* null and aborts deep inside ggml with an opaque assert. Sibling archs behave the same, so this is not blocking; a clear "this is a draft-head-only file" error would be a nice follow-up.

On the PR shape

• (point 10) The speculative.cpp fix is a standalone master bugfix; the rest depends on 27836 landing. Consider splitting the draft-path fix into its own tiny PR so the RAM-eating regression is not held hostage by the larger review, as you did with 28389.

This review was generated automatically by pi coding agent using zai-org/GLM-5.3. It may contain mistakes. Maintainers make the final call.

[okigan]

Rebased this together with #27836 onto current master (465e49b9c as of 2026-09-07) and tested it end-to-end on a real 4×V100 box — wanted to share in case it's useful for landing both.

Rebase: 3 conflicts in src/models/qwen4exp.cpp, all from master's n_ff_exp having moved from a scalar to a per-layer array (n_ff_exp_arr / hparams.n_ff_exp(il)) since this branch diverged:

• Kept master's array-based accessor, added this branch's flags/trunk_flags (MTP-block-skip + optional-trunk-tensor) logic on top of it.
• Dropped this branch's LLM_KV_NEXTN_PREDICT_LAYERS read in load_arch_hparams — master's base load_hparams() already reads it generically before dispatching to the arch-specific loader, so it was a harmless but redundant second read.
• Kept master's tensor-shape-based PLE row computation, just added the !mtp_only guard from this PR.

Testing: unsloth/Qwen3.8-Flash-Next-GGUF (UD-IQ3_XXS target + the MTP/mtp-Qwen3.8-Flash-Next-Q8_0.gguf draft), --spec-type draft-mtp --spec-draft-n-max 2, full GPU offload on 4×V100. Loads cleanly (no more output_hc_norm.weight not found), and generation is genuinely speculating — sustained across three separate real requests (500/900/1600 max_tokens, one code-explanation prompt and two creative-writing prompts):

test	draft_n	accepted	acceptance	gen speed
explain + code	384	307	79.9%	51.3 tok/s
story (900 budget)	830	484	58.3%	42.3 tok/s
story (1600 budget)	1327	935	70.5%	47.0 tok/s

Output stayed coherent and correct throughout (checked all three in full — no corruption, no repetition loops).

Rebased branch for reference, if helpful: https://github.com/okigan/llama.cpp/tree/qwen4exp-mtp-fix (same commits/authors as #27836 and this PR, just replayed onto current master — happy to open a PR from it if that's useful, or drop it if you'd rather keep this moving under your own branches).

Thanks for the work on this — MTP for Qwen3.8-Flash-Next is a nice speedup once it's wired up.