nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

qwen4exp: MTP (NextN) self-speculative decoding support (#2369) 

* qwen4exp: MTP (NextN) self-speculative decoding support

Wires the checkpoint's 2.6B MTP head into the existing MTP speculative
framework for the qwen4exp architecture, following the deepseek4
conventions: block_count includes the NextN tail, the draft pass walks
[n_layer - nextn, n_layer) through the same layer-loop body, and the
target pass emits the pre-final-mixer wide stream as result_mtp_embd.

- hparams: read nextn_predict_layers; the tail is forced non-recurrent
  (the full-attention-interval pattern would classify it as GDN) and
  inherits the last main QSA layer's compress ratio
- loader: the tail layers and seven mtp.* entry/mixer tensors load only
  when an MTP stage is requested (TENSOR_SKIP otherwise), so nextn-
  bearing GGUFs serve normally on plain configs
- graph: MTP entry = per-reference residual_linear_shared fusion (joint
  RMS over the flattened hc*n_embd stream, shared fc_hidden per branch,
  fc_embedding added as a residual to every branch); exit reuses the
  hc_head-shaped mixer with dedicated weights; delta-net construction
  skips the draft pass (the draft context has no recurrent slots)
- llama: qwen4exp joins the MTP-capable arch allowlist and
  llama_mtp_tail_uses_layer_cache; MTP feature width is the wide stream
  (n_embd * hc_mult); the nextn pooling path gets its required
  result_embd tensor

Measured on a 125B-A6B GGUF carrying the head (RTX 5090 + 24-core Zen5
hybrid, -ncmoe 38): draft acceptance 0.93-0.99, per-position 95%/pos on
code at n_max=4; draft passes cost ~2ms wall (131-144 nodes).

* cuda: dispatch the ABS and SGN unary ops

The kernels exist in unary.cu but neither the execution switch nor
supports_op routed them, so the qwen4exp PLE gate split to the CPU
scheduler twice per forward pass once its layer is offloaded (the cost
is the round trip, not the arithmetic — noted in the port's PR text).

* refactor(model): add qwen4exp NextN mixer tensors

* fix(qwen4exp): resolve standard NextN head layouts

* qwen4exp: decline per-step spec checkpoints (PLE conv tail not covered)

The per-step checkpoint path sizes its buffers from the ssm_* hparams
alone, but qwen4exp appends the PLE n-gram convolution history to the
tail of the same recurrent state row. A rejected draft token would roll
back the delta-net state and leave the PLE tail advanced. Decline
per-step (like openPangu) so the mode resolves to the whole-slot
gpu-fallback shadow, which covers the full row.

* qwen4exp: keep legacy mtp.* files loadable without MTP

Pre-rework files carried the NextN head under fork-local mtp.* names.
Register those seven names as known-and-skipped so such files still load
for non-MTP use instead of failing tensor accounting; an MTP request
against them reports the standard missing-layout error.

* loader: summarize skipped tensors in one line (per-tensor detail moves to debug)

An unused NextN/MTP head produces 33 ignore warnings per load; collapse
to a single summary at done_getting_tensors.

* spec: fall back to generic copy when checkpoint rows resolve no backend

Partial offload leaves host-resident recurrent state rows (CUDA_Host) matching
no backend's default buffer type between graphs, so checkpoint save/restore
asserted on a null backend. Drain the scheduler once and use the generic
tensor copy for unresolvable rows; the async fast path is unchanged. Also
guards the split-save branch that previously passed a null backend unchecked.

Diagnosed, patched, and validated by @joelfarthing (round-trip exact on
-ngl 6/8, -ncmoe control unchanged).

Co-authored-by: joelfarthing <joelfarthing@users.noreply.github.com>

* qwen4exp: support separate MTP head via -md (companion pair)

Lets a plain community quant load a predictor-only GGUF as the draft, so MTP
works without re-downloading a merged file and the head can sit on GPU at
full precision while the target runs partially offloaded. Relaxes the five
sites that assumed the target carries the NextN block: predict-layer count,
optional trunk tensors for mtp-only files, hc_attn_norm trunk marker,
n_embd_out derivation, and the wide-embd export predicate, plus qwen4exp
companion pairing enforcement.

Authored by @joelfarthing, replicated by @treo (3090, full 262K context).
Known limitation: llama-quantize on a predictor-only file trips the
attention-tensor count assertion; quantize before extraction.

Co-authored-by: joelfarthing <joelfarthing@users.noreply.github.com>

* server: name the stage chain in the single-slot refusal and say the fix

* qwen4exp: condense verbose NextN comments per CONTRIBUTING.md

* qwen4exp: accept the Dzannotti NextN head layout (predictor-only)

A predictor-only (-md) file may omit blk.N.nextn.shared_head_norm and carry the
head mixer norm only in output_hc_norm (the Dzannotti convention, with the +1
baked in as the direct-multiply grouped-RMS kernel expects). Fall back to
model.hc_head_norm as the head norm in that case, gated to mtp_only files where
output_hc_* is the head's own mixer and not the trunk's. Quimmedes-style raw
(non-baked) tensors remain intentionally unsupported.

Reported by @SamuelOliveirads.

* spec: merge the deepseek4/qwen4exp wide-stream feature-width branches

Both arches return n_embd_out under the same condition; one combined branch.
Suggested by @SamuelOliveirads.

---------

Co-authored-by: joelfarthing <joelfarthing@users.noreply.github.com>
This commit is contained in:
James 2026-09-02 12:49:42 -04:00 committed by GitHub
parent 68bf92bfc3
commit 563b798ac9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
14 changed files with 438 additions and 89 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
common/speculative.cpp
View File

@ -147,7 +147,7 @@ static bool common_speculative_are_compatible(
static bool common_speculative_target_has_appended_mtp_contract(const llama_model * model) { static bool common_speculative_target_has_appended_mtp_contract(const llama_model * model) {
return llama_model_is_step35(model) || llama_model_is_deepseek4(model) || return llama_model_is_step35(model) || llama_model_is_deepseek4(model) ||
llama_model_is_qwen35_family(model); llama_model_is_qwen35_family(model) || llama_model_is_qwen4exp(model);
} }
static bool common_speculative_has_recognized_mtp_companion( static bool common_speculative_has_recognized_mtp_companion(
@ -160,8 +160,13 @@ static bool common_speculative_has_recognized_mtp_companion(
return llama_model_arch_string(target) != nullptr && return llama_model_arch_string(target) != nullptr &&
std::strcmp(llama_model_arch_string(target), "gemma4") == 0; std::strcmp(llama_model_arch_string(target), "gemma4") == 0;
} }
return common_speculative_target_has_appended_mtp_contract(target) && if (llama_model_mtp_package(companion) != LLAMA_MTP_PACKAGE_COMPANION) {
llama_model_mtp_package(companion) == LLAMA_MTP_PACKAGE_COMPANION; return false;
}
if (llama_model_is_qwen4exp(target) || llama_model_is_qwen4exp(companion)) {
return llama_model_is_qwen4exp(target) && llama_model_is_qwen4exp(companion);
}
return common_speculative_target_has_appended_mtp_contract(target);
} }
// state of an implementation of speculative decoding // state of an implementation of speculative decoding
@ -2221,6 +2226,18 @@ bool common_speculative_finalize_startup(
__func__, n_heads); __func__, n_heads);
return false; return false;
} }
} else if (llama_model_is_qwen4exp(model)) {
if (!llama_model_is_qwen4exp(companion)) {
LOG_ERR("%s: Qwen4Exp MTP requires a Qwen4Exp companion\n", __func__);
return false;
}
const int32_t n_heads = llama_model_n_nextn_layer(companion);
if (n_heads != 1) {
LOG_ERR("%s: Qwen4Exp MTP companion requires exactly one predictor layer, got %d\n",
__func__, n_heads);
return false;
}
} }
if (common_speculative_has_recognized_mtp_companion(model, companion)) { if (common_speculative_has_recognized_mtp_companion(model, companion)) {

3
examples/server/server-context.cpp
View File

@ -218,8 +218,9 @@ bool server_context::load_model(const gpt_params& params_) {
common_speculative_prepare_startup(params_base, false); common_speculative_prepare_startup(params_base, false);
if (server_speculative_requires_single_slot(params_base.speculative) && params_base.n_parallel > 1) { if (server_speculative_requires_single_slot(params_base.speculative) && params_base.n_parallel > 1) {
LOG_ERROR("Speculative decoding is currently limited to a single server slot (-np 1).\n", { LOG_ERROR("Speculative decoding is currently limited to a single server slot: restart with -np 1, or drop the --spec-type stage(s) to keep parallel slots.\n", {
{"n_parallel", params_base.n_parallel}, {"n_parallel", params_base.n_parallel},
{"stage_chain", common_speculative_stage_chain_to_str(params_base.speculative)},
}); });
return false; return false;
} }

8
ggml/src/ggml-cuda.cu
View File

@ -3827,6 +3827,12 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
case GGML_UNARY_OP_GELU_ERF: case GGML_UNARY_OP_GELU_ERF:
ggml_cuda_op_gelu_erf(ctx, dst); ggml_cuda_op_gelu_erf(ctx, dst);
break; break;
case GGML_UNARY_OP_ABS:
ggml_cuda_op_abs(ctx, dst);
break;
case GGML_UNARY_OP_SGN:
ggml_cuda_op_sgn(ctx, dst);
break;
case GGML_UNARY_OP_SILU: case GGML_UNARY_OP_SILU:
ggml_cuda_op_silu(ctx, dst); ggml_cuda_op_silu(ctx, dst);
break; break;
@ -4802,6 +4808,8 @@ GGML_CALL static bool ggml_backend_cuda_supports_op(ggml_backend_t backend, cons
switch (ggml_get_unary_op(op)) { switch (ggml_get_unary_op(op)) {
case GGML_UNARY_OP_GELU: case GGML_UNARY_OP_GELU:
case GGML_UNARY_OP_GELU_ERF: case GGML_UNARY_OP_GELU_ERF:
case GGML_UNARY_OP_ABS:
case GGML_UNARY_OP_SGN:
case GGML_UNARY_OP_SILU: case GGML_UNARY_OP_SILU:
case GGML_UNARY_OP_SWIGLU: case GGML_UNARY_OP_SWIGLU:
case GGML_UNARY_OP_SWIGLU_OAI: case GGML_UNARY_OP_SWIGLU_OAI:

2
include/llama.h
View File

@ -716,6 +716,8 @@ extern "C" {
LLAMA_API bool llama_model_is_qwen35_family(const struct llama_model * model); LLAMA_API bool llama_model_is_qwen35_family(const struct llama_model * model);
LLAMA_API bool llama_model_is_qwen4exp(const struct llama_model * model);
LLAMA_API bool llama_is_gemma4_mtp_file(const char * path); LLAMA_API bool llama_is_gemma4_mtp_file(const char * path);
LLAMA_API bool llama_model_is_split_mode_graph(const struct llama_model * model); LLAMA_API bool llama_model_is_split_mode_graph(const struct llama_model * model);

95
src/graphs/build_qwen4exp.cpp
View File

@ -3,6 +3,8 @@
#include "../llama-context.h" #include "../llama-context.h"
#include "../llama-delta-net.h" #include "../llama-delta-net.h"
#include <optional>
// the [hc_dim] gamma is wider than the per-stream reduction, so ggml_fused_rms_norm cannot // the [hc_dim] gamma is wider than the per-stream reduction, so ggml_fused_rms_norm cannot
// express this and the two ops stay separate // express this and the two ops stay separate
static ggml_tensor * qwen4exp_grouped_rms( static ggml_tensor * qwen4exp_grouped_rms(
@ -428,26 +430,67 @@ ggml_cgraph * llm_build_context::build_qwen4exp() {
ggml_cgraph * gf = new_graph_custom(); ggml_cgraph * gf = new_graph_custom();
delta_net delta(lctx, batch); const bool is_mtp = lctx.cparams.mtp_op_type != MTP_OP_NONE;
// the MTP pass walks only the QSA tail: no recurrent state, so the draft
// context has zero qnext state slots and the delta-net ctor must not run
std::optional<delta_net> delta_opt;
if (!is_mtp) {
delta_opt.emplace(lctx, batch);
}
const int32_t n_embd_head = hparams.n_embd_head_v(0); const int32_t n_embd_head = hparams.n_embd_head_v(0);
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0));
const int32_t hc = hparams.dsv4_hc_mult; const int32_t hc = hparams.dsv4_hc_mult;
ggml_tensor * inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb); const int n_layer_begin = is_mtp ? n_layer - hparams.nextn_predict_layers : 0;
const int n_layer_end = is_mtp ? n_layer : n_layer - hparams.nextn_predict_layers;
ggml_tensor * inp_pos = build_inp_pos(); ggml_tensor * inp_pos = build_inp_pos();
ggml_tensor * inp_out_ids = n_tokens > 1 ? build_inp_out_ids() : nullptr; ggml_tensor * inp_out_ids = (is_mtp || n_tokens > 1) ? build_inp_out_ids() : nullptr;
ggml_tensor * KQ_mask = build_inp_KQ_mask(); ggml_tensor * KQ_mask = build_inp_KQ_mask();
float KQ_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
ggml_tensor * res_hc = nullptr;
if (is_mtp) {
// the fill sites assert on these buffers; the MTP graph consumes neither
lctx.inp_s_seq_qnext = nullptr;
lctx.inp_ple_rows = nullptr;
const int32_t hc_dim = hc * n_embd;
ggml_tensor * hidden = build_inp_mtp_states(hc_dim); // target's pre-mixer wide stream
ggml_tensor * tok = build_inp_embd_mtp(model.tok_embd); // [n_embd, n_tokens]
const auto & nextn = model.layers[n_layer_begin].nextn;
// normalize the embedding (n_embd) and the wide hidden stream (hc*n_embd)
ggml_tensor * e = ggml_rms_norm(ctx0, tok, hparams.f_norm_rms_eps);
e = ggml_mul(ctx0, e, nextn.enorm);
ggml_tensor * h = ggml_rms_norm(ctx0, hidden, hparams.f_norm_rms_eps);
h = ggml_mul(ctx0, h, nextn.hnorm);
// eh_proj is [fc_embedding | fc_hidden] along its input dimension.
// Concatenating e_norm and h_norm per stream therefore preserves the old
// fc_embedding@e_norm + fc_hidden@h_norm entry result in one matmul.
e = ggml_repeat_4d(ctx0,
ggml_reshape_3d(ctx0, e, n_embd, 1, n_tokens), n_embd, hc, n_tokens, 1);
h = ggml_reshape_3d(ctx0, h, n_embd, hc, n_tokens);
ggml_tensor * eh = ggml_concat(ctx0, e, h, 0);
cb(eh, "mtp_eh_concat", n_layer_begin);
res_hc = llm_build_lora_mm(lctx, ctx0, nextn.eh_proj, eh);
cb(res_hc, "mtp_eh_proj", n_layer_begin);
} else {
ggml_tensor * inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb);
lctx.inp_s_seq_qnext = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, 1, n_tokens); lctx.inp_s_seq_qnext = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, 1, n_tokens);
cb(lctx.inp_s_seq_qnext, "inp_s_seq_qnext", -1); cb(lctx.inp_s_seq_qnext, "inp_s_seq_qnext", -1);
ggml_set_input(lctx.inp_s_seq_qnext); ggml_set_input(lctx.inp_s_seq_qnext);
float KQ_scale = hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;
// the wide residual starts as hc identical copies of the embedding // the wide residual starts as hc identical copies of the embedding
ggml_tensor * res_hc = ggml_repeat_4d(ctx0, res_hc = ggml_repeat_4d(ctx0,
ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens), ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens),
n_embd, hc, n_tokens, 1); n_embd, hc, n_tokens, 1);
cb(res_hc, "hc_residual", -1); cb(res_hc, "hc_residual", -1);
@ -459,6 +502,7 @@ ggml_cgraph * llm_build_context::build_qwen4exp() {
} else { } else {
lctx.inp_ple_rows = nullptr; lctx.inp_ple_rows = nullptr;
} }
}
// the same test the delta-net path uses for its own state // the same test the delta-net path uses for its own state
const bool ple_reset_state = batch.pos != nullptr && batch.pos[0] == 0; const bool ple_reset_state = batch.pos != nullptr && batch.pos[0] == 0;
@ -467,11 +511,11 @@ ggml_cgraph * llm_build_context::build_qwen4exp() {
ple_reset_pos[i] = batch.pos[i] == 0; ple_reset_pos[i] = batch.pos[i] == 0;
} }
for (int il = 0; il < n_layer; ++il) { for (int il = n_layer_begin; il < n_layer_end; ++il) {
ggml_tensor * inject = nullptr; ggml_tensor * inject = nullptr;
if (hparams.is_ple(il)) { if (hparams.is_ple(il)) {
res_hc = qwen4exp_ple(*this, ctx0, gf, lctx, model, hparams, delta, res_hc, lctx.inp_ple_rows, res_hc = qwen4exp_ple(*this, ctx0, gf, lctx, model, hparams, *delta_opt, res_hc, lctx.inp_ple_rows,
kv_self.s_l[il], ple_reset_state, ple_reset_pos, n_embd, n_tokens, il, cb); kv_self.s_l[il], ple_reset_state, ple_reset_pos, n_embd, n_tokens, il, cb);
} }
@ -481,7 +525,7 @@ ggml_cgraph * llm_build_context::build_qwen4exp() {
&inject, n_embd, il, cb); &inject, n_embd, il, cb);
if (hparams.is_recurrent(il)) { if (hparams.is_recurrent(il)) {
cur = delta.build_layer_attn_linear(ctx0, gf, cur, nullptr, il, cb, /* external_residual */ true, cur = delta_opt->build_layer_attn_linear(ctx0, gf, cur, nullptr, il, cb, /* external_residual */ true,
GGML_UNARY_OP_SIGMOID); GGML_UNARY_OP_SIGMOID);
} else { } else {
// the indexer reads the same block input as q/k/v, and returns the causal mask // the indexer reads the same block input as q/k/v, and returns the causal mask
@ -522,6 +566,36 @@ ggml_cgraph * llm_build_context::build_qwen4exp() {
cb(res_hc, "l_out", il); cb(res_hc, "l_out", il);
} }
if (is_mtp) {
// wide stream out: the next draft step's hidden input (scheme A)
ggml_tensor * flat = ggml_reshape_2d(ctx0, res_hc, hc * n_embd, n_tokens);
ggml_tensor * h_next = inp_out_ids ? ggml_get_rows(ctx0, flat, inp_out_ids) : flat;
cb(h_next, "result_mtp_embd", -1);
ggml_set_output(h_next);
ggml_build_forward_expand(gf, h_next);
// exit mixer is hc_head-shaped: per-stream norm + low-rank gate, no inject
const auto & nextn = model.layers[n_layer_begin].nextn;
ggml_tensor * cur = qwen4exp_hc_mix(*this, ctx0, lctx, hparams,
ggml_reshape_3d(ctx0, h_next, n_embd, hc, h_next->ne[1]),
nextn.hc_head_norm, nextn.hc_head_down, nextn.hc_head_up,
nullptr, nullptr, n_embd, -1, cb);
cur = llm_build_lora_mm(lctx, ctx0, model.output, cur);
cb(cur, "result_output", -1);
ggml_build_forward_expand(gf, cur);
return gf;
}
if (lctx.cparams.mtp && hparams.n_embd_out > hparams.n_embd) {
// hand the draft the pre-mixer wide stream for its first step
ggml_tensor * flat = ggml_reshape_2d(ctx0, res_hc, hc * n_embd, n_tokens);
ggml_tensor * h_nextn = inp_out_ids ? ggml_get_rows(ctx0, flat, inp_out_ids) : flat;
cb(h_nextn, "result_mtp_embd", -1);
ggml_set_output(h_nextn);
ggml_build_forward_expand(gf, h_nextn);
}
ggml_tensor * cur = qwen4exp_hc_mix(*this, ctx0, lctx, hparams, res_hc, ggml_tensor * cur = qwen4exp_hc_mix(*this, ctx0, lctx, hparams, res_hc,
model.hc_head_norm, model.hc_head_down, model.hc_head_up, model.hc_head_norm, model.hc_head_down, model.hc_head_up,
nullptr, nullptr, n_embd, -1, cb); nullptr, nullptr, n_embd, -1, cb);
@ -530,6 +604,9 @@ ggml_cgraph * llm_build_context::build_qwen4exp() {
cur = ggml_get_rows(ctx0, cur, inp_out_ids); cur = ggml_get_rows(ctx0, cur, inp_out_ids);
} }
// name the mixed stream for append_pooling (the nextn tail needs a named embd)
cb(cur, "result_embd", -1);
cur = llm_build_lora_mm(lctx, ctx0, model.output, cur); cur = llm_build_lora_mm(lctx, ctx0, model.output, cur);
cb(cur, "result_output", -1); cb(cur, "result_output", -1);

3
src/llama-arch.h
View File

@ -409,6 +409,9 @@ enum llm_tensor {
LLM_TENSOR_NEXTN_HNORM, LLM_TENSOR_NEXTN_HNORM,
LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD,
LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, // 95 LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, // 95
LLM_TENSOR_NEXTN_HC_HEAD_NORM,
LLM_TENSOR_NEXTN_HC_HEAD_DOWN,
LLM_TENSOR_NEXTN_HC_HEAD_UP,
LLM_TENSOR_INDEXER_K_NORM, LLM_TENSOR_INDEXER_K_NORM,
LLM_TENSOR_INDEXER_PROJ, LLM_TENSOR_INDEXER_PROJ,
LLM_TENSOR_INDEXER_ATTN_K, LLM_TENSOR_INDEXER_ATTN_K,

31
src/llama-hparams.cpp
View File

@ -629,20 +629,45 @@ void llm_load_hparams(
ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult);
ml.get_key(LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank); ml.get_key(LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank);
// the MTP handover is the wide pre-mixer residual, so the width must not depend on an appended NextN block
if (hparams.dsv4_hc_mult > 1) {
const uint32_t wide = hparams.n_embd * hparams.dsv4_hc_mult;
if (hparams.n_embd_out != hparams.n_embd && hparams.n_embd_out != wide) {
throw std::runtime_error("qwen4exp: embedding_length_out must equal n_embd * hyper-connection count");
}
hparams.n_embd_out = wide;
}
ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false);
if (hparams.nextn_predict_layers >= hparams.n_layer) {
throw std::runtime_error("qwen4exp: nextn_predict_layers must be smaller than block_count");
}
{ {
uint32_t full_attn_interval = 4; uint32_t full_attn_interval = 4;
ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false); ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false);
for (uint32_t i = 0; i < hparams.n_layer; ++i) { for (uint32_t i = 0; i < hparams.n_layer; ++i) {
hparams.recurrent_layer_arr[i] = ((i + 1) % full_attn_interval != 0); hparams.recurrent_layer_arr[i] = ((i + 1) % full_attn_interval != 0);
} }
// the MTP tail is a full-attention (QSA) layer regardless of the interval pattern
for (uint32_t i = hparams.n_layer - hparams.nextn_predict_layers; i < hparams.n_layer; ++i) {
hparams.recurrent_layer_arr[i] = false;
}
} }
{ {
// block_count includes the nextn tail, but the converter's ratio array covers
// only the main layers; the tail inherits the last main QSA layer's ratio
const uint32_t n_main = hparams.n_layer - hparams.nextn_predict_layers;
uint32_t n_ratios = 0; uint32_t n_ratios = 0;
if (ml.get_arr_n(LLM_KV_ATTENTION_COMPRESS_RATIOS, n_ratios, false) && n_ratios >= hparams.n_layer) { if (ml.get_arr_n(LLM_KV_ATTENTION_COMPRESS_RATIOS, n_ratios, false) && n_ratios >= n_main) {
std::vector<uint32_t> ratios; std::vector<uint32_t> ratios;
ml.get_arr(ml.llm_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS), ratios); ml.get_arr(ml.llm_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS), ratios);
std::copy_n(ratios.begin(), hparams.n_layer, hparams.dsv4_compress_ratios.begin()); const uint32_t n_copy = std::min<uint32_t>(n_ratios, hparams.n_layer);
std::copy_n(ratios.begin(), n_copy, hparams.dsv4_compress_ratios.begin());
for (uint32_t i = n_copy; i < hparams.n_layer; ++i) {
hparams.dsv4_compress_ratios[i] = ratios[n_main - 1];
}
} }
} }
@ -692,7 +717,7 @@ void llm_load_hparams(
} }
} }
switch (hparams.n_layer) { switch (hparams.n_layer - hparams.nextn_predict_layers) {
case 48: model.type = e_model::MODEL_125B_A6B; break; case 48: model.type = e_model::MODEL_125B_A6B; break;
default: model.type = e_model::MODEL_UNKNOWN; default: model.type = e_model::MODEL_UNKNOWN;
} }

