Public/koboldcpp
1
0
Fork 0
mirror of https://github.com/LostRuins/koboldcpp.git synced 2026-09-11 15:29:12 +02:00
Code Issues Packages Projects Releases Wiki Activity

llama : add Gemma4 MTP (#23398)

Browse Source
This commit is contained in:
Aman Gupta
2026-06-07 20:50:54 +08:00
committed by GitHub
parent 8a091c47ab
commit 04eb4c446d
31 changed files with 644 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/models/gemma4.cpp
+18 -4
View File
@@ -155,12 +155,14 @@ public:
}
virtual ~llm_graph_input_logits_bias() = default;
void set_input(const llama_ubatch *) override {
void set_input(const llama_ubatch * /*ubatch*/) override {
const int64_t n_vocab = arr.size();
ggml_backend_tensor_set(logits_bias, arr.data(), 0, n_vocab*ggml_element_size(logits_bias));
}
// bool can_reuse(const llm_graph_params & params) override;
bool can_reuse(const llm_graph_params & /*params*/) override {
return true;
}
ggml_tensor * logits_bias = nullptr; // F32 [n_vocab]
@@ -270,7 +272,8 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
}
// TODO @ngxson : strip unused token right after the last KV layer to speed up prompt processing
if (il == n_layer - 1 && inp_out_ids) {
// keep all rows when extracting unmasked nextn embeddings (MTP target needs the hidden state for every token)
if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
}
@@ -370,7 +373,7 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
ggml_tensor * inp_this_layer = ggml_view_2d_slice(ctx0, inp_per_layer, il); // [n_embd_per_layer, n_tokens]
// TODO @ngxson : improve this
if (il == n_layer - 1 && inp_out_ids) {
if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
inp_this_layer = ggml_get_rows(ctx0, inp_this_layer, inp_out_ids);
}
@@ -401,6 +404,17 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
model.output_norm, nullptr,
LLM_NORM_RMS, -1);
// Expose the post-output-norm hidden state (the LM-head input feature) so that
// MTP draft contexts can read it via llama_get_embeddings_nextn_ith() as the
// recurrent h input. This matches the reference (transformers/vLLM/SGLang),
// which feeds the drafter the target's post-final-norm hidden state.
cb(cur, "h_nextn", -1);
res->t_h_nextn = cur;
if (!cparams.embeddings_nextn_masked && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
}
cb(cur, "result_norm", -1);
res->t_embd = cur;