From 9a26522af234f8db079ae3735f35ab6c20fe2c66 Mon Sep 17 00:00:00 2001 From: joelfarthing Date: Thu, 7 May 2026 07:46:41 -0500 Subject: [PATCH] qwen35moe : support MTP tail layer (#1745) Co-authored-by: Joel Farthing <262452229+joelfarthing@users.noreply.github.com> --- src/graphs/build_qwen35.cpp | 174 ++++++++++++++++++++++++++++-------- src/llama-build-context.cpp | 3 +- src/llama-build-context.h | 8 ++ src/llama-hparams.cpp | 18 +++- src/llama-load-tensors.cpp | 74 ++++++++++----- src/llama-model.cpp | 4 + src/llama.cpp | 12 ++- 7 files changed, 225 insertions(+), 68 deletions(-) diff --git a/src/graphs/build_qwen35.cpp b/src/graphs/build_qwen35.cpp index 1a978f4c..93c18793 100644 --- a/src/graphs/build_qwen35.cpp +++ b/src/graphs/build_qwen35.cpp @@ -7,56 +7,80 @@ ggml_cgraph * llm_build_context::build_qwen35moe() { struct ggml_cgraph * gf = ggml_new_graph_custom(ctx0, model.max_nodes(n_tokens), false); - delta_net delta(lctx, batch); - const int64_t n_embd_head = hparams.n_embd_head_v(0); GGML_ASSERT(n_embd_head == hparams.n_embd_head_k(0)); - ggml_tensor * inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb); ggml_tensor * inp_pos = build_inp_pos(); - ggml_tensor * inp_out_ids = n_tokens > 1 ? build_inp_out_ids() : nullptr; - ggml_tensor * KQ_mask = build_inp_KQ_mask(); - - 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); - 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; ggml_tensor * cur = nullptr; - for (int il = 0; il < n_layer; ++il) { - - if (hparams.is_recurrent(il)) { - cur = delta.build_layer_attn_linear(ctx0, gf, inpL, il == n_layer - 1 ? inp_out_ids : nullptr, il, cb); + if (cparams.mtp_op_type != MTP_OP_NONE) { + ggml_tensor * hidden_states_from_main_model; + if (cparams.mtp_op_type == MTP_OP_WARMUP || cparams.mtp_op_type == MTP_OP_UPDATE_ACCEPTED) { + hidden_states_from_main_model = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens); } else { - cur = build_std_attention(gf, model.layers[il].attn_norm, inpL, inp_pos, il == n_layer - 1 ? inp_out_ids : nullptr, nullptr, - KQ_mask, nullptr, nullptr, KQ_scale, 0.0f, 0, il, true, false, true, false, true); + hidden_states_from_main_model = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, hparams.n_embd); + } + ggml_set_name(hidden_states_from_main_model, "inp_mtp_states"); + ggml_set_input(hidden_states_from_main_model); + lctx.inp_mtp_states = hidden_states_from_main_model; + + const int il_mtp = hparams.n_layer - 1; + const auto & mtp_layer = model.layers[il_mtp]; + + cur = build_qwen35moe_mtp(mtp_layer, hidden_states_from_main_model, n_embd_head, gf, inp_pos); + } else { + delta_net delta(lctx, batch); + + ggml_tensor * inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb); + ggml_tensor * inp_out_ids = (n_tokens > 1 && !lctx.cparams.mtp) ? build_inp_out_ids() : nullptr; + ggml_tensor * KQ_mask = build_inp_KQ_mask(); + + 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); + 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; + + const int n_transformer_layers = n_layer - hparams.nextn_predict_layers; + for (int il = 0; il < n_transformer_layers; ++il) { + + if (hparams.is_recurrent(il)) { + cur = delta.build_layer_attn_linear(ctx0, gf, inpL, il == n_transformer_layers - 1 ? inp_out_ids : nullptr, il, cb); + } else { + cur = build_std_attention(gf, model.layers[il].attn_norm, inpL, inp_pos, il == n_transformer_layers - 1 ? inp_out_ids : nullptr, nullptr, + KQ_mask, nullptr, nullptr, KQ_scale, 0.0f, 0, il, true, false, true, false, true); + } + + cur = llm_build_std_moe_ffn(ctx0, lctx, model.layers[il].ffn_norm, cur, + model.layers[il].ffn_gate_inp, nullptr, + model.layers[il].ffn_up_exps, nullptr, + model.layers[il].ffn_gate_exps, nullptr, + model.layers[il].ffn_down_exps, nullptr, + nullptr, + model.layers[il].ffn_up_shexp, nullptr, // we don't have shared expert biases? + model.layers[il].ffn_gate_shexp, nullptr, + model.layers[il].ffn_down_shexp, nullptr, + n_expert, n_expert_used, + LLM_FFN_SILU, true, false, 0.0f, + LLM_EXPERT_GATING_FUNC_SOFTMAX, + LLM_FFN_SILU, cb, il, gf, true, model.layers[il].ffn_up_gate_exps, nullptr, model.layers[il].ffn_gate_inp_shexp); + + cur = lctx.cvec.apply_to(ctx0, cur, il); + cb(cur, "l_out", il); + + inpL = cur; } - cur = llm_build_std_moe_ffn(ctx0, lctx, model.layers[il].ffn_norm, cur, - model.layers[il].ffn_gate_inp, nullptr, - model.layers[il].ffn_up_exps, nullptr, - model.layers[il].ffn_gate_exps, nullptr, - model.layers[il].ffn_down_exps, nullptr, - nullptr, - model.layers[il].ffn_up_shexp, nullptr, // we don't have shared expert biases? - model.layers[il].ffn_gate_shexp, nullptr, - model.layers[il].ffn_down_shexp, nullptr, - n_expert, n_expert_used, - LLM_FFN_SILU, true, false, 0.0f, - LLM_EXPERT_GATING_FUNC_SOFTMAX, - LLM_FFN_SILU, cb, il, gf, true, model.layers[il].ffn_up_gate_exps, nullptr, model.layers[il].ffn_gate_inp_shexp); + if (lctx.cparams.mtp) { + cb(inpL, "result_mtp_embd", -1); + ggml_set_output(inpL); + } - cur = lctx.cvec.apply_to(ctx0, cur, il); - cb(cur, "l_out", il); - - inpL = cur; + cur = build_output(lctx, ctx0, inpL, model.output, model.output_norm, cb); + cb(cur, "result_output", -1); } - cur = build_output(lctx, ctx0, inpL, model.output, model.output_norm, cb); - cb(cur, "result_output", -1); - ggml_build_forward_expand(gf, cur); return gf; @@ -144,6 +168,82 @@ ggml_cgraph * llm_build_context::build_qwen35() { return gf; } +struct ggml_tensor * llm_build_context::build_qwen35moe_mtp( + const llama_layer & mtp_layer, + struct ggml_tensor * prev_embeddings, + int64_t n_embd_head, + struct ggml_cgraph * gf, + struct ggml_tensor * inp_pos) { + + const int il = hparams.n_layer - 1; + + struct ggml_tensor * KQ_mask = build_inp_KQ_mask(); + struct ggml_tensor * inp_out_ids = (n_tokens > 1 && n_outputs < n_tokens) ? build_inp_out_ids() : nullptr; + + ggml_tensor * token_emb = build_inp_embd_mtp(model.tok_embd); + + ggml_tensor * token_emb_norm = llm_build_norm(ctx0, token_emb, hparams, mtp_layer.nextn.enorm, NULL, LLM_NORM_RMS, cb, il); + ggml_tensor * hidden_state_norm = llm_build_norm(ctx0, prev_embeddings, hparams, mtp_layer.nextn.hnorm, NULL, LLM_NORM_RMS, cb, il); + + ggml_tensor * cur; + if (mtp_layer.nextn.eh_proj != nullptr) { + ggml_tensor * combined = ggml_concat(ctx0, token_emb_norm, hidden_state_norm, 0); + cb(combined, "mtp_concat", il); + cur = llm_build_lora_mm(lctx, ctx0, mtp_layer.nextn.eh_proj, combined); + } else { + cur = ggml_add(ctx0, token_emb_norm, hidden_state_norm); + } + cb(cur, "mtp_fused", il); + + GGML_ASSERT(il < (int)kv_self.k_l.size() && il < (int)kv_self.v_l.size()); + if (!kv_self.k_l[il] || !kv_self.v_l[il]) { + LLAMA_LOG_ERROR("%s: KV cache not allocated for MTP layer %d (k=%p, v=%p)\n", + __func__, il, (void*)kv_self.k_l[il], (void*)kv_self.v_l[il]); + GGML_ABORT("KV cache not allocated for MTP layer"); + } + if (!mtp_layer.wq || !mtp_layer.wk || !mtp_layer.wv || !mtp_layer.wo) { + LLAMA_LOG_ERROR("%s: Missing attention weights for MTP layer %d (wq=%p, wk=%p, wv=%p, wo=%p)\n", + __func__, il, (void*)mtp_layer.wq, (void*)mtp_layer.wk, + (void*)mtp_layer.wv, (void*)mtp_layer.wo); + GGML_ABORT("Missing attention weights for MTP layer"); + } + + const float kq_scale = 1.0f / sqrtf(float(n_embd_head)); + + cur = build_std_attention(gf, mtp_layer.attn_norm, cur, + inp_pos, nullptr, nullptr, + KQ_mask, nullptr, nullptr, + kq_scale, 0.0f, 0, il, true, false, true, false, true, nullptr); + + if (inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + + cur = llm_build_std_moe_ffn(ctx0, lctx, mtp_layer.ffn_norm, cur, + mtp_layer.ffn_gate_inp, nullptr, + mtp_layer.ffn_up_exps, nullptr, + mtp_layer.ffn_gate_exps, nullptr, + mtp_layer.ffn_down_exps, nullptr, + nullptr, + mtp_layer.ffn_up_shexp, nullptr, + mtp_layer.ffn_gate_shexp, nullptr, + mtp_layer.ffn_down_shexp, nullptr, + n_expert, n_expert_used, + LLM_FFN_SILU, true, false, 0.0f, + LLM_EXPERT_GATING_FUNC_SOFTMAX, + LLM_FFN_SILU, cb, il, gf, true, mtp_layer.ffn_up_gate_exps, nullptr, mtp_layer.ffn_gate_inp_shexp); + + cur = lctx.cvec.apply_to(ctx0, cur, il); + cb(cur, "ffn_out", il); + + cb(cur, "result_norm", -1); + + cur = build_output(lctx, ctx0, cur, model.output, mtp_layer.nextn.shared_head_norm, cb); + cb(cur, "result_output", -1); + + return cur; +} + struct ggml_tensor * llm_build_context::build_qwen35_mtp( const llama_layer & mtp_layer, struct ggml_tensor * prev_embeddings, diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index df265893..9a3c245f 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -2059,7 +2059,8 @@ ggml_tensor * llm_build_context::build_output(llama_context & lctx, ggml_context int idx = lctx.model.default_layer_device[lctx.model.hparams.n_layer]; int idx_out = ggml_backend_sched_get_backend_idx(lctx.sched, lctx.model.output->buffer); if (idx_out >= 0) idx = idx_out; - const bool is_qwen_mtp = lctx.model.arch == LLM_ARCH_QWEN35 && lctx.cparams.mtp; + const bool is_qwen_mtp = (lctx.model.arch == LLM_ARCH_QWEN35 || + lctx.model.arch == LLM_ARCH_QWEN35MOE) && lctx.cparams.mtp; if (cur->op == GGML_OP_REDUCE && cur->src[idx] && !is_qwen_mtp) { // avoid copy to main GPU cur->view_src = cur->src[idx]; diff --git a/src/llama-build-context.h b/src/llama-build-context.h index f3b2b2fb..375a203a 100644 --- a/src/llama-build-context.h +++ b/src/llama-build-context.h @@ -470,4 +470,12 @@ llm_expert_gating_func_type gating_op, struct ggml_cgraph * gf, struct ggml_tensor * inp_pos ); + + struct ggml_tensor * build_qwen35moe_mtp( + const struct llama_layer & mtp_layer, + struct ggml_tensor * prev_embeddings, + int64_t n_embd_head, + struct ggml_cgraph * gf, + struct ggml_tensor * inp_pos + ); }; diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 7d2c4b7c..530e105e 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -495,6 +495,13 @@ void llm_load_hparams( ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, true); + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false); + if (model.mtp) { + hparams.n_layer_kv_from_start = hparams.n_layer; + } else { + hparams.n_layer_kv_from_start = hparams.n_layer - hparams.nextn_predict_layers; + } + // Load linear attention (gated delta net) parameters ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); ml.get_key(LLM_KV_SSM_INNER_SIZE, hparams.ssm_d_inner); @@ -506,13 +513,20 @@ void llm_load_hparams( { uint32_t full_attn_interval = 4; ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false); + const uint32_t n_main_layers = hparams.n_layer - hparams.nextn_predict_layers; for (uint32_t i = 0; i < hparams.n_layer; ++i) { - hparams.recurrent_layer_arr[i] = ((i + 1) % full_attn_interval != 0); + if (i < n_main_layers) { + hparams.recurrent_layer_arr[i] = ((i + 1) % full_attn_interval != 0); + } else { + hparams.recurrent_layer_arr[i] = false; + } } } switch (hparams.n_layer) { - case 40: model.type = e_model::MODEL_35B_A3B; break; + case 40: + case 41: + model.type = e_model::MODEL_35B_A3B; break; case 48: model.type = e_model::MODEL_122B_A10B; break; case 60: model.type = e_model::MODEL_397B_A17B; break; default: model.type = e_model::MODEL_UNKNOWN; diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index 58eda981..af847e4e 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -1531,49 +1531,75 @@ bool create_tensors_helper::create_qwen35moe_tensors(const LLM_TN & tn) { const int64_t conv_dim = key_dim * 2 + value_dim; for (int i = 0; i < n_layer; ++i) { - auto ctx_split = ctx_for_layer_split(i); + const bool is_mtp_layer = hparams.nextn_predict_layers > 0 && + static_cast(i) >= n_layer - hparams.nextn_predict_layers; + + auto ctx_split = is_mtp_layer ? ctx_for_layer(i) : ctx_for_layer_split(i); auto & layer = model.layers[i]; - layer.attn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0); - layer.attn_post_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), { n_embd }, 0); + int flags = 0; + if (!model.mtp && is_mtp_layer) { + flags |= llama_model_loader::TENSOR_SKIP; + } + + layer.attn_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, flags); + layer.attn_post_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), { n_embd }, flags); layer.ffn_norm = layer.attn_post_norm; if (!hparams.is_recurrent(i)) { // Attention layers - layer.wq = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q, "weight", i), { n_embd, n_embd_head_k * n_head * 2 }, 0); - layer.wk = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K, "weight", i), { n_embd, n_embd_k_gqa }, 0); - layer.wv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V, "weight", i), { n_embd, n_embd_v_gqa }, 0); - layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0); + layer.wq = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q, "weight", i), { n_embd, n_embd_head_k * n_head * 2 }, flags); + layer.wk = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K, "weight", i), { n_embd, n_embd_k_gqa }, flags); + layer.wv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_V, "weight", i), { n_embd, n_embd_v_gqa }, flags); + layer.wo = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, flags); // Q/K normalization for attention layers - layer.attn_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0); - layer.attn_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, 0); + layer.attn_q_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, flags); + layer.attn_k_norm = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, flags); } else { // Linear attention (gated delta net) specific tensors // Create tensors with calculated dimensions - layer.wqkv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_QKV, "weight", i), { n_embd, key_dim * 2 + value_dim }, 0); - layer.wqkv_gate = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_GATE, "weight", i), { n_embd, value_dim }, 0); - layer.ssm_conv1d = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_CONV1D, "weight", i), { hparams.ssm_d_conv, conv_dim }, 0); - layer.ssm_dt = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_DT, "bias", i), { hparams.ssm_dt_rank }, 0); - layer.ssm_a = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_A_NOSCAN, i), { hparams.ssm_dt_rank }, 0); - layer.ssm_beta = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_BETA, "weight", i), { n_embd, n_v_heads }, 0); - layer.ssm_alpha = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_ALPHA, "weight", i), { n_embd, n_v_heads }, 0); - layer.ssm_norm = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_NORM, "weight", i), { head_v_dim }, 0); - layer.ssm_out = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_OUT, "weight", i), { value_dim, n_embd }, 0); + layer.wqkv = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_QKV, "weight", i), { n_embd, key_dim * 2 + value_dim }, flags); + layer.wqkv_gate = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_GATE, "weight", i), { n_embd, value_dim }, flags); + layer.ssm_conv1d = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_CONV1D, "weight", i), { hparams.ssm_d_conv, conv_dim }, flags); + layer.ssm_dt = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_DT, "bias", i), { hparams.ssm_dt_rank }, flags); + layer.ssm_a = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_A_NOSCAN, i), { hparams.ssm_dt_rank }, flags); + layer.ssm_beta = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_BETA, "weight", i), { n_embd, n_v_heads }, flags); + layer.ssm_alpha = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_ALPHA, "weight", i), { n_embd, n_v_heads }, flags); + layer.ssm_norm = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_NORM, "weight", i), { head_v_dim }, flags); + layer.ssm_out = create_tensor(ctx_split, tn(LLM_TENSOR_SSM_OUT, "weight", i), { value_dim, n_embd }, flags); } - layer.ffn_gate_inp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert }, 0); - use_mmap_buffer &= !create_std_ffn_exps(n_embd, tn, i, 0, n_ff_exp); + layer.ffn_gate_inp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert }, flags); + use_mmap_buffer &= !create_std_ffn_exps(n_embd, tn, i, flags, n_ff_exp); // Shared experts const int64_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff; - layer.ffn_gate_inp_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", i), { n_embd }, 0); - layer.ffn_gate_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), { n_embd, n_ff_shexp }, 0); - layer.ffn_up_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), { n_embd, n_ff_shexp }, 0); - layer.ffn_down_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_shexp, n_embd }, 0); + layer.ffn_gate_inp_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", i), { n_embd }, flags); + layer.ffn_gate_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), { n_embd, n_ff_shexp }, flags); + layer.ffn_up_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), { n_embd, n_ff_shexp }, flags); + layer.ffn_down_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_shexp, n_embd }, flags); + if (is_mtp_layer) { + layer.nextn.eh_proj = create_tensor(ctx_split, + tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), + { 2 * n_embd, n_embd }, + flags | llama_model_loader::TENSOR_NOT_REQUIRED); + layer.nextn.enorm = create_tensor(ctx_split, + tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), + { n_embd }, + flags); + layer.nextn.hnorm = create_tensor(ctx_split, + tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), + { n_embd }, + flags); + layer.nextn.shared_head_norm = create_tensor(ctx_split, + tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), + { n_embd }, + flags | llama_model_loader::TENSOR_NOT_REQUIRED); + } } return use_mmap_buffer; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 10d02475..8cc7f6ae 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -503,6 +503,10 @@ static const std::map> LLM_TENSOR_NA { LLM_TENSOR_FFN_GATE_SHEXP, "blk.%d.ffn_gate_shexp" }, { LLM_TENSOR_FFN_DOWN_SHEXP, "blk.%d.ffn_down_shexp" }, { LLM_TENSOR_FFN_UP_SHEXP, "blk.%d.ffn_up_shexp" }, + { 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" }, }, }, { diff --git a/src/llama.cpp b/src/llama.cpp index 5cfa1216..02915969 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -803,7 +803,8 @@ static bool llama_kv_cache_init( // count used buffer types std::map buft_layer_count; if (offload) { - const bool qwen_mtp = model.arch == LLM_ARCH_QWEN35 && hparams.nextn_predict_layers > 0; + const bool qwen_mtp = (model.arch == LLM_ARCH_QWEN35 || + model.arch == LLM_ARCH_QWEN35MOE) && hparams.nextn_predict_layers > 0; const int64_t n_mtp_first = n_layer - hparams.nextn_predict_layers; for (int64_t i = 0; i < n_layer; ++i) { const bool is_mtp_tail = qwen_mtp && i >= n_mtp_first; @@ -893,7 +894,8 @@ static bool llama_kv_cache_init( const uint32_t n_head_kv = hparams.n_head_kv(i); const uint32_t n_embd_head_k= hparams.n_embd_head_k(i); - const bool is_mtp_tail_layer = model.arch == LLM_ARCH_QWEN35 && + const bool is_mtp_tail_layer = (model.arch == LLM_ARCH_QWEN35 || + model.arch == LLM_ARCH_QWEN35MOE) && hparams.nextn_predict_layers > 0 && i >= (int)n_mtp_first_layer; //struct ggml_context * ctx = split_cache && !qnext_recurrent ? ctx_map.at(model.buft_layer[i].buft_matrix) : offload ? ctx_map.at(model.buft_layer[i].buft) : cache.ctxs.front(); struct ggml_context * ctx = (split_cache && !is_mtp_tail_layer) ? ctx_map.at(model.buft_layer[i].buft_matrix) : offload ? ctx_map.at(model.buft_layer[i].buft) : cache.ctxs.front(); @@ -4528,7 +4530,8 @@ static int llama_decode_internal( } else { const bool has_mtp = lctx.model.hparams.nextn_predict_layers > 0 && lctx.model.mtp; - const bool use_qwen_mtp_embd = has_mtp && lctx.model.arch == LLM_ARCH_QWEN35; + const bool use_qwen_mtp_embd = has_mtp && (lctx.model.arch == LLM_ARCH_QWEN35 || + lctx.model.arch == LLM_ARCH_QWEN35MOE); if (cparams.embeddings || has_mtp) { for (int i = gf->n_nodes - 1; i >= 0; --i) { if (use_qwen_mtp_embd && strcmp(gf->nodes[i]->name, "result_mtp_embd") == 0) { @@ -6012,7 +6015,8 @@ struct llama_context * llama_init_from_model( } } - if (model->arch != LLM_ARCH_GLM4_MOE && model->arch != LLM_ARCH_QWEN35 && cparams.mtp != 0) { + if (model->arch != LLM_ARCH_GLM4_MOE && model->arch != LLM_ARCH_QWEN35 && + model->arch != LLM_ARCH_QWEN35MOE && cparams.mtp != 0) { cparams.mtp = 0; }