Commit 6cb3e672e for llama.cpp
commit 6cb3e672e8d10fd4feab1a048bf2967c40623133
Author: Max Krasnyansky <maxk@qti.qualcomm.com>
Date: Tue Nov 3 10:30:37 2026 -0800
gemma4: graph optimizations (massive gains on pp for E2B and E4B) (#31356)
* Optimizing Gemma4 graph for prompt processing
TODO:
- Currently complete microbatch needs to contain tokens which does not require the logits; need to optimize this further
- There are a few more operations that are performed after the kv cache is calculated, need to safely optimize them.
* Not filtering the rows when they are unmasked
* Cropping out the attention mask too.
* gemma-4: improved pruning, mask cropping and PLE scaling
Prune before the first KV-shared layer with the single-stream, causal, non-embedding safety gates.
Use cropped masks only during shared-layer attention.
Preserve requested layer-input embeddings in the KV prefix.
Remove the PLE contiguous materialization and fold equivalent PLE scaling into RMSNorm parameters/weights.
* spec-mtp: Enable masked h_nextn extraction only when the draft context shares the target KV cache
This lets Gemma 4 prune unused prompt rows before its shared layers so that we keep the full perf P.P gains with MTP.
Avoid extracting target hidden states for output-free prompt chunks in that shared-MTP path.
When a batch has a selected final token, extract only that row unless the whole verification batch is selected.
Keep ordinary MTP and classic speculative decoding behavior unchanged: they still use unmasked h_nextn because the draft
context must catch up on all target rows.
speculative-simple needs an update to select the final prompt row only for shared-KV MTP,
so the optimization works there without changing other speculative modes.
And it looks like the server already selects its final prompt row for first-token sampling.
* spec-mtp: make sure to not drop first draft round
* gemma4: enable pruning based on static cparams.kv_unified/n_seq_max=1
* gemma4-mtp: simplify prunning logic and avoid graph branches/re-allocs
* gemma4: update pruning gates to use only n_seq_max==1, kv_unified can now be used for multi-seq
* spec-simple: revert spec-simple changes for now
* gemma4: remove causal_attn from pruning gate (last source of reallocs)
* cont : adapt to llm_attn_mask changes
* cont : add TODO
---------
Co-authored-by: Meet Amin <meetamin378@gmail.com>
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
diff --git a/common/speculative.cpp b/common/speculative.cpp
index f82dc879f..e6290eaad 100644
--- a/common/speculative.cpp
+++ b/common/speculative.cpp
@@ -1613,12 +1613,12 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
}
}
- llama_set_embeddings_nextn(ctx_tgt, true, /*masked*/ false);
- llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true);
-
is_mem_shared = llama_get_ctx_other(ctx_dft) == ctx_tgt;
chain_heads = n_mtp_layers > 1 && !is_mem_shared;
+ llama_set_embeddings_nextn(ctx_tgt, true, /*masked*/ is_mem_shared);
+ llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true);
+
if (chain_heads) {
this->params.n_max = std::min(this->params.n_max, n_mtp_layers);
@@ -1772,16 +1772,28 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
}
const int32_t n_rows = i_batch_end[seq_id] - i_batch_beg[seq_id] + 1;
- verify_h_rows[seq_id] = n_rows;
- verify_h[seq_id].resize((size_t) n_rows * n_embd);
+ if (is_mem_shared && !batch_in.tokens[i_batch_end[seq_id]].output) {
+ verify_h_rows[seq_id] = 0;
+ continue;
+ }
+
+ const bool all_rows_output = std::all_of(
+ batch_in.tokens.begin() + i_batch_beg[seq_id],
+ batch_in.tokens.begin() + i_batch_end[seq_id] + 1,
+ [](const common_batch::token & token) { return token.output; });
+ const int32_t h_rows = is_mem_shared && !all_rows_output ? 1 : n_rows;
+ const int32_t h_offset = is_mem_shared && !all_rows_output ? i_batch_end[seq_id] : i_batch_beg[seq_id];
+
+ verify_h_rows[seq_id] = h_rows;
+ verify_h[seq_id].resize((size_t) h_rows * n_embd);
- for (int32_t i = 0; i < n_rows; ++i) {
- const float * h = llama_get_embeddings_nextn_ith(ctx_tgt, i_batch_beg[seq_id] + i);
+ for (int32_t i = 0; i < h_rows; ++i) {
+ const float * h = llama_get_embeddings_nextn_ith(ctx_tgt, h_offset + i);
std::memcpy(verify_h[seq_id].data() + (size_t) i * n_embd, h, row_bytes);
}
std::memcpy(pending_h[seq_id].data(),
- verify_h[seq_id].data() + (size_t) (n_rows - 1) * n_embd, row_bytes);
+ verify_h[seq_id].data() + (size_t) (h_rows - 1) * n_embd, row_bytes);
}
return true;
diff --git a/src/models/gemma4.cpp b/src/models/gemma4.cpp
index d0b2b3a5b..df245f811 100644
--- a/src/models/gemma4.cpp
+++ b/src/models/gemma4.cpp
@@ -154,8 +154,7 @@ std::unique_ptr<llm_graph_context> llama_model_gemma4::build_arch_graph(const ll
// get 2D slice view from a 3D tensor, the idx corresponds to the 3rd dim
static ggml_tensor * gemma4_view_2d_slice(ggml_context * ctx0, ggml_tensor * x, int idx) {
GGML_ASSERT(idx < (int) x->ne[2]);
- return ggml_view_2d(ctx0, x, x->ne[0], x->ne[1], ggml_row_size(x->type, x->ne[0]),
- idx * x->ne[0] * x->ne[1] * ggml_element_size(x));
+ return ggml_view_2d(ctx0, x, x->ne[0], x->ne[1], x->nb[1], idx * x->nb[2]);
}
llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_params & params) :
@@ -183,10 +182,14 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
// TODO: is causal == true correct? might need some changes
auto * inp_attn = build_attn_inp_kv_iswa();
+ auto inp_attn_full = *inp_attn;
+ auto inp_attn_prune = *inp_attn;
+
ggml_tensor * inp_out_ids = build_inp_out_ids();
if (model.per_layer_tok_embd) {
- const float tok_embd_scale = sqrtf((float) n_embd_per_layer);
+ // fold per_layer_input_scale (1/sqrt(2)) from project_per_layer_inputs() into tok_embd_scale
+ const float tok_embd_scale = sqrtf((float) n_embd_per_layer) / sqrtf(2.0f);
inp_per_layer = ggml_scale (ctx0, inp_per_layer, tok_embd_scale);
inp_per_layer = ggml_reshape_3d(ctx0, inp_per_layer, n_embd_per_layer, n_layer, inp_per_layer->ne[1]);
@@ -195,7 +198,36 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
inp_per_layer = project_per_layer_inputs(inpL, inp_per_layer);
}
+ // KV-shared layers do not write to the KV cache, so prune non-output token rows before the first KV-shared layer
+ bool prune_shared = crop_before_nextn(inp_out_ids) &&
+ hparams.n_layer_kv_from_start > 0 && hparams.n_layer_kv_from_start < n_layer &&
+ !cparams.embeddings && cparams.n_seq_max == 1;
+ for (int il = (int) hparams.n_layer_kv_from_start; prune_shared && il <= n_layer; ++il) {
+ prune_shared = !cparams.embeddings_layer_inp[il];
+ }
+
+ int64_t n_rows = n_tokens;
+
for (int il = 0; il < n_layer; ++il) {
+ if (prune_shared && il == (int) hparams.n_layer_kv_from_start) {
+ inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
+ inp_pos = ggml_reshape_1d(ctx0, ggml_get_rows(ctx0, ggml_reshape_2d(ctx0, inp_pos, 1, n_tokens), inp_out_ids), n_outputs);
+
+ // TODO: instead of doing get_rows here, extend llm_attn_mask to have a `out` mask that
+ // is construted during set_input to contain just the pruned output cells
+ // this can be used to enable pruning with more than 1 sequence, both with unified and split KV cache
+ // ref: https://github.com/ggml-org/llama.cpp/pull/31356#discussion_r4455286628
+ inp_attn_prune.self_attn_mask.raw = ggml_get_rows(ctx0, inp_attn_full.self_attn_mask.raw, inp_out_ids);
+ inp_attn_prune.self_attn_mask_swa.raw = ggml_get_rows(ctx0, inp_attn_full.self_attn_mask_swa.raw, inp_out_ids);
+
+ inp_attn_prune.self_attn_mask.raw = ggml_cast(ctx0, inp_attn_prune.self_attn_mask.raw, GGML_TYPE_F16);
+ inp_attn_prune.self_attn_mask_swa.raw = ggml_cast(ctx0, inp_attn_prune.self_attn_mask_swa.raw, GGML_TYPE_F16);
+
+ n_rows = n_outputs;
+
+ inp_attn = &inp_attn_prune;
+ }
+
const int64_t n_embd_head = hparams.n_embd_head_k(il);
GGML_ASSERT(n_embd_head == hparams.n_embd_head_v(il));
@@ -226,14 +258,14 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
qkv_fused = build_lora_mm(model.layers[il].wqkv, cur, model.layers[il].wqkv_s);
cb(qkv_fused, "wqkv", il);
const int64_t q_dim = n_embd_head * n_head;
- Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv_fused, q_dim, n_tokens, qkv_fused->nb[1], 0));
+ Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv_fused, q_dim, n_rows, qkv_fused->nb[1], 0));
} else {
Qcur = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s);
}
{
cb(Qcur, "Qcur", il);
- Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_tokens);
+ Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head, n_rows);
Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il);
cb(Qcur, "Qcur_normed", il);
@@ -287,9 +319,8 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
Qcur, nullptr, nullptr, nullptr, nullptr, nullptr, hparams.f_attention_scale, il);
}
- // TODO @ngxson : strip unused token right after the last KV layer to speed up prompt processing
// keep all rows when extracting unmasked nextn embeddings (MTP target needs the hidden state for every token)
- if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
+ if (!prune_shared && il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
}
@@ -331,7 +362,7 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
ggml_tensor * tmp = ggml_rms_norm(ctx0, attn_out, hparams.f_norm_rms_eps);
tmp = ggml_scale(ctx0, tmp, 1.0f / sqrtf((float) n_embd));
tmp = ggml_mul(ctx0, tmp, model.layers[il].ffn_gate_inp_s);
- ggml_tensor * logits = build_lora_mm(model.layers[il].ffn_gate_inp, tmp); // [n_expert, n_tokens]
+ ggml_tensor * logits = build_lora_mm(model.layers[il].ffn_gate_inp, tmp); // [n_expert, n_rows]
cb(logits, "ffn_moe_logits", il);
cur_moe = build_moe_ffn(cur_moe,
@@ -383,18 +414,17 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
ggml_tensor * pe_in = cur;
cb(cur, "pe_in", il);
- cur = build_lora_mm(model.layers[il].per_layer_inp_gate, cur); // [n_embd_per_layer, n_tokens]
+ cur = build_lora_mm(model.layers[il].per_layer_inp_gate, cur); // [n_embd_per_layer, n_rows]
cur = ggml_gelu(ctx0, cur);
ggml_tensor * inp_this_layer = gemma4_view_2d_slice(ctx0, inp_per_layer, il); // [n_embd_per_layer, n_tokens]
- // TODO @ngxson : improve this
- if (il == n_layer - 1 && crop_before_nextn(inp_out_ids)) {
+ if ((prune_shared && !hparams.has_kv(il)) || (il == n_layer - 1 && crop_before_nextn(inp_out_ids))) {
inp_this_layer = ggml_get_rows(ctx0, inp_this_layer, inp_out_ids);
}
cur = ggml_mul(ctx0, cur, inp_this_layer);
- cur = build_lora_mm(model.layers[il].per_layer_proj, cur); // [n_embd, n_tokens]
+ cur = build_lora_mm(model.layers[il].per_layer_proj, cur); // [n_embd, n_rows]
cur = build_norm(cur, model.layers[il].per_layer_post_norm, nullptr, LLM_NORM_RMS, il);
cb(cur, "per_layer_embd_out", il);
@@ -522,19 +552,22 @@ ggml_tensor * llama_model_gemma4::graph::project_per_layer_inputs(ggml_tensor *
const float per_layer_input_scale = 1.0f / sqrtf(2.0f);
// note: this matrix multiplication will be performed in the input layer (i.e. on the CPU)
+ // fold per_layer_input_scale into proj_norm_w and per_layer_projection_scale into proj_eps
+ ggml_tensor * proj_norm_w = ggml_scale(ctx0, model.per_layer_proj_norm, per_layer_input_scale);
+ ggml_build_forward_expand(gf, proj_norm_w);
+
ggml_tensor * per_layer_proj;
per_layer_proj = ggml_mul_mat (ctx0, model.per_layer_model_proj, inp_batch);
- per_layer_proj = ggml_scale (ctx0, per_layer_proj, per_layer_projection_scale);
per_layer_proj = ggml_reshape_3d(ctx0, per_layer_proj, n_embd_per_layer, n_layer, n_tokens);
- per_layer_proj = build_norm(per_layer_proj, model.per_layer_proj_norm, nullptr, LLM_NORM_RMS, -1);
+ const float proj_eps = hparams.f_norm_rms_eps / (per_layer_projection_scale * per_layer_projection_scale);
+ per_layer_proj = ggml_rms_norm(ctx0, per_layer_proj, proj_eps);
+ per_layer_proj = ggml_mul(ctx0, per_layer_proj, proj_norm_w);
cb(per_layer_proj, "per_layer_proj", -1);
- inp_per_layer = ggml_add (ctx0, per_layer_proj, inp_per_layer);
- inp_per_layer = ggml_scale(ctx0, inp_per_layer, per_layer_input_scale);
+ inp_per_layer = ggml_add(ctx0, per_layer_proj, inp_per_layer);
cb(inp_per_layer, "inp_per_layer", -1);
// permute to shape: [n_embd_per_layer, n_tokens, n_layer]
- inp_per_layer = ggml_cont(ctx0, ggml_permute(ctx0, inp_per_layer, 0, 2, 1, 3));
- return inp_per_layer;
+ return ggml_permute(ctx0, inp_per_layer, 0, 2, 1, 3);
}