219
src/llama-load-tensors.cpp
View File

@ -1681,19 +1681,41 @@ bool create_tensors_helper::create_qwen4exp_tensors(const LLM_TN & tn) {
const int32_t hc_dim = hc * n_embd; const int32_t hc_dim = hc * n_embd;
const int32_t hc_rank = hparams.hc_low_rank; const int32_t hc_rank = hparams.hc_low_rank;
// a companion target passes the head via -md and legitimately declares zero NextN blocks
if (model.mtp && hparams.nextn_predict_layers > 1) {
throw std::runtime_error(format(
"qwen4exp: MTP supports at most one appended NextN block; nextn_predict_layers is %u",
hparams.nextn_predict_layers));
}
model.tok_embd = create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}); model.tok_embd = create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab});
// a predictor-only GGUF reports the full block count but ships only the NextN block
const bool mtp_only = hparams.nextn_predict_layers > 0 &&
ml.get_tensor_meta(tn(LLM_TENSOR_HC_ATTN_NORM, "weight", 0).c_str()) == nullptr;
const int trunk_flags = mtp_only
? llama_model_loader::TENSOR_SKIP | llama_model_loader::TENSOR_NOT_REQUIRED : 0;
// The wide residual is normalised and collapsed by a hyper-connection mix rather // The wide residual is normalised and collapsed by a hyper-connection mix rather
// than by an output_norm, so this architecture carries none. // than by an output_norm, so this architecture carries none.
model.hc_head_norm = create_tensor(ctx_output, tn(LLM_TENSOR_HC_HEAD_NORM, "weight"), {hc_dim}); const int out_mixer_flags = mtp_only ? llama_model_loader::TENSOR_NOT_REQUIRED : 0;
model.hc_head_down = create_tensor(ctx_output, tn(LLM_TENSOR_HC_HEAD_DOWN, "weight"), {hc_dim, hc_rank}); model.hc_head_norm = create_tensor(ctx_output, tn(LLM_TENSOR_HC_HEAD_NORM, "weight"), {hc_dim}, out_mixer_flags);
model.hc_head_up = create_tensor(ctx_output, tn(LLM_TENSOR_HC_HEAD_UP, "weight"), {hc_rank, hc_dim}); model.hc_head_down = create_tensor(ctx_output, tn(LLM_TENSOR_HC_HEAD_DOWN, "weight"), {hc_dim, hc_rank}, out_mixer_flags);
model.hc_head_up = create_tensor(ctx_output, tn(LLM_TENSOR_HC_HEAD_UP, "weight"), {hc_rank, hc_dim}, out_mixer_flags);
model.output = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_NOT_REQUIRED); model.output = create_tensor(ctx_output, tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_NOT_REQUIRED);
if (model.output == NULL) { if (model.output == NULL) {
model.output = create_tensor(ctx_output, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_DUPLICATED); model.output = create_tensor(ctx_output, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, llama_model_loader::TENSOR_DUPLICATED);
} }
if (hparams.ple_n_heads > 0) { bool ple_table_needed = hparams.ple_n_heads > 0;
if (ple_table_needed && mtp_only) {
// only the predictor tail loads from this file; the table is needed only if a tail block uses PLE
ple_table_needed = false;
for (uint32_t il = hparams.n_layer - hparams.nextn_predict_layers; il < hparams.n_layer; ++il) {
ple_table_needed = ple_table_needed || hparams.is_ple(il);
}
}
if (ple_table_needed) {
// The row count comes from the table itself. The converter shards the n-gram // The row count comes from the table itself. The converter shards the n-gram
// table, so the sum of ple_head_vocab_sizes is not the stored row count. // table, so the sum of ple_head_vocab_sizes is not the stored row count.
const std::string ple_name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"); const std::string ple_name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight");
@ -1738,60 +1760,171 @@ bool create_tensors_helper::create_qwen4exp_tensors(const LLM_TN & tn) {
auto & layer = model.layers[i]; auto & layer = model.layers[i];
layer.hc_attn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_NORM, "weight", i), {hc_dim}); // the nextn (MTP) tail loads only when speculative MTP is requested
layer.hc_attn_down = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_DOWN, "weight", i), {hc_dim, hc_rank}); const bool is_nextn = hparams.nextn_predict_layers > 0 &&
layer.hc_attn_up = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_UP, "weight", i), {hc_rank, hc_dim}); (uint32_t) i >= hparams.n_layer - hparams.nextn_predict_layers;
layer.hc_attn_inject = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_INJECT, "weight", i), {hc_dim, hc}); int lf = is_nextn && !model.mtp ? llama_model_loader::TENSOR_SKIP : 0;
layer.hc_ffn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_NORM, "weight", i), {hc_dim}); if (mtp_only && !is_nextn) {
layer.hc_ffn_down = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_DOWN, "weight", i), {hc_dim, hc_rank}); lf |= trunk_flags;
layer.hc_ffn_up = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_UP, "weight", i), {hc_rank, hc_dim}); }
layer.hc_ffn_inject = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_INJECT, "weight", i), {hc_dim, hc});
layer.hc_attn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_NORM, "weight", i), {hc_dim}, lf);
layer.hc_attn_down = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_DOWN, "weight", i), {hc_dim, hc_rank}, lf);
layer.hc_attn_up = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_UP, "weight", i), {hc_rank, hc_dim}, lf);
layer.hc_attn_inject = create_tensor(ctx_split, tn(LLM_TENSOR_HC_ATTN_INJECT, "weight", i), {hc_dim, hc}, lf);
layer.hc_ffn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_NORM, "weight", i), {hc_dim}, lf);
layer.hc_ffn_down = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_DOWN, "weight", i), {hc_dim, hc_rank}, lf);
layer.hc_ffn_up = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_UP, "weight", i), {hc_rank, hc_dim}, lf);
layer.hc_ffn_inject = create_tensor(ctx_split, tn(LLM_TENSOR_HC_FFN_INJECT, "weight", i), {hc_dim, hc}, lf);
if (hparams.is_ple(i)) { if (hparams.is_ple(i)) {
layer.ple_key = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_KEY, "weight", i), {n_embd, hc_dim}); layer.ple_key = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_KEY, "weight", i), {n_embd, hc_dim}, lf);
layer.ple_value = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_VALUE, "weight", i), {n_embd, n_embd}); layer.ple_value = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_VALUE, "weight", i), {n_embd, n_embd}, lf);
layer.ple_norm_key = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_NORM_KEY, "weight", i), {hc_dim}); layer.ple_norm_key = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_NORM_KEY, "weight", i), {hc_dim}, lf);
layer.ple_norm_query = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_NORM_QUERY, "weight", i), {hc_dim}); layer.ple_norm_query = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_NORM_QUERY, "weight", i), {hc_dim}, lf);
layer.ple_norm_conv = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_NORM_CONV, "weight", i), {hc_dim}); layer.ple_norm_conv = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_NORM_CONV, "weight", i), {hc_dim}, lf);
layer.ple_conv1d = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_CONV1D, "weight", i), {hparams.ple_conv_kernel, hc_dim}); layer.ple_conv1d = create_tensor(ctx_split, tn(LLM_TENSOR_PLE_CONV1D, "weight", i), {hparams.ple_conv_kernel, hc_dim}, lf);
} }
if (!hparams.is_recurrent(i)) { if (!hparams.is_recurrent(i)) {
// wq carries the query and an equal-width gate // wq carries the query and an equal-width gate
layer.wq = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head_k * n_head * 2}); layer.wq = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_embd_head_k * n_head * 2}, lf);
layer.wk = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k_gqa}); layer.wk = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, n_embd_k_gqa}, lf);
layer.wv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v_gqa}); layer.wv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, n_embd_v_gqa}, lf);
layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}); layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, lf);
layer.attn_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}); layer.attn_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, lf);
layer.attn_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}); layer.attn_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, lf);
// one indexer key is shared across the indexer heads // one indexer key is shared across the indexer heads
layer.indexer_q_proj = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", i), {n_embd, idx_head * idx_n_head}); layer.indexer_q_proj = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", i), {n_embd, idx_head * idx_n_head}, lf);
layer.indexer_k_proj = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", i), {n_embd, idx_head}); layer.indexer_k_proj = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", i), {n_embd, idx_head}, lf);
layer.indexer_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", i), {idx_head}); layer.indexer_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", i), {idx_head}, lf);
layer.indexer_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {idx_head}); layer.indexer_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {idx_head}, lf);
} else { } else {
layer.wqkv = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, key_dim * 2 + value_dim}); layer.wqkv = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, key_dim * 2 + value_dim}, lf);
layer.wqkv_gate = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, value_dim}); layer.wqkv_gate = create_tensor(ctx_layer, tn(LLM_TENSOR_ATTN_GATE, "weight", i), {n_embd, value_dim}, lf);
layer.ssm_conv1d = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {hparams.ssm_d_conv, conv_dim}); layer.ssm_conv1d = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {hparams.ssm_d_conv, conv_dim}, lf);
layer.ssm_dt = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_DT, "bias", i), {hparams.ssm_dt_rank}); layer.ssm_dt = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_DT, "bias", i), {hparams.ssm_dt_rank}, lf);
layer.ssm_a = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_A_NOSCAN, i), {hparams.ssm_dt_rank}); layer.ssm_a = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_A_NOSCAN, i), {hparams.ssm_dt_rank}, lf);
layer.ssm_beta = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, num_v_heads}); layer.ssm_beta = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, num_v_heads}, lf);
layer.ssm_alpha = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_ALPHA, "weight", i), {n_embd, num_v_heads}); layer.ssm_alpha = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_ALPHA, "weight", i), {n_embd, num_v_heads}, lf);
layer.ssm_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_v_dim}); layer.ssm_norm = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_v_dim}, lf);
layer.ssm_out = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_OUT, "weight", i), {value_dim, n_embd}); layer.ssm_out = create_tensor(ctx_layer, tn(LLM_TENSOR_SSM_OUT, "weight", i), {value_dim, n_embd}, lf);
} }
auto ffn_ctx = ctx_split; auto ffn_ctx = ctx_split;
layer.ffn_gate_inp = create_tensor(ffn_ctx, tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}); layer.ffn_gate_inp = create_tensor(ffn_ctx, tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, lf);
use_mmap_buffer &= !create_std_ffn_exps(n_embd, tn, i, 0, n_ff_exp); use_mmap_buffer &= !create_std_ffn_exps(n_embd, tn, i, lf, n_ff_exp);
layer.ffn_gate_inp_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", i), {n_embd}); layer.ffn_gate_inp_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", i), {n_embd}, lf);
layer.ffn_gate_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}); layer.ffn_gate_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}, lf);
layer.ffn_up_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}); layer.ffn_up_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, lf);
layer.ffn_down_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}); layer.ffn_down_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, lf);
}
if (hparams.nextn_predict_layers > 0) {
const int bid = n_layer - 1;
auto & nextn = model.layers[bid].nextn;
auto * head_ctx = ctx_for_layer_split(bid);
// Head-level tensors are optional at file-load time. A normal load skips any
// that are present, while an MTP load resolves and validates one complete layout below.
int flags = llama_model_loader::TENSOR_NOT_REQUIRED;
if (!model.mtp) {
flags |= llama_model_loader::TENSOR_SKIP;
}
const std::string eh_proj_name = tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", bid);
const std::string enorm_name = tn(LLM_TENSOR_NEXTN_ENORM, "weight", bid);
const std::string hnorm_name = tn(LLM_TENSOR_NEXTN_HNORM, "weight", bid);
const std::string shared_norm_name = tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", bid);
const std::string hc_norm_name = tn(LLM_TENSOR_NEXTN_HC_HEAD_NORM, "weight", bid);
const std::string hc_down_name = tn(LLM_TENSOR_NEXTN_HC_HEAD_DOWN, "weight", bid);
const std::string hc_up_name = tn(LLM_TENSOR_NEXTN_HC_HEAD_UP, "weight", bid);
nextn.eh_proj = create_tensor(head_ctx, eh_proj_name, {2 * n_embd, n_embd}, flags);
nextn.enorm = create_tensor(head_ctx, enorm_name, {n_embd}, flags);
nextn.hnorm = create_tensor(head_ctx, hnorm_name, {hc_dim}, flags);
// Flavor A: the upstream converter's head-local mixer.
nextn.hc_head_norm = create_tensor(head_ctx, hc_norm_name, {hc_dim}, flags);
nextn.hc_head_down = create_tensor(head_ctx, hc_down_name, {hc_dim, hc_rank}, flags);
nextn.hc_head_up = create_tensor(head_ctx, hc_up_name, {hc_rank, hc_dim}, flags);
// Flavor B: community merge-script files retain only the head norm and
// use the trunk's already-loaded output_hc_down/output_hc_up pair.
nextn.shared_head_norm = create_tensor(head_ctx, shared_norm_name, {hc_dim}, flags);
// Legacy files carry the head under fork-local mtp.* names: register them
// known-and-skipped so a non-MTP load balances tensor accounting.
{
const int legacy = llama_model_loader::TENSOR_NOT_REQUIRED | llama_model_loader::TENSOR_SKIP;
create_tensor(head_ctx, "mtp.fc_embd.weight", {n_embd, n_embd}, legacy);
create_tensor(head_ctx, "mtp.fc_hidden.weight", {n_embd, n_embd}, legacy);
create_tensor(head_ctx, "mtp.pre_norm_embd.weight", {n_embd}, legacy);
create_tensor(head_ctx, "mtp.pre_norm_hidden.weight", {hc_dim}, legacy);
create_tensor(head_ctx, "mtp.mixer_norm.weight", {hc_dim}, legacy);
create_tensor(head_ctx, "mtp.mixer_down.weight", {hc_dim, hc_rank}, legacy);
create_tensor(head_ctx, "mtp.mixer_up.weight", {hc_rank, hc_dim}, legacy);
}
if (model.mtp) {
std::string missing;
auto require = [&](const ggml_tensor * tensor, const std::string & name) {
if (tensor == nullptr) {
if (!missing.empty()) {
missing += ", ";
}
missing += name;
}
};
require(nextn.eh_proj, eh_proj_name);
require(nextn.enorm, enorm_name);
require(nextn.hnorm, hnorm_name);
// Any dedicated-mixer tensor selects flavor A. Do not hide a malformed
// flavor A by mixing it with flavor B.
const bool has_dedicated_mixer = nextn.hc_head_norm != nullptr ||
nextn.hc_head_down != nullptr || nextn.hc_head_up != nullptr;
if (has_dedicated_mixer) {
require(nextn.hc_head_norm, hc_norm_name);
require(nextn.hc_head_down, hc_down_name);
require(nextn.hc_head_up, hc_up_name);
} else {
// Flavor B': a predictor-only file may omit shared_head_norm and
// carry the head norm only in output_hc_norm (Dzannotti layout).
// Safe only when mtp_only, where output_hc_* is the head's own mixer.
if (nextn.shared_head_norm == nullptr && mtp_only) {
require(model.hc_head_norm, "output_hc_norm.weight");
} else {
require(nextn.shared_head_norm, shared_norm_name);
}
require(model.hc_head_down, "output_hc_down.weight");
require(model.hc_head_up, "output_hc_up.weight");
}
if (!missing.empty()) {
throw std::runtime_error(format(
"qwen4exp: MTP requested but the NextN head is incomplete; missing tensor(s): %s. "
"Accepted layouts for block %d are flavor A: %s, %s, %s, %s, %s, %s; "
"or flavor B: %s, %s, %s, %s with trunk output_hc_down.weight and output_hc_up.weight",
missing.c_str(), bid,
eh_proj_name.c_str(), enorm_name.c_str(), hnorm_name.c_str(),
hc_norm_name.c_str(), hc_down_name.c_str(), hc_up_name.c_str(),
eh_proj_name.c_str(), enorm_name.c_str(), hnorm_name.c_str(), shared_norm_name.c_str()));
}
if (!has_dedicated_mixer) {
nextn.hc_head_norm = nextn.shared_head_norm != nullptr
? nextn.shared_head_norm : model.hc_head_norm;
nextn.hc_head_down = model.hc_head_down;
nextn.hc_head_up = model.hc_head_up;
LLAMA_LOG_WARN("qwen4exp: NextN head has no dedicated hc mixer; using the trunk output_hc mixer\n");
}
}
} }
return use_mmap_buffer; return use_mmap_buffer;

