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 <samueloliveira32df@gmail.com>
Date: Mon Apr 20 13:30:21 2026 -0300
recurrent model: refactor api
commit dfac8f19f6edc0014b4116041b89c1e0dfb173c7
Author: SamuelOliveirads <samueloliveira32df@gmail.com>
Date: Mon Apr 20 12:22:29 2026 -0300
recurrent model: implement recurrent kernel checkpoint
commit 9c44b117f93e9060030907e1250106358c9ccf47
Author: SamuelOliveirads <samueloliveira32df@gmail.com>
Date: Sat Apr 18 11:52:39 2026 -0300
speculative: fix sampler for checkpoints
commit e7006393bca20adcd86e2d77021e0b41d7bd9db1
Author: SamuelOliveirads <samueloliveira32df@gmail.com>
Date: Fri Apr 17 14:08:25 2026 -0300
server: refactor checkpoint state logic
commit 57eabf04df5185cab19539185b8f4b85e578905b
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 dc4797b72363482bb35750bb5edc87068116dc0f
Author: SamuelOliveirads <samueloliveira32df@gmail.com>
Date: Fri Apr 17 13:12:40 2026 -0300
reset ngram mod state for rejected tokens
commit 8ff2d943a31b3d54440698db41e62ea121d661be
Author: SamuelOliveirads <samueloliveira32df@gmail.com>
Date: Fri Apr 17 13:08:04 2026 -0300
server: snapshot recurrent state in tensor
commit d93dfb5e6b78a822e7331b33feedbdc47eb5ec79
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 d670cf85cd59f23a339faf6fd773889063176801
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
GitHub
parent
d6f3e4e28f
commit
67e6346225
|
|
@ -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));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
|
@ -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];
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
);
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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" },
|
||||
},
|
||||
},
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue