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

Apply helper functions to the MTP graph (#2238)

Browse Source 
* refactor: streamline MTP state handling and input building functions

* move reshape logic outside mtp helper function
This commit is contained in:
Samuel Oliveira Alves
2026-08-03 19:27:19 +03:00
committed by GitHub
parent 276e4ea1eb
commit 60389410a1
8 changed files with 66 additions and 122 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/graphs/build_deepseek2.cpp
+2 -18
View File
@@ -1285,17 +1285,7 @@ ggml_cgraph * llm_build_context::build_deepseek2() {
GGML_ABORT("MTP tail is only wired for GLM_DSA models with NextN layers enabled");
}
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 {
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;
ggml_tensor * hidden_states_from_main_model = build_inp_mtp_states(hparams.n_embd);
const int il_mtp = hparams.n_layer - 1;
const auto & mtp_layer = model.layers[il_mtp];
@@ -1452,17 +1442,11 @@ struct ggml_tensor * llm_build_context::build_deepseek2_mtp(
}
ggml_tensor * token_emb = build_inp_embd_mtp(mtp_embd_weights);
// Normalize and project
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);
if (mtp_layer.nextn.eh_proj == nullptr) {
GGML_ABORT("GLM_DSA MTP requires nextn.eh_proj");
}
ggml_tensor * combined = ggml_concat(ctx0, token_emb_norm, hidden_state_norm, 0);
cb(combined, "mtp_concat", il);
ggml_tensor * cur = llm_build_lora_mm(lctx, ctx0, mtp_layer.nextn.eh_proj, combined);
ggml_tensor * cur = build_mtp_input(mtp_layer, prev_embeddings, token_emb, il, nullptr);
struct ggml_tensor * inpSA = cur;
src/graphs/build_deepseek4.cpp
+11 -24
View File
@@ -1270,35 +1270,22 @@ ggml_cgraph * llm_build_context::build_deepseek4() {
GGML_ASSERT(n_layer > hparams.nextn_predict_layers);
const int64_t n_hidden = n_embd * hc;
ggml_tensor * hidden_state = nullptr;
if (lctx.cparams.mtp_op_type == MTP_OP_WARMUP || lctx.cparams.mtp_op_type == MTP_OP_UPDATE_ACCEPTED) {
hidden_state = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_hidden, n_tokens);
} else {
hidden_state = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_hidden);
}
ggml_set_name(hidden_state, "inp_mtp_states");
ggml_set_input(hidden_state);
lctx.inp_mtp_states = hidden_state;
ggml_tensor * hidden_state = build_inp_mtp_states(n_hidden);
ggml_tensor * tok_embd = build_inp_embd_mtp(model.tok_embd);
const int il_mtp = n_layer - hparams.nextn_predict_layers;
const auto & mtp_layer = model.layers[il_mtp];
ggml_tensor * h_state = ggml_reshape_3d(ctx0, hidden_state, n_embd, hc, n_tokens);
cb(h_state, "mtp_h_state", il_mtp);
ggml_tensor * h_norm = llm_build_norm(ctx0, h_state, hparams, mtp_layer.nextn.hnorm,
nullptr, LLM_NORM_RMS, cb, il_mtp);
cb(h_norm, "mtp_hnorm", il_mtp);
ggml_tensor * e_norm = llm_build_norm(ctx0, tok_embd, hparams, mtp_layer.nextn.enorm,
nullptr, LLM_NORM_RMS, cb, il_mtp);
e_norm = ggml_reshape_3d(ctx0, e_norm, n_embd, 1, n_tokens);
e_norm = ggml_repeat_4d(ctx0, e_norm, n_embd, hc, n_tokens, 1);
cb(e_norm, "mtp_enorm", il_mtp);
ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, 0);
cb(concat, "mtp_concat", il_mtp);
inpL = llm_build_lora_mm(lctx, ctx0, mtp_layer.nextn.eh_proj, concat);
hidden_state = ggml_reshape_2d(ctx0, hidden_state, n_embd, hc * n_tokens);
tok_embd = ggml_reshape_3d(ctx0, tok_embd, n_embd, 1, n_tokens);
tok_embd = ggml_repeat_4d(ctx0, tok_embd, n_embd, hc, n_tokens, 1);
tok_embd = ggml_reshape_2d(ctx0, tok_embd, n_embd, hc * n_tokens);
inpL = build_mtp_input(mtp_layer, hidden_state, tok_embd, il_mtp);
GGML_ASSERT(inpL->ne[0] == n_embd);
GGML_ASSERT(inpL->ne[1] == hc * n_tokens);
GGML_ASSERT(inpL->ne[2] == 1);
GGML_ASSERT(inpL->ne[3] == 1);
inpL = ggml_reshape_3d(ctx0, inpL, n_embd, hc, n_tokens);
cb(inpL, "mtp_eh_proj", il_mtp);
} else {
ggml_tensor * inp = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb);
src/graphs/build_gemma4.cpp
+1 -4
View File
@@ -541,10 +541,7 @@ ggml_cgraph * llm_build_context::build_gemma4_mtp() {
GGML_ASSERT(n_backbone > 0);
ggml_tensor * hidden_state = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_backbone, n_tokens);
ggml_set_name(hidden_state, "inp_mtp_states");
ggml_set_input(hidden_state);
lctx.inp_mtp_states = hidden_state;
ggml_tensor * hidden_state = build_inp_mtp_states(n_backbone);
if (!has_target_ctx || !batch.token) {
ggml_tensor * cur = ggml_view_2d(ctx0, hidden_state, n_embd, n_tokens,
src/graphs/build_glm4.cpp
+2 -17
View File
@@ -18,17 +18,7 @@ ggml_cgraph * llm_build_context::build_glm4_moe() {
ext_factor, attn_factor, beta_fast, beta_slow) : nullptr;
if (cparams.mtp_op_type != MTP_OP_NONE) {
ggml_tensor* hidden_states_from_main_model;
if (cparams.mtp_op_type == MTP_OP_WARMUP || cparams.mtp_op_type == MTP_OP_UPDATE_ACCEPTED) {
hidden_states_from_main_model = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
} else {
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;
ggml_tensor * hidden_states_from_main_model = build_inp_mtp_states(hparams.n_embd);
const int il_mtp = hparams.n_layer - 1;
const auto & mtp_layer = model.layers[il_mtp];
@@ -307,12 +297,7 @@ struct ggml_tensor * llm_build_context::build_glm4_moe_mtp(
}
ggml_tensor * token_emb = build_inp_embd_mtp(mtp_embd_weights);
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 * combined = ggml_concat(ctx0, token_emb_norm, hidden_state_norm, 0);
cb(combined, "mtp_concat", il);
ggml_tensor* cur = llm_build_lora_mm(lctx, ctx0, mtp_layer.nextn.eh_proj, combined);
ggml_tensor * cur = build_mtp_input(mtp_layer, prev_embeddings, token_emb, il, nullptr);
// Self-Attention
const float kq_scale = 1.0f / sqrtf(float(n_embd_head));
src/graphs/build_openpangu.cpp
+2 -15
View File
@@ -922,12 +922,7 @@ ggml_tensor * llm_build_context::build_openpangu_mtp(
ggml_tensor * token_emb = ggml_get_rows(ctx0, mtp_embd_weights, inp_tokens);
cb(token_emb, "inp_embd", il);
ggml_tensor * emb_norm = llm_build_norm(ctx0, token_emb, hparams, mtp_layer.nextn.enorm, NULL, LLM_NORM_RMS, cb, il);
ggml_tensor * hid_norm = llm_build_norm(ctx0, prev_embeddings, hparams, mtp_layer.nextn.hnorm, NULL, LLM_NORM_RMS, cb, il);
// reference order: cat([inputs_embeds, previous_hidden_states], -1)
ggml_tensor * combined = ggml_concat(ctx0, emb_norm, hid_norm, 0);
ggml_tensor * cur = llm_build_lora_mm(lctx, ctx0, mtp_layer.nextn.eh_proj, combined);
ggml_tensor * cur = build_mtp_input(mtp_layer, prev_embeddings, token_emb, il, nullptr);
cb(cur, "mtp_eh_proj", il);
// --- attention sublayer (plain residual) ---
@@ -1022,15 +1017,7 @@ ggml_cgraph * llm_build_context::build_openpangu() {
"OpenPangu MTP graph requested without NextN layers loaded");
GGML_ASSERT(batch.token && "openPangu MTP graphs decode token batches");
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 {
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;
ggml_tensor * hidden_states_from_main_model = build_inp_mtp_states(hparams.n_embd);
// shared batch inputs, created exactly once per graph (see build_openpangu_mtp)
ggml_tensor * inp_pos = build_inp_pos();
src/graphs/build_qwen35.cpp
+4 -44
View File
@@ -15,15 +15,7 @@ ggml_cgraph * llm_build_context::build_qwen35moe() {
ggml_tensor * cur = nullptr;
if (cparams.mtp_op_type != MTP_OP_NONE) {
ggml_tensor * hidden_states_from_main_model;
if (cparams.mtp_op_type == MTP_OP_WARMUP || cparams.mtp_op_type == MTP_OP_UPDATE_ACCEPTED) {
hidden_states_from_main_model = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
} else {
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;
ggml_tensor * hidden_states_from_main_model = build_inp_mtp_states(hparams.n_embd);
const int il_mtp = hparams.n_layer - 1;
const auto & mtp_layer = model.layers[il_mtp];
@@ -99,15 +91,7 @@ ggml_cgraph * llm_build_context::build_qwen35() {
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 {
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;
ggml_tensor * hidden_states_from_main_model = build_inp_mtp_states(hparams.n_embd);
const int il_mtp = hparams.n_layer - 1;
const auto & mtp_layer = model.layers[il_mtp];
@@ -182,18 +166,7 @@ struct ggml_tensor * llm_build_context::build_qwen35moe_mtp(
ggml_tensor * token_emb = build_inp_embd_mtp(model.tok_embd);
ggml_tensor * token_emb_norm = llm_build_norm(ctx0, token_emb, hparams, mtp_layer.nextn.enorm, NULL, LLM_NORM_RMS, cb, il);
ggml_tensor * hidden_state_norm = llm_build_norm(ctx0, prev_embeddings, hparams, mtp_layer.nextn.hnorm, NULL, LLM_NORM_RMS, cb, il);
ggml_tensor * cur;
if (mtp_layer.nextn.eh_proj != nullptr) {
ggml_tensor * combined = ggml_concat(ctx0, token_emb_norm, hidden_state_norm, 0);
cb(combined, "mtp_concat", il);
cur = llm_build_lora_mm(lctx, ctx0, mtp_layer.nextn.eh_proj, combined);
} else {
cur = ggml_add(ctx0, token_emb_norm, hidden_state_norm);
}
cb(cur, "mtp_fused", il);
ggml_tensor * cur = build_mtp_input(mtp_layer, prev_embeddings, token_emb, il, "mtp_fused");
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]) {
@@ -259,20 +232,7 @@ struct ggml_tensor * llm_build_context::build_qwen35_mtp(
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);
ggml_tensor * cur = build_mtp_input(mtp_layer, prev_embeddings, token_emb, il, "mtp_fused");
// 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());
src/llama-build-context.cpp
+35
View File
@@ -467,6 +467,41 @@ struct ggml_tensor * llm_build_context::build_inp_embd_mtp(struct ggml_tensor *
return cur;
}
struct ggml_tensor * llm_build_context::build_inp_mtp_states(int64_t n_hidden) {
struct ggml_tensor * hidden_state = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_hidden, n_tokens);
cb(hidden_state, "inp_mtp_states", -1);
ggml_set_input(hidden_state);
lctx.inp_mtp_states = hidden_state;
return hidden_state;
}
struct ggml_tensor * llm_build_context::build_mtp_input(
const struct llama_layer & mtp_layer,
struct ggml_tensor * hidden_state,
struct ggml_tensor * token_embd,
int il,
const char * output_name) {
GGML_ASSERT(hidden_state->ne[0] == n_embd);
GGML_ASSERT(ggml_are_same_shape(hidden_state, token_embd));
struct ggml_tensor * hidden_state_norm = llm_build_norm(ctx0, hidden_state, hparams,
mtp_layer.nextn.hnorm, nullptr, LLM_NORM_RMS, cb, il);
struct ggml_tensor * token_emb_norm = llm_build_norm(ctx0, token_embd, hparams,
mtp_layer.nextn.enorm, nullptr, LLM_NORM_RMS, cb, il);
struct ggml_tensor * result;
if (mtp_layer.nextn.eh_proj != nullptr) {
struct ggml_tensor * combined = ggml_concat(ctx0, token_emb_norm, hidden_state_norm, 0);
cb(combined, "mtp_concat", il);
result = llm_build_lora_mm(lctx, ctx0, mtp_layer.nextn.eh_proj, combined);
} else {
result = ggml_add(ctx0, token_emb_norm, hidden_state_norm);
}
if (output_name != nullptr) {
cb(result, output_name, il);
}
return result;
}
ggml_tensor * llm_build_context::build_inp_pos() {
int n_pos_per_embd = hparams.rope_type == LLAMA_ROPE_TYPE_MROPE || hparams.rope_type == LLAMA_ROPE_TYPE_IMROPE ? 4 : 1;
lctx.inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, int64_t(n_tokens)*n_pos_per_embd);
src/llama-build-context.h
+9
View File
@@ -129,6 +129,15 @@ struct llm_build_context {
struct ggml_tensor * build_inp_embd_mtp(struct ggml_tensor * mtp_tok_embd);
struct ggml_tensor * build_inp_mtp_states(int64_t n_hidden);
struct ggml_tensor * build_mtp_input(
const struct llama_layer & mtp_layer,
struct ggml_tensor * hidden_state,
struct ggml_tensor * token_embd,
int il,
const char * output_name = "mtp_eh_proj");
ggml_tensor * build_inp_pos();
ggml_tensor * build_input_scale(int n_tokens);