Skip to content

gemma4: graph optimizations (massive gains on pp for E2B and E4B) #31356

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ggerganov merged 13 commits into from
Nov 3, 2026
+71 −26
Merged

gemma4: graph optimizations (massive gains on pp for E2B and E4B) #31356

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
22756d3
Optimizing Gemma4 graph for prompt processing
Meet91721 Sep 20, 2026
e1c1ca2
Not filtering the rows when they are unmasked
Meet91721 Oct 28, 2026
7da7a2c
Cropping out the attention mask too.
Meet91721 Oct 29, 2026
262cbca
gemma-4: improved pruning, mask cropping and PLE scaling
max-krasnyansky Oct 29, 2026
a9b3dc6
spec-mtp: Enable masked h_nextn extraction only when the draft contex…
max-krasnyansky Oct 30, 2026
8e7d454
spec-mtp: make sure to not drop first draft round
max-krasnyansky Oct 30, 2026
3341212
gemma4: enable pruning based on static cparams.kv_unified/n_seq_max=1
max-krasnyansky Oct 30, 2026
f27b68f
gemma4-mtp: simplify prunning logic and avoid graph branches/re-allocs
max-krasnyansky Oct 30, 2026
d842bfb
gemma4: update pruning gates to use only n_seq_max==1, kv_unified can…
max-krasnyansky Nov 3, 2026
d14ee0b
spec-simple: revert spec-simple changes for now
max-krasnyansky Nov 3, 2026
0455252
gemma4: remove causal_attn from pruning gate (last source of reallocs)
max-krasnyansky Nov 3, 2026
0971a7a
cont : adapt to llm_attn_mask changes
ggerganov Nov 3, 2026
12b8b54
cont : add TODO
ggerganov Nov 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 20 additions & 8 deletions common/speculative.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down Expand Up @@ -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;
Expand Down
69 changes: 51 additions & 18 deletions src/models/gemma4.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) :
Expand Down Expand Up @@ -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]);
Expand All @@ -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;

@ggerganov ggerganov Nov 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@max-krasnyansky You mentioned that the approach works with unified KV cache, but here you actually have a stronger restriction: cparams.n_seq_max == 1. So for example, the default server config uses 4 sequences with unified KV cache and it does not benefit from the pruning because of this check.

@max-krasnyansky max-krasnyansky Nov 3, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ggerganov
The earlier check was n_seq_max == 1 || unified_kv == true but it stopped working after #29510 got merged.
It probably worked earlier just because unified_kv was not used witn n_seq_max > 1, or perhaps I misread the code.

As things are right now it only works with the Single-seq configurations.
Are you OK merging with this restriction or should we try to make it work with Multi-seq first?
It seems pretty tricky :) and I'm not sure there is a good use-case for Multi-seq with these smaller models. It'd be very useful for 26B but all layers feed the KV$ in 26B so we can't prune anything.

@ggerganov ggerganov Nov 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's not very difficult to implement - I added a TODO in the code. If you prefer, you can merge it like this and resolve it later.

@max-krasnyansky max-krasnyansky Nov 3, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok. I'd say let's merge to enable the default case, and I'll do another pass to resolve the TODO.
Approvals expired. You'd have to merge or approve :)

@Meet91721 Meet91721 Nov 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @max-krasnyansky ,
If you are not planning to implement that change in this PR, I would like to take it up. I was trying to understand a few changes that you made on top of my commits.

@max-krasnyansky max-krasnyansky Nov 3, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @max-krasnyansky , If you are not planning to implement that change in this PR, I would like to take it up. I was trying to understand a few changes that you made on top of my commits.

Go for it!
I think we're in a very good shape for the default use cases: single-seq with/without MTP.
Multi-seq support is what's missing. If you have cycles please give it a go.
Please tag me in the next PR just so that it shows up in my feed (thank you thank you).

for (int il = (int) hparams.n_layer_kv_from_start; prune_shared && il <= n_layer; ++il) {
prune_shared = !cparams.embeddings_layer_inp[il];
}
Comment on lines +201 to +207

@ggerganov ggerganov Oct 29, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When the MTP is enabled, there will be no benefit because we still need to extract the full embeddings.

Also, this will create branches in the compute graph when using split KV cache. Can we avoid the inp_attn->get_kq_mask()->ne[3] == 1; condition somehow?

@max-krasnyansky max-krasnyansky Oct 29, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah. That's one of the use-cases I missed.
Will play with that and see if we can avoid branches.


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));

Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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);
}