LLM-Inference/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-08-12 22:29:39 +04:00
Code Issues Packages Projects Releases Wiki Activity

Support for Qwen 3.5 MTP (dense models only) (#1698)

Browse Source 
* 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 <samueloliveira32df@gmail.com>
Date:   Mon Apr 20 13:30:21 2026 -0300

    recurrent model: refactor api

commit dfac8f19f6
Author: SamuelOliveirads <samueloliveira32df@gmail.com>
Date:   Mon Apr 20 12:22:29 2026 -0300

    recurrent model: implement recurrent kernel checkpoint

commit 9c44b117f9
Author: SamuelOliveirads <samueloliveira32df@gmail.com>
Date:   Sat Apr 18 11:52:39 2026 -0300

    speculative: fix sampler for checkpoints

commit e7006393bc
Author: SamuelOliveirads <samueloliveira32df@gmail.com>
Date:   Fri Apr 17 14:08:25 2026 -0300

    server: refactor checkpoint state logic

commit 57eabf04df
Merge: dc4797b7 64234e3c
Author: SamuelOliveirads <samueloliveira32df@gmail.com>
Date:   Fri Apr 17 13:53:41 2026 -0300

    Merge branch 'main' into fix/hybrid-cache-speculative

commit dc4797b723
Author: SamuelOliveirads <samueloliveira32df@gmail.com>
Date:   Fri Apr 17 13:12:40 2026 -0300

    reset ngram mod state for rejected tokens

commit 8ff2d943a3
Author: SamuelOliveirads <samueloliveira32df@gmail.com>
Date:   Fri Apr 17 13:08:04 2026 -0300

    server: snapshot recurrent state in tensor

commit d93dfb5e6b
Author: SamuelOliveirads <samueloliveira32df@gmail.com>
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 d670cf85cd
Author: SamuelOliveirads <samueloliveira32df@gmail.com>
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
This commit is contained in:
Samuel Oliveira Alves
2026-04-28 07:47:50 +02:00
committed by GitHub
parent d6f3e4e28f
commit 67e6346225
10 changed files with 401 additions and 131 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/server/server-context.cpp
+69 -31
View File
@@ -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<llama_token> & ids, int n_draft) {
const std::vector<llama_token> & ids, int n_draft,
const std::vector<float> & 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<float> 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));
}
}
}
}
ggml/src/ggml-alloc.c
+7 -1
View File
@@ -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;
}
src/graphs/build_glm4.cpp
+4 -4
View File
@@ -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,
src/graphs/build_qwen35.cpp
+136 -32
View File
@@ -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;
}
src/llama-build-context.cpp
+5 -1
View File
@@ -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];
}
src/llama-build-context.h
+9 -1
View File
@@ -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
);
};
src/llama-hparams.cpp
+24 -4
View File
@@ -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;
src/llama-load-tensors.cpp
+92 -29
View File
@@ -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<uint32_t>(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<uint32_t>(i) >= n_layer - hparams.nextn_predict_layers;
int flags = 0;
if (hparams.nextn_predict_layers > 0 && static_cast<uint32_t>(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<uint32_t>(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<uint32_t>(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<uint32_t>(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<uint32_t>(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<uint32_t>(i) >= n_layer - hparams.nextn_predict_layers;
int flags = 0;
if (hparams.nextn_predict_layers > 0 && static_cast<uint32_t>(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<uint32_t>(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<uint32_t>(il) >= static_cast<uint32_t>(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);
src/llama-model.cpp
+27 -23
View File
@@ -508,29 +508,33 @@ static const std::map<llm_arch, std::map<llm_tensor, std::string>> 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" },
},
},
{
src/llama.cpp
+28 -5
View File
@@ -778,8 +778,11 @@ static bool llama_kv_cache_init(
// count used buffer types
std::map<ggml_backend_buffer_type_t, int> 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;
}