13
src/llama-model-loader.cpp
View File

@ -998,10 +998,15 @@ struct ggml_tensor * llama_model_loader::create_tensor(struct ggml_context * ctx
return NULL; return NULL;
} }
// skip unused tensors // skip unused tensors (per-tensor detail at debug level; one summary at
// done_getting_tensors so ordinary loads are not flooded by an unused
// NextN/MTP head)
if (flags & TENSOR_SKIP) { if (flags & TENSOR_SKIP) {
const size_t nbytes = ggml_nbytes(cur); const size_t nbytes = ggml_nbytes(cur);
LLAMA_LOG_WARN("model has unused tensor %s (size = %zu bytes) -- ignoring\n", name.c_str(), nbytes); LLAMA_LOG_DEBUG("model has unused tensor %s (size = %zu bytes) -- ignoring\n", name.c_str(), nbytes);
n_skipped++;
size_skipped += nbytes;
size_data -= nbytes; size_data -= nbytes;
n_created++; n_created++;
@ -1042,6 +1047,10 @@ struct ggml_tensor * llama_model_loader::create_tensor_as_view(struct ggml_conte
} }
void llama_model_loader::done_getting_tensors() const { void llama_model_loader::done_getting_tensors() const {
if (n_skipped > 0) {
LLAMA_LOG_WARN("%s: skipped %d unused tensors (%.1f MiB) -- e.g. a NextN/MTP head when MTP is not requested; per-tensor detail at debug log level\n",
__func__, n_skipped, size_skipped / 1024.0 / 1024.0);
}
if (n_created != n_tensors) { if (n_created != n_tensors) {
throw std::runtime_error(format("%s: wrong number of tensors; expected %d, got %d", __func__, n_tensors, n_created)); throw std::runtime_error(format("%s: wrong number of tensors; expected %d, got %d", __func__, n_tensors, n_created));
} }

2
src/llama-model-loader.h
View File

@ -38,6 +38,8 @@ struct llama_model_loader {
int n_kv = 0; int n_kv = 0;
int n_tensors = 0; int n_tensors = 0;
int n_created = 0; int n_created = 0;
int n_skipped = 0;
size_t size_skipped = 0;
int64_t n_elements = 0; int64_t n_elements = 0;
size_t n_bytes = 0; size_t n_bytes = 0;

24
src/llama-model.cpp
View File

@ -490,6 +490,13 @@ static const std::map<llm_arch, std::map<llm_tensor, std::string>> LLM_TENSOR_NA
{ LLM_TENSOR_HC_HEAD_NORM, "output_hc_norm" }, { LLM_TENSOR_HC_HEAD_NORM, "output_hc_norm" },
{ LLM_TENSOR_HC_HEAD_DOWN, "output_hc_down" }, { LLM_TENSOR_HC_HEAD_DOWN, "output_hc_down" },
{ LLM_TENSOR_HC_HEAD_UP, "output_hc_up" }, { LLM_TENSOR_HC_HEAD_UP, "output_hc_up" },
{ LLM_TENSOR_NEXTN_EH_PROJ, "blk.%d.nextn.eh_proj" },
{ LLM_TENSOR_NEXTN_ENORM, "blk.%d.nextn.enorm" },
{ LLM_TENSOR_NEXTN_HNORM, "blk.%d.nextn.hnorm" },
{ LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "blk.%d.nextn.shared_head_norm" },
{ LLM_TENSOR_NEXTN_HC_HEAD_NORM, "blk.%d.nextn.hc_head_norm" },
{ LLM_TENSOR_NEXTN_HC_HEAD_DOWN, "blk.%d.nextn.hc_head_down" },
{ LLM_TENSOR_NEXTN_HC_HEAD_UP, "blk.%d.nextn.hc_head_up" },
{ LLM_TENSOR_ATTN_Q, "blk.%d.attn_q" }, { LLM_TENSOR_ATTN_Q, "blk.%d.attn_q" },
{ LLM_TENSOR_ATTN_K, "blk.%d.attn_k" }, { LLM_TENSOR_ATTN_K, "blk.%d.attn_k" },
{ LLM_TENSOR_ATTN_V, "blk.%d.attn_v" }, { LLM_TENSOR_ATTN_V, "blk.%d.attn_v" },
@ -2401,6 +2408,15 @@ bool llama_model_is_step35(const llama_model * model) {
return model && model->arch == LLM_ARCH_STEP35; return model && model->arch == LLM_ARCH_STEP35;
} }
bool llama_model_is_qwen4exp(const llama_model * model) {
return model && model->arch == LLM_ARCH_QWEN4EXP;
}
// qwen4exp marks a present trunk block with hc_attn_norm, every other arch with attn_norm
static bool llama_model_trunk_block_present(const llama_layer & layer) {
return layer.attn_norm != nullptr || layer.hc_attn_norm != nullptr;
}
bool llama_model_is_qwen35_family(const llama_model * model) { bool llama_model_is_qwen35_family(const llama_model * model) {
return model && (model->arch == LLM_ARCH_QWEN35 || model->arch == LLM_ARCH_QWEN35MOE); return model && (model->arch == LLM_ARCH_QWEN35 || model->arch == LLM_ARCH_QWEN35MOE);
} }
@ -2417,7 +2433,7 @@ enum llama_mtp_package llama_model_mtp_package(const llama_model * model) {
const size_t n_nextn = model->hparams.nextn_predict_layers; const size_t n_nextn = model->hparams.nextn_predict_layers;
const bool has_common_package_contract = const bool has_common_package_contract =
llama_model_is_step35(model) || llama_model_is_deepseek4(model) || llama_model_is_step35(model) || llama_model_is_deepseek4(model) ||
llama_model_is_qwen35_family(model) || llama_model_is_qwen35_family(model) || llama_model_is_qwen4exp(model) ||
llama_model_is_gemma4_mtp_assistant(model); llama_model_is_gemma4_mtp_assistant(model);
if (!has_common_package_contract) { if (!has_common_package_contract) {
return n_nextn > 0 ? LLAMA_MTP_PACKAGE_EMBEDDED : LLAMA_MTP_PACKAGE_NONE; return n_nextn > 0 ? LLAMA_MTP_PACKAGE_EMBEDDED : LLAMA_MTP_PACKAGE_NONE;
@ -2425,9 +2441,9 @@ enum llama_mtp_package llama_model_mtp_package(const llama_model * model) {
if (n_nextn == 0) { if (n_nextn == 0) {
if (llama_model_is_step35(model) || llama_model_is_deepseek4(model) || if (llama_model_is_step35(model) || llama_model_is_deepseek4(model) ||
llama_model_is_qwen35_family(model)) { llama_model_is_qwen35_family(model) || llama_model_is_qwen4exp(model)) {
for (const auto & layer : model->layers) { for (const auto & layer : model->layers) {
if (layer.attn_norm != nullptr) { if (llama_model_trunk_block_present(layer)) {
return LLAMA_MTP_PACKAGE_TARGET_ONLY; return LLAMA_MTP_PACKAGE_TARGET_ONLY;
} }
} }
@ -2448,7 +2464,7 @@ enum llama_mtp_package llama_model_mtp_package(const llama_model * model) {
bool has_trunk = false; bool has_trunk = false;
for (size_t il = 0; il < first; ++il) { for (size_t il = 0; il < first; ++il) {
if (model->layers[il].attn_norm != nullptr) { if (llama_model_trunk_block_present(model->layers[il])) {
has_trunk = true; has_trunk = true;
break; break;
} }

6
src/llama-model.h
View File

@ -139,6 +139,12 @@ struct llama_layer_nextn {
struct ggml_tensor * hnorm = nullptr; struct ggml_tensor * hnorm = nullptr;
struct ggml_tensor * shared_head_head = nullptr; struct ggml_tensor * shared_head_head = nullptr;
struct ggml_tensor * shared_head_norm = nullptr; struct ggml_tensor * shared_head_norm = nullptr;
// qwen4exp: the NextN head's final hyper-connection mixer. For merged
// community files these are resolved to shared_head_norm + the trunk mixer.
struct ggml_tensor * hc_head_norm = nullptr;
struct ggml_tensor * hc_head_down = nullptr;
struct ggml_tensor * hc_head_up = nullptr;
}; };
// TODO: separate into "llama_layer_enc" and "llama_layer_dec" // TODO: separate into "llama_layer_enc" and "llama_layer_dec"

5
src/llama-spec-features.cpp
View File

@ -27,7 +27,10 @@ uint32_t llama_model_mtp_feature_width(const struct llama_model * model) {
hparams.mtp_backbone_n_embd > 0) { hparams.mtp_backbone_n_embd > 0) {
return hparams.mtp_backbone_n_embd; return hparams.mtp_backbone_n_embd;
} }
if (model->arch == LLM_ARCH_DEEPSEEK4 && hparams.n_embd_out > hparams.n_embd) { if ((model->arch == LLM_ARCH_DEEPSEEK4 ||
model->arch == LLM_ARCH_QWEN4EXP) &&
hparams.n_embd_out > hparams.n_embd) {
// the pre-final-mixer wide stream; n_embd_out holds even when the NextN block lives in a companion file
return hparams.n_embd_out; return hparams.n_embd_out;
} }
return hparams.n_embd; return hparams.n_embd;

53
src/llama.cpp
View File

@ -1156,6 +1156,7 @@ static bool llama_mtp_tail_uses_layer_cache(const llama_model & model) {
return model.hparams.nextn_predict_layers > 0 && return model.hparams.nextn_predict_layers > 0 &&
(model.arch == LLM_ARCH_GLM_DSA || (model.arch == LLM_ARCH_GLM_DSA ||
model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_QWEN35MOE ||
model.arch == LLM_ARCH_QWEN4EXP ||
model.arch == LLM_ARCH_STEP35); model.arch == LLM_ARCH_STEP35);
} }
@ -2029,6 +2030,15 @@ bool llama_kv_cache::checkpoint_save(ggml_backend_sched_t sched) {
ckpt.used_snapshot = used; ckpt.used_snapshot = used;
std::unordered_set<ggml_backend_t> backends_to_sync; std::unordered_set<ggml_backend_t> backends_to_sync;
// a row with no resolvable backend is copied generically, outside any stream;
// drain queued work once so the copy cannot race an in-flight graph
bool sched_synced = false;
auto sync_sched_once = [&]() {
if (!sched_synced) {
ggml_backend_sched_synchronize(sched);
sched_synced = true;
}
};
for (uint32_t il = 0; il < n_layer; ++il) { for (uint32_t il = 0; il < n_layer; ++il) {
if (s_l[il] == nullptr) { if (s_l[il] == nullptr) {
@ -2040,20 +2050,30 @@ bool llama_kv_cache::checkpoint_save(ggml_backend_sched_t sched) {
auto & shadow_split = ckpt.split_s_l_shadow[il]; auto & shadow_split = ckpt.split_s_l_shadow[il];
for (int d = 0; d < split_info->n_device; ++d) { for (int d = 0; d < split_info->n_device; ++d) {
if (split_info->splits[d] && shadow_split[d]) { if (split_info->splits[d] && shadow_split[d]) {
//ggml_backend_tensor_copy(split_info->splits[d], shadow_split[d]);
auto src_backend = ggml_backend_sched_get_tensor_backend(sched, split_info->splits[d]); auto src_backend = ggml_backend_sched_get_tensor_backend(sched, split_info->splits[d]);
if (src_backend == nullptr) {
// a host-resident row (e.g. CUDA_Host) matches no backend's default
// buffer type between graphs; the generic copy handles it
sync_sched_once();
ggml_backend_tensor_copy(split_info->splits[d], shadow_split[d]);
} else {
ggml_backend_tensor_copy_async(src_backend, src_backend, split_info->splits[d], shadow_split[d]); ggml_backend_tensor_copy_async(src_backend, src_backend, split_info->splits[d], shadow_split[d]);
backends_to_sync.insert(src_backend); backends_to_sync.insert(src_backend);
} }
} }
}
} else { } else {
GGML_ASSERT(ckpt.s_l_shadow[il] != nullptr); GGML_ASSERT(ckpt.s_l_shadow[il] != nullptr);
auto src_backend = ggml_backend_sched_get_tensor_backend(sched, s_l[il]); auto src_backend = ggml_backend_sched_get_tensor_backend(sched, s_l[il]);
GGML_ASSERT(src_backend != nullptr); if (src_backend == nullptr) {
sync_sched_once();
ggml_backend_tensor_copy(s_l[il], ckpt.s_l_shadow[il]);
} else {
ggml_backend_tensor_copy_async(src_backend, src_backend, s_l[il], ckpt.s_l_shadow[il]); ggml_backend_tensor_copy_async(src_backend, src_backend, s_l[il], ckpt.s_l_shadow[il]);
backends_to_sync.insert(src_backend); backends_to_sync.insert(src_backend);
} }
} }
}
for (auto backend : backends_to_sync) { for (auto backend : backends_to_sync) {
ggml_backend_synchronize(backend); ggml_backend_synchronize(backend);
@ -2078,6 +2098,15 @@ bool llama_kv_cache::checkpoint_restore(ggml_backend_sched_t sched) {
used = ckpt.used_snapshot; used = ckpt.used_snapshot;
std::unordered_set<ggml_backend_t> backends_to_sync; std::unordered_set<ggml_backend_t> backends_to_sync;
// a row with no resolvable backend is copied generically, outside any stream;
// drain queued work once so the copy cannot race an in-flight graph
bool sched_synced = false;
auto sync_sched_once = [&]() {
if (!sched_synced) {
ggml_backend_sched_synchronize(sched);
sched_synced = true;
}
};
for (uint32_t il = 0; il < n_layer; ++il) { for (uint32_t il = 0; il < n_layer; ++il) {
if (s_l[il] == nullptr) { if (s_l[il] == nullptr) {
@ -2090,19 +2119,28 @@ bool llama_kv_cache::checkpoint_restore(ggml_backend_sched_t sched) {
for (int d = 0; d < split_info->n_device; ++d) { for (int d = 0; d < split_info->n_device; ++d) {
if (split_info->splits[d] && shadow_split[d]) { if (split_info->splits[d] && shadow_split[d]) {
auto dst_backend = ggml_backend_sched_get_tensor_backend(sched, split_info->splits[d]); auto dst_backend = ggml_backend_sched_get_tensor_backend(sched, split_info->splits[d]);
if (dst_backend == nullptr) {
sync_sched_once();
ggml_backend_tensor_copy(shadow_split[d], split_info->splits[d]);
} else {
ggml_backend_tensor_copy_async(dst_backend, dst_backend, shadow_split[d], split_info->splits[d]); ggml_backend_tensor_copy_async(dst_backend, dst_backend, shadow_split[d], split_info->splits[d]);
backends_to_sync.insert(dst_backend); backends_to_sync.insert(dst_backend);
} }
} }
}
} else { } else {
GGML_ASSERT(ckpt.s_l_shadow[il] != nullptr); GGML_ASSERT(ckpt.s_l_shadow[il] != nullptr);
GGML_ASSERT(ggml_nbytes(ckpt.s_l_shadow[il]) == ggml_nbytes(s_l[il])); GGML_ASSERT(ggml_nbytes(ckpt.s_l_shadow[il]) == ggml_nbytes(s_l[il]));
auto dst_backend = ggml_backend_sched_get_tensor_backend(sched, s_l[il]); auto dst_backend = ggml_backend_sched_get_tensor_backend(sched, s_l[il]);
GGML_ASSERT(dst_backend != nullptr); if (dst_backend == nullptr) {
sync_sched_once();
ggml_backend_tensor_copy(ckpt.s_l_shadow[il], s_l[il]);
} else {
ggml_backend_tensor_copy_async(dst_backend, dst_backend, ckpt.s_l_shadow[il], s_l[il]); ggml_backend_tensor_copy_async(dst_backend, dst_backend, ckpt.s_l_shadow[il], s_l[il]);
backends_to_sync.insert(dst_backend); backends_to_sync.insert(dst_backend);
} }
} }
}
for (auto backend : backends_to_sync) { for (auto backend : backends_to_sync) {
ggml_backend_synchronize(backend); ggml_backend_synchronize(backend);
@ -8632,6 +8670,7 @@ struct llama_context * llama_init_from_model(
model->arch != LLM_ARCH_STEP35 && model->arch != LLM_ARCH_STEP35 &&
model->arch != LLM_ARCH_GEMMA4_ASSISTANT && model->arch != LLM_ARCH_GEMMA4_ASSISTANT &&
model->arch != LLM_ARCH_OPENPANGU && model->arch != LLM_ARCH_OPENPANGU &&
model->arch != LLM_ARCH_QWEN4EXP &&
cparams.mtp != 0) { cparams.mtp != 0) {
cparams.mtp = 0; cparams.mtp = 0;
} }
@ -9650,6 +9689,14 @@ static bool spec_ckpt_try_per_step(llama_kv_cache & kv, const llama_model & mode
return false; return false;
} }
// qwen4exp's recurrent row carries a PLE conv-history tail the per-step
// checkpoint doesn't size, so it would leave that tail advanced by rejected
// drafts. Decline; the whole-slot (gpu-fallback) shadow covers the full row.
if (model.arch == LLM_ARCH_QWEN4EXP) {
kv.save_per_step_ssm = false;
return false;
}
// Split recurrent tensors are supported as long as each layer exposes // Split recurrent tensors are supported as long as each layer exposes
// concrete backend buffers for the per-step tensors. CPU-only and mixed // concrete backend buffers for the per-step tensors. CPU-only and mixed
// CPU/GPU recurrent placement are also allowed. // CPU/GPU recurrent placement are also allowed.