Fix Gemma4 MTP

Author: ikawrakow

src/graphs/build_gemma4.cpp

@@ -663,7 +663,7 @@ ggml_cgraph * llm_build_context::build_gemma4_mtp() {
     // not required for correct inference — the full-vocab matmul against the tied output
     // weight still yields valid per-token logits.
     {
-        logits = build_output(lctx, ctx0, cur, model.output, model.output_norm, cb);
+        logits = build_output(lctx, ctx0, cur, model.output, model.output_norm, cb, false);
         cb(logits, "result_output", -1);
     }
     ggml_build_forward_expand(gf, logits);

src/llama-build-context.cpp

@@ -2048,7 +2048,7 @@ ggml_tensor * llm_build_context::build_output(llama_context & lctx, ggml_context
 }
 
 ggml_tensor * llm_build_context::build_output(llama_context & lctx, ggml_context * ctx, ggml_tensor * cur,
-        ggml_tensor * output, ggml_tensor * output_norm, const llm_build_cb & cb) {
+        ggml_tensor * output, ggml_tensor * output_norm, const llm_build_cb & cb, bool add_normed_name) {
     // lm_head
     if (output->extra) {
         auto split_output = (ggml_split_tensor_t *)output->extra;
@@ -2069,7 +2069,7 @@ ggml_tensor * llm_build_context::build_output(llama_context & lctx, ggml_context
                 o.push_back(llm_build_context::llm_build_lora_mm(lctx, ctx, split, cur));
             }
             cb(o.back(), "output", id);
-            if (last_norm) {
+            if (add_normed_name && last_norm) {
                 cb(last_norm, "result_norm", -1);
             }
         }
@@ -2095,7 +2095,9 @@ ggml_tensor * llm_build_context::build_output(llama_context & lctx, ggml_context
         }
         if (output_norm) {
             cur = llm_build_context::llm_build_norm(ctx, cur, lctx.model.hparams, output_norm, NULL, LLM_NORM_RMS, cb, -1);
-            cb(cur, "result_norm", -1);
+            if (add_normed_name) {
+                cb(cur, "result_norm", -1);
+            }
         }
         cur = llm_build_context::llm_build_lora_mm(lctx, ctx, output, cur);
     }

src/llama-build-context.h

@@ -448,7 +448,7 @@ llm_expert_gating_func_type   gating_op,
     static ggml_tensor * build_output(llama_context & lctx, ggml_context * ctx, ggml_tensor * cur, ggml_tensor * output, const llm_build_cb & cb);
 
     static ggml_tensor * build_output(llama_context & lctx, ggml_context * ctx, ggml_tensor * cur,
-            ggml_tensor * output, ggml_tensor * output_norm, const llm_build_cb & cb);
+            ggml_tensor * output, ggml_tensor * output_norm, const llm_build_cb & cb, bool add_normed_name = true);
 
     static ggml_tensor * do_split_norm(ggml_context * ctx, ggml_tensor * cur, ggml_tensor * the_norm, const llama_hparams & hparams,
         const llm_build_cb & cb, int id, int il_cb, bool is_norm);