gemma4: enable pruning based on static cparams.kv_unified/n_seq_max=1

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.

common/speculative.cpp

@@ -1574,12 +1574,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);
 
@@ -1733,16 +1733,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;
+            }
 
-            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);
+            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 < 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;

examples/speculative-simple/speculative-simple.cpp

@@ -5,6 +5,8 @@
 #include "log.h"
 #include "llama.h"
 
+#include "../../src/llama-ext.h"
+
 #include <algorithm>
 #include <clocale>
 #include <cstdio>
@@ -64,6 +66,7 @@ int main(int argc, char ** argv) {
     }
 
     llama_context * ctx_dft = params.speculative.draft.ctx_dft;
+    const bool is_mtp_mem_shared = ctx_dft && llama_get_ctx_other(ctx_dft) == ctx_tgt;
 
     // check if the context supports partial sequence removal
     const bool use_ckpt_tgt = common_context_can_seq_rm(ctx_tgt) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL;
@@ -130,6 +133,9 @@ int main(int argc, char ** argv) {
             batch_prompt.add(inp[i], i, seq_id, false);
         }
 
+        if (is_mtp_mem_shared && batch_prompt.size() > 0) {
+            batch_prompt.set_output(batch_prompt.size() - 1, true);
+        }
         llama_process(ctx_tgt, LLAMA_PROCESS_TYPE_DECODE, batch_prompt.get());
 
         if (!common_speculative_process(spec, batch_prompt)) {
@@ -228,7 +234,6 @@ int main(int argc, char ** argv) {
                 batch_tgt.add(draft[i], n_past + i, seq_id, true);
             }
 
-
             llama_process(ctx_tgt, LLAMA_PROCESS_TYPE_DECODE, batch_tgt.get());
         }
 

src/models/gemma4.cpp

@@ -148,8 +148,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) :
@@ -180,7 +179,8 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
     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]);
@@ -189,7 +189,27 @@ 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) && n_outputs < n_tokens && cparams.causal_attn &&
+                        hparams.n_layer_kv_from_start > 0 && hparams.n_layer_kv_from_start < n_layer &&
+                        !cparams.embeddings && (cparams.kv_unified || 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];
+    }
+    ggml_tensor * pos_cur         = inp_pos;
+    ggml_tensor * kq_mask_out     = nullptr;
+    ggml_tensor * kq_mask_swa_out = nullptr;
+    int64_t       n_rows          = n_tokens;
+
     for (int il = 0; il < n_layer; ++il) {
+        if (prune_shared && n_rows == n_tokens && !hparams.has_kv(il)) {
+            inpL            = ggml_get_rows(ctx0, inpL, inp_out_ids);
+            pos_cur         = ggml_reshape_1d(ctx0, ggml_get_rows(ctx0, ggml_reshape_2d(ctx0, inp_pos, 1, n_tokens), inp_out_ids), n_outputs);
+            kq_mask_out     = ggml_cast(ctx0, ggml_get_rows(ctx0, inp_attn->get_kq_mask(),     inp_out_ids), inp_attn->get_kq_mask()->type);
+            kq_mask_swa_out = ggml_cast(ctx0, ggml_get_rows(ctx0, inp_attn->get_kq_mask_swa(), inp_out_ids), inp_attn->get_kq_mask_swa()->type);
+            n_rows          = n_outputs;
+        }
+
         const int64_t n_embd_head = hparams.n_embd_head_k(il);
         GGML_ASSERT(n_embd_head == hparams.n_embd_head_v(il));
 
@@ -220,19 +240,19 @@ 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);
 
-            Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, freq_factors, n_rot_l, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
+            Qcur = ggml_rope_ext(ctx0, Qcur, pos_cur, freq_factors, n_rot_l, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,
                                  ext_factor, attn_factor, beta_fast, beta_slow);
             cb(Qcur, "Qcur_pos", il);
         }
@@ -276,14 +296,21 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
                     hparams.f_attention_scale, il);
         } else {
             // reuse KV cache of earlier layers
+            ggml_tensor * mask_full     = inp_attn->self_kq_mask_cnv;
+            ggml_tensor * mask_swa_full = inp_attn->self_kq_mask_swa_cnv;
+            if (kq_mask_out) {
+                inp_attn->self_kq_mask_cnv     = kq_mask_out;
+                inp_attn->self_kq_mask_swa_cnv = kq_mask_swa_out;
+            }
             cur = build_attn(inp_attn,
                     model.layers[il].wo, nullptr, model.layers[il].wo_s,
                     Qcur, nullptr, nullptr, nullptr, nullptr, nullptr, hparams.f_attention_scale, il);
+            inp_attn->self_kq_mask_cnv     = mask_full;
+            inp_attn->self_kq_mask_swa_cnv = mask_swa_full;
         }
 
-        // 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 (il == n_layer - 1 && crop_before_nextn(inp_out_ids) && n_rows == n_tokens) {
             cur  = ggml_get_rows(ctx0,  cur, inp_out_ids);
             inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
         }
@@ -325,7 +352,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,
@@ -377,18 +404,19 @@ 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 (inp_out_ids && n_rows != n_tokens) {
+                inp_this_layer = ggml_get_rows(ctx0, inp_this_layer, inp_out_ids);
+            } else if (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);
 
@@ -421,7 +449,7 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
     cb(cur, "h_nextn", -1);
     res->t_h_nextn = cur;
 
-    if (crop_after_nextn(inp_out_ids)) {
+    if (crop_after_nextn(inp_out_ids) && n_rows == n_tokens) {
         cur = ggml_get_rows(ctx0, cur, inp_out_ids);
     }
 
@@ -516,19 +544,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);
 }