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; }