From 67e634622534f06b2df5ee7332823acb669c3202 Mon Sep 17 00:00:00 2001 From: Samuel Oliveira Alves <107287165+SamuelOliveirads@users.noreply.github.com> Date: Tue, 28 Apr 2026 02:47:50 -0300 Subject: [PATCH] Support for Qwen 3.5 MTP (dense models only) (#1698) * qwen-mtp: add dense mtp for one draft * add support for smaller qwen mtp commit * qwen-mtp: fix graph for qwen dense variants * Squashed commit of the following: commit a92a154b38c7fddc84460f8852c900f8d6ce907e Author: SamuelOliveirads Date: Mon Apr 20 13:30:21 2026 -0300 recurrent model: refactor api commit dfac8f19f6edc0014b4116041b89c1e0dfb173c7 Author: SamuelOliveirads Date: Mon Apr 20 12:22:29 2026 -0300 recurrent model: implement recurrent kernel checkpoint commit 9c44b117f93e9060030907e1250106358c9ccf47 Author: SamuelOliveirads Date: Sat Apr 18 11:52:39 2026 -0300 speculative: fix sampler for checkpoints commit e7006393bca20adcd86e2d77021e0b41d7bd9db1 Author: SamuelOliveirads Date: Fri Apr 17 14:08:25 2026 -0300 server: refactor checkpoint state logic commit 57eabf04df5185cab19539185b8f4b85e578905b Merge: dc4797b7 64234e3c Author: SamuelOliveirads Date: Fri Apr 17 13:53:41 2026 -0300 Merge branch 'main' into fix/hybrid-cache-speculative commit dc4797b72363482bb35750bb5edc87068116dc0f Author: SamuelOliveirads Date: Fri Apr 17 13:12:40 2026 -0300 reset ngram mod state for rejected tokens commit 8ff2d943a31b3d54440698db41e62ea121d661be Author: SamuelOliveirads Date: Fri Apr 17 13:08:04 2026 -0300 server: snapshot recurrent state in tensor commit d93dfb5e6b78a822e7331b33feedbdc47eb5ec79 Author: SamuelOliveirads Date: Thu Apr 16 22:36:37 2026 -0300 fix: save/restore sampler state during speculative checkpoint When speculative decoding rejects draft tokens and restores the recurrent state checkpoint, the sampler (RNG, grammar, prev tokens) must also be restored to maintain consistency. Without this, the sampler state reflects the rejected draft tokens, leading to potential divergence. Uses common_sampler_clone() to snapshot the sampler before the speculative batch decode, and restores it on rejection. commit d670cf85cd59f23a339faf6fd773889063176801 Author: SamuelOliveirads Date: Thu Apr 16 21:53:52 2026 -0300 server: spec checkpoints for recurrent models * server: fix leak context between requests * qwen3: allow mtp to run with split graph * qwen3 mtp: selects rows before the ffn --- examples/server/server-context.cpp | 100 +++++++++++------ ggml/src/ggml-alloc.c | 8 +- src/graphs/build_glm4.cpp | 8 +- src/graphs/build_qwen35.cpp | 168 +++++++++++++++++++++++------ src/llama-build-context.cpp | 6 +- src/llama-build-context.h | 10 +- src/llama-hparams.cpp | 28 ++++- src/llama-load-tensors.cpp | 121 ++++++++++++++++----- src/llama-model.cpp | 50 +++++---- src/llama.cpp | 33 +++++- 10 files changed, 401 insertions(+), 131 deletions(-) diff --git a/examples/server/server-context.cpp b/examples/server/server-context.cpp index e5e1ec1d..782d1547 100644 --- a/examples/server/server-context.cpp +++ b/examples/server/server-context.cpp @@ -3265,11 +3265,14 @@ void server_context::apply_checkpoint(server_slot & slot) { if (do_reset) { if (has_recurrent) { - // Hybrid/recurrent: do NOT zero n_past. The prompt prefix is already in cache_tokens - // and update_slots() reprocesses from slot.n_past_prompt; dropping to 0 forces a full - // recompute on every turn and — combined with cached state — trips llama_decode ret=-3. - SLT_WRN(slot, "no usable hybrid/recurrent checkpoint; preserving slot state (n_past = %d, n_past_prompt = %d)\n", - (int)slot.n_past, (int)slot.n_past_prompt); + // Without a usable recurrent checkpoint, preserving prefix state leaks stale recurrent memory + // from prior requests into the current prompt. Force a full prompt re-processing fallback. + SLT_WRN(slot, "%s", "no usable hybrid/recurrent checkpoint; forcing full prompt re-processing\n"); + slot.n_past = 0; + slot.n_past_prompt = 0; + slot.n_past_se = 0; + slot.ga_i = 0; + common_sampler_reset(slot.ctx_sampling); } else { SLT_WRN(slot, "forcing full prompt re-processing due to lack of cache data (likely due to SWA, see %s)\n", "https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055"); @@ -3720,7 +3723,8 @@ void server_context::extend_context(const int32_t n_tokens) { // Restore recurrent state and re-decode accepted tokens after speculative-decode rejection. static void restore_speculative_checkpoint( server_slot & slot, llama_context * ctx, llama_model * model, - const std::vector & ids, int n_draft) { + const std::vector & ids, int n_draft, + const std::vector & mtp_hidden_state_pre, int32_t mtp_n_past_base) { if (slot.spec_ckpt.per_step_enabled) { const int step = (int)ids.size() - 1; llama_spec_ckpt_restore(ctx, slot.id, slot.spec_ckpt.n_past, step); @@ -3732,6 +3736,15 @@ static void restore_speculative_checkpoint( common_sampler_accept(slot.ctx_sampling, ctx, id, true); } + // Update MTP KV cache and hidden state using embeddings collected before checkpoint restore. + if (slot.has_mtp && !mtp_hidden_state_pre.empty()) { + slot.mtp_hidden_state = mtp_hidden_state_pre; + llama_context * mtp_ctx = common_speculative_get_mtp_ctx(slot.spec); + llama_context * mtp_target = mtp_ctx ? mtp_ctx : ctx; + llama_set_draft_input_hidden_state(mtp_target, slot.mtp_hidden_state.data()); + mtp_accept_tokens(mtp_target, ids, mtp_n_past_base, slot.id); + } + SLT_DBG(slot, "per-step restore: step=%d (rejected %d drafts)\n", step, (int)(n_draft - (ids.size() - 1))); } else { @@ -3752,6 +3765,9 @@ static void restore_speculative_checkpoint( } if (slot.has_mtp) { + for (int j = 0; j < re_batch.n_tokens; j++) { + re_batch.logits[j] = true; + } llama_set_embeddings(ctx, true); } @@ -3759,15 +3775,29 @@ static void restore_speculative_checkpoint( if (ret != 0) { SLT_ERR(slot, "failed to re-decode accepted tokens after checkpoint restore: %d\n", ret); } - if (slot.has_mtp) { - llama_set_embeddings(ctx, false); const int n_embd = llama_model_n_embd(llama_get_model(ctx)); - const float * emb = llama_get_embeddings_ith(ctx, -1); - if (emb) { - slot.mtp_hidden_state.resize(n_embd); - memcpy(slot.mtp_hidden_state.data(), emb, n_embd * sizeof(float)); + + const int n_accepted = (int)ids.size(); + slot.mtp_hidden_state.resize(n_accepted * n_embd); + for (int j = 0; j < n_accepted; j++) { + const float * emb_j = llama_get_embeddings_ith(ctx, j); + if (emb_j) { + memcpy(slot.mtp_hidden_state.data() + j * n_embd, emb_j, n_embd * sizeof(float)); + } } + + llama_context * mtp_ctx_rej = common_speculative_get_mtp_ctx(slot.spec); + llama_context * mtp_target_rej = mtp_ctx_rej ? mtp_ctx_rej : ctx; + llama_set_draft_input_hidden_state(mtp_target_rej, slot.mtp_hidden_state.data()); + mtp_accept_tokens(mtp_target_rej, ids, slot.spec_ckpt.n_past, slot.id); + + if (n_accepted > 1) { + memmove(slot.mtp_hidden_state.data(), + slot.mtp_hidden_state.data() + (n_accepted - 1) * n_embd, + n_embd * sizeof(float)); + } + slot.mtp_hidden_state.resize(n_embd); } for (llama_token id : ids) { @@ -3795,30 +3825,28 @@ void server_context::speculative_decoding_accept() { // the accepted tokens from the speculation const auto ids = common_sampler_sample_and_accept_n(slot.ctx_sampling, ctx, slot.i_batch_dft, slot.drafted); - + + int32_t mtp_n_past_base = 0; + std::vector mtp_hidden_state_pre; if (slot.has_mtp) { - llama_context * mtp_ctx = common_speculative_get_mtp_ctx(slot.spec); - llama_context * mtp_target = mtp_ctx ? mtp_ctx : ctx; + mtp_n_past_base = slot.n_past - (slot.drafted.size() + 1); const int n_embd = llama_model_n_embd(llama_get_model(ctx)); if (!ids.empty()) { - const float* emb = llama_get_embeddings(ctx); - if (emb) { - slot.mtp_hidden_state.resize(ids.size() * n_embd); - memcpy(slot.mtp_hidden_state.data(), emb, ids.size() * n_embd * sizeof(float)); + mtp_hidden_state_pre.resize(ids.size() * n_embd); + for (size_t i = 0; i < ids.size(); i++) { + const float* emb_i = llama_get_embeddings_ith(ctx, slot.i_batch_dft[i]); + if (emb_i) { + memcpy(mtp_hidden_state_pre.data() + i * n_embd, emb_i, n_embd * sizeof(float)); + } } } else { const float* emb0 = llama_get_embeddings_ith(ctx, 0); if (emb0) { - slot.mtp_hidden_state.resize(n_embd); - memcpy(slot.mtp_hidden_state.data(), emb0, n_embd * sizeof(float)); + mtp_hidden_state_pre.resize(n_embd); + memcpy(mtp_hidden_state_pre.data(), emb0, n_embd * sizeof(float)); } } - - llama_set_draft_input_hidden_state(mtp_target, slot.mtp_hidden_state.data()); - - int32_t n_past_base = slot.n_past - (slot.drafted.size() + 1); - mtp_accept_tokens(mtp_target, ids, n_past_base, slot.id); } slot.i_batch_dft.clear(); @@ -3846,8 +3874,16 @@ void server_context::speculative_decoding_accept() { // for recurrent/hybrid models: if any drafts were rejected, restore recurrent state const bool any_rejected = (ids.size() - 1) < n_draft; if (any_rejected && slot.spec_ckpt.valid) { - restore_speculative_checkpoint(slot, ctx, model, ids, n_draft); + restore_speculative_checkpoint(slot, ctx, model, ids, n_draft, mtp_hidden_state_pre, mtp_n_past_base); } else { + if (slot.has_mtp && !mtp_hidden_state_pre.empty()) { + llama_context * mtp_ctx = common_speculative_get_mtp_ctx(slot.spec); + llama_context * mtp_target = mtp_ctx ? mtp_ctx : ctx; + + slot.mtp_hidden_state = std::move(mtp_hidden_state_pre); + llama_set_draft_input_hidden_state(mtp_target, slot.mtp_hidden_state.data()); + mtp_accept_tokens(mtp_target, ids, mtp_n_past_base, slot.id); + } llama_kv_cache_seq_rm(ctx, slot.id, slot.n_past, -1); discard_speculative_checkpoint(slot, ctx); } @@ -4232,12 +4268,14 @@ void server_context::process_batch_tokens(int32_t & n_batch) { } } if (mtp_warmup_needed) { - const float* emb = llama_get_embeddings(ctx); const int n_embd = llama_model_n_embd(llama_get_model(ctx)); const int n_toks = batch_view.n_tokens; - if (emb) { - batch_mtp_hidden_state.resize(n_toks * n_embd); - memcpy(batch_mtp_hidden_state.data(), emb, n_toks * n_embd * sizeof(float)); + batch_mtp_hidden_state.resize(n_toks * n_embd); + for (int t = 0; t < n_toks; t++) { + const float* emb_t = llama_get_embeddings_ith(ctx, t); + if (emb_t) { + memcpy(batch_mtp_hidden_state.data() + t * n_embd, emb_t, n_embd * sizeof(float)); + } } } } diff --git a/ggml/src/ggml-alloc.c b/ggml/src/ggml-alloc.c index 5bdfa942..70433f98 100644 --- a/ggml/src/ggml-alloc.c +++ b/ggml/src/ggml-alloc.c @@ -815,7 +815,13 @@ static void ggml_gallocr_init_tensor(ggml_gallocr_t galloc, struct ggml_tensor * } static bool ggml_gallocr_node_needs_realloc(ggml_gallocr_t galloc, struct ggml_tensor * node, struct tensor_alloc * talloc) { - size_t node_size = (node->data || node->view_src) ? 0 : ggml_backend_buft_get_alloc_size(galloc->bufts[talloc->buffer_id], node); + if (node->data || node->view_src) { + return true; + } + if (talloc->buffer_id < 0 || talloc->buffer_id >= galloc->n_buffers) { + return false; + } + size_t node_size = ggml_backend_buft_get_alloc_size(galloc->bufts[talloc->buffer_id], node); return talloc->size_max >= node_size; } diff --git a/src/graphs/build_glm4.cpp b/src/graphs/build_glm4.cpp index f60c8258..a0dbc455 100644 --- a/src/graphs/build_glm4.cpp +++ b/src/graphs/build_glm4.cpp @@ -26,7 +26,7 @@ ggml_cgraph * llm_build_context::build_glm4_moe() { } else { hidden_states_from_main_model = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, hparams.n_embd); } - ggml_set_name(hidden_states_from_main_model, "result_embd_pooled"); + 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; @@ -34,7 +34,7 @@ ggml_cgraph * llm_build_context::build_glm4_moe() { const int il_mtp = hparams.n_layer - 1; const auto & mtp_layer = model.layers[il_mtp]; - cur = build_mtp_tail(mtp_layer, hidden_states_from_main_model, n_embd_head, gf, inp_pos, rope_cache); + cur = build_glm4_moe_mtp(mtp_layer, hidden_states_from_main_model, n_embd_head, gf, inp_pos, rope_cache); } else { struct ggml_tensor * inpL; @@ -288,7 +288,7 @@ ggml_cgraph * llm_build_context::build_glm4() { return gf; } -struct ggml_tensor * llm_build_context::build_mtp_tail( +struct ggml_tensor * llm_build_context::build_glm4_moe_mtp( const llama_layer & mtp_layer, struct ggml_tensor * prev_embeddings, int64_t n_embd_head, @@ -350,7 +350,7 @@ struct ggml_tensor * llm_build_context::build_mtp_tail( cb(ffn_inp, "mtp_ffn_inp", il); } - // FFN + // MoE FFN cur = llm_build_std_moe_ffn(ctx0, lctx, mtp_layer.ffn_norm, ffn_inp, mtp_layer.ffn_gate_inp, NULL, mtp_layer.ffn_up_exps, NULL, diff --git a/src/graphs/build_qwen35.cpp b/src/graphs/build_qwen35.cpp index fa1e6f5f..fb19d679 100644 --- a/src/graphs/build_qwen35.cpp +++ b/src/graphs/build_qwen35.cpp @@ -66,51 +66,155 @@ ggml_cgraph * llm_build_context::build_qwen35() { 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 * cur; + 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) { + // MTP tail-only graph + 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_qwen35_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; + + cur = nullptr; + + 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_ffn(ctx0, lctx, model.layers[il].ffn_norm, cur, + model.layers[il].ffn_up, NULL, NULL, + model.layers[il].ffn_gate, NULL, NULL, + model.layers[il].ffn_down, NULL, NULL, + NULL, + LLM_FFN_SILU, LLM_FFN_PAR, cb, il, gf, true, false); + + cur = lctx.cvec.apply_to(ctx0, cur, il); + cb(cur, "l_out", il); + + inpL = cur; } - cur = llm_build_ffn(ctx0, lctx, model.layers[il].ffn_norm, cur, - model.layers[il].ffn_up, NULL, NULL, - model.layers[il].ffn_gate, NULL, NULL, - model.layers[il].ffn_down, NULL, NULL, - NULL, - LLM_FFN_SILU, LLM_FFN_PAR, cb, il, gf, true, false); + if (lctx.cparams.mtp) { + struct ggml_tensor * embd_copy = ggml_dup(ctx0, inpL); + cb(embd_copy, "result_mtp_embd", -1); + ggml_set_output(embd_copy); + } - 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; } +struct ggml_tensor * llm_build_context::build_qwen35_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_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) { + // Full fusion: concat + project (27B, 4B, 2B, 0.8B) + 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 { + // 9B — no fc/eh_proj + cur = ggml_add(ctx0, token_emb_norm, hidden_state_norm); + } + cb(cur, "mtp_fused", il); + + // Self-Attention (wq may be shared from main model's last layer) + 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 (!model.layers[il].wq || !model.layers[il].wk || !model.layers[il].wv || !model.layers[il].wo) { + LLAMA_LOG_ERROR("%s: Missing attention weights for MTP layer %d (wq=%p, wk=%p, wv=%p, wo=%p)\n", + __func__, il, (void*)model.layers[il].wq, (void*)model.layers[il].wk, + (void*)model.layers[il].wv, (void*)model.layers[il].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); + } + + // Dense FFN — optional (9B and 4B don't have FFN in MTP layer) + if (mtp_layer.ffn_gate != nullptr) { + cur = llm_build_ffn(ctx0, lctx, mtp_layer.ffn_norm, cur, + mtp_layer.ffn_up, NULL, NULL, + mtp_layer.ffn_gate, NULL, NULL, + mtp_layer.ffn_down, NULL, NULL, + NULL, + LLM_FFN_SILU, LLM_FFN_PAR, cb, il, gf, true, false); + } + + cur = lctx.cvec.apply_to(ctx0, cur, il); + cb(cur, "ffn_out", il); + + cur = llm_build_norm(ctx0, cur, hparams, mtp_layer.nextn.shared_head_norm, NULL, LLM_NORM_RMS, cb, il); + cb(cur, "result_norm", -1); + + cur = build_output(lctx, ctx0, cur, model.output, nullptr, cb); + cb(cur, "result_output", -1); + + return cur; +} \ No newline at end of file diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index 351efe07..087114be 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -1558,6 +1558,9 @@ static ggml_tensor * llm_build_kqv( auto v_cache = lctx.model.hparams.has_kv(il) ? kv.v_l[il] : lctx.model.hparams.swa_layers[il] ? kv.v_l[hparams.n_layer_kv_from_start-2] : kv.v_l[hparams.n_layer_kv_from_start-1]; + GGML_ASSERT(k_cache != nullptr && "k_cache is null in llm_build_kqv"); + GGML_ASSERT(v_cache != nullptr && "v_cache is null in llm_build_kqv"); + struct ggml_tensor * k = ggml_view_3d(ctx, k_cache, n_embd_head_k, n_kv, n_head_kv, @@ -2056,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; - if (cur->op == GGML_OP_REDUCE && cur->src[idx]) { + const bool is_qwen_mtp = lctx.model.arch == LLM_ARCH_QWEN35 && 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 cd9f35e1..f3b2b2fb 100644 --- a/src/llama-build-context.h +++ b/src/llama-build-context.h @@ -454,7 +454,7 @@ llm_expert_gating_func_type gating_op, static uint32_t llama_kv_qnext_state_slots(const llama_kv_cache & kv_self); - struct ggml_tensor * build_mtp_tail( + struct ggml_tensor * build_glm4_moe_mtp( const struct llama_layer & mtp_layer, struct ggml_tensor * prev_embeddings, int64_t n_embd_head, @@ -462,4 +462,12 @@ llm_expert_gating_func_type gating_op, struct ggml_tensor * inp_pos, struct ggml_tensor * rope_cache ); + + struct ggml_tensor * build_qwen35_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 ad13fe7b..7053952c 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -523,6 +523,14 @@ void llm_load_hparams( ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, true); + // NextN/MTP parameters + 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); @@ -531,18 +539,30 @@ void llm_load_hparams( ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group); // Mark recurrent layers (linear attention layers) + // MTP layers always use standard attention, not delta-net { 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 24: model.type = hparams.n_embd == 1024 ? e_model::MODEL_0_8B : e_model::MODEL_2B; break; - case 32: model.type = hparams.n_embd == 2560 ? e_model::MODEL_4B : e_model::MODEL_9B; break; - case 64: model.type = e_model::MODEL_27B; break; + case 24: // without MTP layer + case 25: // with MTP layer (24 main + 1 MTP) + model.type = hparams.n_embd == 1024 ? e_model::MODEL_0_8B : e_model::MODEL_2B; break; + case 32: // without MTP layer + case 33: // with MTP layer (32 main + 1 MTP) + model.type = hparams.n_embd == 2560 ? e_model::MODEL_4B : e_model::MODEL_9B; break; + case 64: // without MTP layer + case 65: // with MTP layer (64 main + 1 MTP) + model.type = e_model::MODEL_27B; break; default: model.type = e_model::MODEL_UNKNOWN; } } break; diff --git a/src/llama-load-tensors.cpp b/src/llama-load-tensors.cpp index 7d7b385b..58eda981 100644 --- a/src/llama-load-tensors.cpp +++ b/src/llama-load-tensors.cpp @@ -1601,42 +1601,90 @@ bool create_tensors_helper::create_qwen35_tensors(const LLM_TN & tn) { const int64_t conv_dim = key_dim * 2 + value_dim; for (int i = 0; i < n_layer; ++i) { - ggml_context * ctx_split = 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); + const bool is_mtp_layer = hparams.nextn_predict_layers > 0 && + static_cast(i) >= n_layer - hparams.nextn_predict_layers; + + // For now only run MTP into the per-layer + ggml_context * ctx_split = is_mtp_layer ? ctx_for_layer(i) : ctx_for_layer_split(i); + + int flags = 0; + // Skip loading MTP layers if the feature is disabled + if (!model.mtp) { + if (is_mtp_layer) { + flags |= llama_model_loader::TENSOR_SKIP; + } + } + const int mtp_opt = is_mtp_layer ? llama_model_loader::TENSOR_NOT_REQUIRED : 0; + + 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); + // Attention layers (MTP layer is always standard attention) + layer.wq = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_Q, "weight", i), { n_embd, n_embd_head_k * n_head * 2 }, flags | mtp_opt); + 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 }, llama_model_loader::TENSOR_NOT_REQUIRED); - layer.wqkv_gate = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_GATE, "weight", i), { n_embd, value_dim }, llama_model_loader::TENSOR_NOT_REQUIRED); - 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 }, llama_model_loader::TENSOR_NOT_REQUIRED | flags); + layer.wqkv_gate = create_tensor(ctx_split, tn(LLM_TENSOR_ATTN_GATE, "weight", i), { n_embd, value_dim }, llama_model_loader::TENSOR_NOT_REQUIRED | 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 = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, 0); - layer.ffn_down = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0); - layer.ffn_up = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, 0); + // 9B and 4B don't have MLP in MTP + layer.ffn_gate = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, flags | mtp_opt); + layer.ffn_down = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, flags | mtp_opt); + layer.ffn_up = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, flags | mtp_opt); + // --- NextN / MTP tensors on the MTP layer --- + if (is_mtp_layer) { + const int final_layer = n_layer - 1; + auto nextn_ctx = ctx_for_layer(final_layer); + // 9B doesn't have fc + layer.nextn.eh_proj = create_tensor(nextn_ctx, + tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", final_layer), + { 2*n_embd, n_embd }, + flags | llama_model_loader::TENSOR_NOT_REQUIRED); + layer.nextn.enorm = create_tensor(nextn_ctx, + tn(LLM_TENSOR_NEXTN_ENORM, "weight", final_layer), + { n_embd }, + flags); + layer.nextn.hnorm = create_tensor(nextn_ctx, + tn(LLM_TENSOR_NEXTN_HNORM, "weight", final_layer), + { n_embd }, + flags); + layer.nextn.shared_head_norm = create_tensor(nextn_ctx, + tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", final_layer), + { n_embd }, + flags | llama_model_loader::TENSOR_NOT_REQUIRED); + } + } + + // 9B shares q_proj + if (model.mtp && hparams.nextn_predict_layers > 0) { + const uint32_t n_main = n_layer - hparams.nextn_predict_layers; + for (uint32_t i = n_main; i < (uint32_t)n_layer; ++i) { + auto & mtp_layer = model.layers[i]; + auto & last_main = model.layers[n_main - 1]; + if (mtp_layer.wq == nullptr) { + mtp_layer.wq = last_main.wq; + } + } } return use_mmap_buffer; @@ -2447,8 +2495,11 @@ bool create_tensors_helper::create_glm_dsa_tensors(const LLM_TN & tn) { } for (int i = 0; i < n_layer; ++i) { + const bool is_mtp_layer = hparams.nextn_predict_layers > 0 && + static_cast(i) >= n_layer - hparams.nextn_predict_layers; + int flags = 0; - if (hparams.nextn_predict_layers > 0 && static_cast(i) >= n_layer - hparams.nextn_predict_layers) { + if (is_mtp_layer) { flags |= llama_model_loader::TENSOR_SKIP | llama_model_loader::TENSOR_NOT_REQUIRED; } ggml_context * ctx_layer = ctx_for_layer(i); @@ -2526,7 +2577,7 @@ bool create_tensors_helper::create_glm_dsa_tensors(const LLM_TN & tn) { layer.ffn_up_shexp = create_tensor(ctx_split, tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags); } - if (hparams.nextn_predict_layers > 0 && static_cast(i) >= n_layer - hparams.nextn_predict_layers) { + 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); 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); @@ -2554,10 +2605,13 @@ bool create_tensors_helper::create_glm4_moe_tensors(const LLM_TN & tn) { ggml_context * ctx_layer = ctx_for_layer(i); ggml_context * 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; + int flags = 0; // Skip loading MTP layers if the feature is disabled if (!model.mtp) { - if (hparams.nextn_predict_layers > 0 && static_cast(i) >= n_layer - hparams.nextn_predict_layers) { + if (is_mtp_layer) { // skip all tensors in the NextN layers flags |= llama_model_loader::TENSOR_SKIP; } @@ -2625,7 +2679,7 @@ bool create_tensors_helper::create_glm4_moe_tensors(const LLM_TN & tn) { layer.ffn_up = create_tensor(ffn_ctx, tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, flags); } // --- NextN / MTP tensors on the final layer --- - if (hparams.nextn_predict_layers > 0 && static_cast(i) >= n_layer - hparams.nextn_predict_layers) { + if (is_mtp_layer) { const int final_layer = n_layer - 1; auto nextn_ctx = ctx_for_layer(final_layer); auto nextn_host_ctx = ctx_input; @@ -3087,8 +3141,11 @@ bool create_tensors_helper::create_bailingmoe2_tensors(const LLM_TN & tn) { ggml_context * ctx_layer = ctx_for_layer(i); ggml_context * 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; + int flags = 0; - if (hparams.nextn_predict_layers > 0 && static_cast(i) >= n_layer - hparams.nextn_predict_layers) { + if (is_mtp_layer) { // skip all tensors in the NextN layers flags |= llama_model_loader::TENSOR_SKIP; } @@ -3122,7 +3179,7 @@ bool create_tensors_helper::create_bailingmoe2_tensors(const LLM_TN & tn) { } // NextN/MTP tensors (preserved but unused) - conditionally load for last nextn_predict_layers - if (hparams.nextn_predict_layers > 0 && static_cast(i) >= n_layer - hparams.nextn_predict_layers) { + 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); layer.nextn.embed_tokens = create_tensor(ctx_split, tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), { n_embd, n_vocab }, llama_model_loader::TENSOR_NOT_REQUIRED | flags); @@ -4102,6 +4159,12 @@ bool create_tensors_helper::create_tensors() { gpu_split_count.resize(model.splits.size(), 0.0f); } for (int il = 0; il < n_layer; ++il) { + // For now only run MTP into the per-layer + if (model.mtp && hparams.nextn_predict_layers > 0 && + static_cast(il) >= static_cast(n_layer) - hparams.nextn_predict_layers) { + LLAMA_LOG_DEBUG("%s: not splitting MTP tail layer %d (forced non-split)\n", __func__, il); + continue; + } int gqa_ratio = hparams.n_head(il) / hparams.n_head_kv(il); if (ggml_backend_buft_is_host(model.buft_layer[il].buft_matrix)) { LLAMA_LOG_INFO("%s: not splitting layer %d because buffer type is host\n", __func__, il); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index d5088eea..4840d44c 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -508,29 +508,33 @@ static const std::map> LLM_TENSOR_NA { LLM_ARCH_QWEN35, { - { LLM_TENSOR_TOKEN_EMBD, "token_embd" }, - { LLM_TENSOR_OUTPUT_NORM, "output_norm" }, - { LLM_TENSOR_OUTPUT, "output" }, - { LLM_TENSOR_ATTN_NORM, "blk.%d.attn_norm" }, - { LLM_TENSOR_ATTN_POST_NORM, "blk.%d.post_attention_norm" }, - { LLM_TENSOR_ATTN_Q, "blk.%d.attn_q" }, - { LLM_TENSOR_ATTN_Q_NORM, "blk.%d.attn_q_norm" }, - { LLM_TENSOR_ATTN_K, "blk.%d.attn_k" }, - { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, - { LLM_TENSOR_ATTN_V, "blk.%d.attn_v" }, - { LLM_TENSOR_ATTN_OUT, "blk.%d.attn_output" }, - { LLM_TENSOR_ATTN_QKV, "blk.%d.attn_qkv" }, - { LLM_TENSOR_ATTN_GATE, "blk.%d.attn_gate" }, - { LLM_TENSOR_SSM_CONV1D, "blk.%d.ssm_conv1d" }, - { LLM_TENSOR_SSM_DT, "blk.%d.ssm_dt" }, - { LLM_TENSOR_SSM_A_NOSCAN, "blk.%d.ssm_a" }, - { LLM_TENSOR_SSM_BETA, "blk.%d.ssm_beta" }, - { LLM_TENSOR_SSM_ALPHA, "blk.%d.ssm_alpha" }, - { LLM_TENSOR_SSM_NORM, "blk.%d.ssm_norm" }, - { LLM_TENSOR_SSM_OUT, "blk.%d.ssm_out" }, - { LLM_TENSOR_FFN_GATE, "blk.%d.ffn_gate" }, - { LLM_TENSOR_FFN_DOWN, "blk.%d.ffn_down" }, - { LLM_TENSOR_FFN_UP, "blk.%d.ffn_up" }, + { LLM_TENSOR_TOKEN_EMBD, "token_embd" }, + { LLM_TENSOR_OUTPUT_NORM, "output_norm" }, + { LLM_TENSOR_OUTPUT, "output" }, + { LLM_TENSOR_ATTN_NORM, "blk.%d.attn_norm" }, + { LLM_TENSOR_ATTN_POST_NORM, "blk.%d.post_attention_norm" }, + { LLM_TENSOR_ATTN_Q, "blk.%d.attn_q" }, + { LLM_TENSOR_ATTN_Q_NORM, "blk.%d.attn_q_norm" }, + { LLM_TENSOR_ATTN_K, "blk.%d.attn_k" }, + { LLM_TENSOR_ATTN_K_NORM, "blk.%d.attn_k_norm" }, + { LLM_TENSOR_ATTN_V, "blk.%d.attn_v" }, + { LLM_TENSOR_ATTN_OUT, "blk.%d.attn_output" }, + { LLM_TENSOR_ATTN_QKV, "blk.%d.attn_qkv" }, + { LLM_TENSOR_ATTN_GATE, "blk.%d.attn_gate" }, + { LLM_TENSOR_SSM_CONV1D, "blk.%d.ssm_conv1d" }, + { LLM_TENSOR_SSM_DT, "blk.%d.ssm_dt" }, + { LLM_TENSOR_SSM_A_NOSCAN, "blk.%d.ssm_a" }, + { LLM_TENSOR_SSM_BETA, "blk.%d.ssm_beta" }, + { LLM_TENSOR_SSM_ALPHA, "blk.%d.ssm_alpha" }, + { LLM_TENSOR_SSM_NORM, "blk.%d.ssm_norm" }, + { LLM_TENSOR_SSM_OUT, "blk.%d.ssm_out" }, + { LLM_TENSOR_FFN_GATE, "blk.%d.ffn_gate" }, + { LLM_TENSOR_FFN_DOWN, "blk.%d.ffn_down" }, + { LLM_TENSOR_FFN_UP, "blk.%d.ffn_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" }, }, }, { diff --git a/src/llama.cpp b/src/llama.cpp index 7a832716..f7b55bbf 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -778,8 +778,11 @@ 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 int64_t n_mtp_first = n_layer - hparams.nextn_predict_layers; for (int64_t i = 0; i < n_layer; ++i) { - if (split_cache) { + const bool is_mtp_tail = qwen_mtp && i >= n_mtp_first; + if (split_cache && !is_mtp_tail) { buft_layer_count[model.buft_layer[i].buft_matrix]++; } else { buft_layer_count[model.buft_layer[i].buft]++; @@ -865,8 +868,10 @@ 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 && + 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 ? 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(); ggml_tensor * k = nullptr; ggml_tensor * v = nullptr; ggml_tensor * s = nullptr; @@ -893,7 +898,9 @@ static bool llama_kv_cache_init( n_mla++; } else { - if (!hparams.has_kv(i)) { + // Qwen 3 MTP layer always need KV cache in the MTP context + const bool is_mtp_layer = (cparams.mtp_op_type != MTP_OP_NONE && i >= (int)n_mtp_first_layer); + if (!hparams.has_kv(i) && !is_mtp_layer) { cache.k_l.push_back(nullptr); cache.v_l.push_back(nullptr); continue; @@ -3304,6 +3311,8 @@ static int llama_model_load(const std::string & fname, llama_model & model, llam model.hparams.vocab_only = params.vocab_only; + model.mtp = params.mtp; + try { llm_load_arch(ml, model); } catch(const std::exception & e) { @@ -4131,7 +4140,15 @@ static size_t llama_output_reserve(llama_context & lctx, size_t n_outputs) { // set all ids as invalid (negative) std::fill(lctx.output_ids.begin(), lctx.output_ids.end(), -1); - ggml_backend_buffer_clear(lctx.buf_output, 0); + if (has_mtp) { + // MTP uses a large output footprint, clear only the active region. + const size_t clear_size = (logits_size + embd_size) * sizeof(float); + if (clear_size > 0 && output_base) { + memset(output_base, 0, clear_size); + } + } else { + ggml_backend_buffer_clear(lctx.buf_output, 0); + } lctx.n_outputs = 0; @@ -4485,8 +4502,14 @@ 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; 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) { + // Qwen 3.5 uses raw hidden state before the final shared-head normalization. + embd = gf->nodes[i]; + break; + } if (strcmp(gf->nodes[i]->name, "result_embd_pooled") == 0) { embd = gf->nodes[i]; break; @@ -5963,7 +5986,7 @@ struct llama_context * llama_init_from_model( } } - if (model->arch != LLM_ARCH_GLM4_MOE && cparams.mtp != 0) { + if (model->arch != LLM_ARCH_GLM4_MOE && model->arch != LLM_ARCH_QWEN35 && cparams.mtp != 0) { cparams.mtp = 0; }