qwen4exp: load the MTP head's hc_head_norm as [n_embd, hc] like the trunk gammas

[LeoBorcherding]

Fixes the GGML_ASSERT(ggml_can_repeat(b, a)) abort that kills the Qwen3.8-Flash-Next drafter on every mix build since b10995-mix-3e83366.

Reported in:

• [Bug] b10995-mix-3e83366: Qwen3.8-Flash-Next MTP aborts at load (nextn.hc_head_norm still [hc_dim] after rebase) unsloth#11143 (macOS/Metal first, then Windows and Linux CUDA; the diagnosis this change comes from)
• MTP drafter crashes with GGML_ASSERT(ggml_can_repeat(b, a)) during graph build (b11007-mix, RTX 5080) unsloth#11219 (Windows CUDA, both sidecar heads, merged single-file target)

What broke

f40f64a81 ("Fix merge conflicts") rebased this branch onto upstream's rms_norm + mul fusion for qwen4exp (ggml-org#28896). That reshapes the hyper-connection gammas: they now load as { n_embd, hc } with TENSOR_ALLOW_RESHAPE, and build_hc_mix multiplies the [n_embd, hc, n_tokens] stream by the gamma directly instead of flattening to hc_dim first.

Three trunk gammas got converted (qwen4exp.cpp:170, :220, :224), and so did the comment sitting above them:

// the gammas load as [n_embd, hc] so the grouped norm multiplies them without a graph reshape
// the MTP head has its own in nextn.hc_head_*

The MTP head's gamma, twenty lines further down at :280, did not:

layer.nextn.hc_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_HC_HEAD_NORM, "weight", il), { hc_dim }, flags);

graph_mtp hands that straight to build_hc_mix (:643), so the first ggml_mul there sees a [10240] weight against a [2560, 4, n_tokens] input. ggml_can_repeat says no and the process aborts, in graph_reserve before a single draft token is decoded.

Nothing about the GGUF changed, which is why the same files work on b10909-mix-bea84f7. The weight was always flat; only the expected shape moved.

The fix

layer.nextn.hc_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_HC_HEAD_NORM, "weight", il), { n_embd, hc }, flags | TENSOR_ALLOW_RESHAPE);

Same treatment the trunk gammas got, and TENSOR_ALLOW_RESHAPE is what lets the flat [10240] tensor in the file satisfy the 2D declaration.

nextn.hnorm at :277 is still { hc_dim } and stays that way. It looks like the same oversight but isn't: graph_mtp reshapes to 2D before multiplying by it (:540-542), so that multiply is well formed either way.

Testing

I have not built this. The change is a code read against f40f64a81, the GGUF header and the backtrace; the run below is someone else's.

@SvenMeyer applied the same change to the shipped ~/.unsloth/llama.cpp tree, rebuilt, and got a working drafter (unslothai/unsloth#11143):

RTX 4070 Laptop 8 GB, Linux CUDA, UD-IQ4_XS + mtp-…-shared-Q4_K_M, 8192 ctx, --cpu-moe --spec-draft-cpu-moe, q8_0 KV, --spec-type draft-mtp --spec-draft-n-max 3

workload	no MTP	MTP n_max=3	accepted
digit counting, 150 tok	8.06 t/s	12.93 t/s	98% (112/114)
short prose, 180 tok	8.06 t/s	6.51 t/s	31% (85/277)

The crash itself is reported on macOS/Metal (M4 Max), Windows CUDA (2×5060 Ti + 3060, and a 5080), and Linux CUDA, with both the shared and the standalone head, flash attention on and off, sharded and gguf-split --merged targets, mmap and --load-mode none. It is not hardware or file specific: any path that reaches graph_mtp hits it.

Also worth knowing

Studio catches the exit, retries once without flash attention, then falls back to loading with no speculative decoding at all. The load API still returns 200, so the only user-facing symptom is that decode got slower. That's tracked separately.

Once this lands, #144's pin in scripts/unsloth/pr-set.json needs a bump to the new head, same shape as #217.

Diagnosis by @zachweyland, fix confirmed by @SvenMeyer.

[zorrobyte]

Independent verification, since this hasn't been built yet: reproduced the crash and confirmed the fix on a from-source CUDA build.

Setup: RTX 5090 (32GB), unslothai/llama.cpp mix b11027-mix-3e83366 (commit f6b9ea7), built locally with -DGGML_CUDA=ON -DCMAKE_CUDA_ARCHITECTURES=120. Main model Qwen3.8-Flash-Next-UD-Q2_K_XL, draft head mtp-Qwen3.8-Flash-Next-shared-Q8_0.gguf, --spec-type draft-mtp --spec-draft-n-max 2.

Before fix: reproduces exactly as described — GGML_ASSERT(ggml_can_repeat(b, a)) in build_hc_mix, called from graph_mtp, abort before any token is decoded.

Fix applied: functionally equivalent patch (reshape w_norm to [n_embd, hc] at the build_hc_mix call site rather than at tensor declaration — same effect on the hc_head_norm shape mismatch). With it, the server loads and serves normally:

draft acceptance = 0.66082 (113 accepted / 171 generated), mean len = 2.31

66.1% acceptance matches Unsloth's own published number for shared-Q8_0 exactly, which is a good sign the fix is correct and not just papering over the crash.

For what it's worth, on this setup (Q2_K_XL, most MoE experts CPU-offloaded — i.e. bandwidth-bound, not the fully GPU-resident B200 case in the docs) the real-world speedup was much smaller than 1.3–1.7x: roughly a wash at 1 stream, and prefill measurably regressed (~-22%, 2440→1900 tok/s at 15k context). Matches the "skip MTP for concurrent serving" guidance in the MTP README, just adding a data point for the CPU-offload case specifically.

Happy to share full logs/config if useful.