graph: skip the LM head for pooled embeddings (#30445)

* graph: skip the LM head for pooled embeddings

A context that extracts pooled embeddings outputs one vector per
sequence, but every model still expanded its LM head, so the graph
computed n_vocab logits per token and the context kept a host buffer
for them.

The models now only declare t_logits, and set_outputs() expands it
unless the context extracts embeddings with a pooling type other than
NONE. Since the toggle changes the graph shape for a pooled context,
set_embeddings() schedules a re-reserve in that case.

* graph: drop the redundant t_embd expand

Address review from @CISC: with embeddings on, build_pooling() already
expands its output for every pooling type, and a decoder graph without
embeddings does not read t_embd, so set_outputs() only expands
t_logits.

* graph: add build_logits() for the LM head output

Address review from @ggerganov: the LM head is expanded by a dedicated
build_logits(), called from llama_model::build_graph() like
build_pooling(), and set_outputs() only marks the outputs again.

* graph: call build_logits() before build_sampling()

Address review from @ggerganov: build_sampling() reads t_logits, so the
LM head is built first. A pooled context then skips the sampling layers
too, since it has no logits to sample from.

* nit

include/llama.h

@@ -405,7 +405,7 @@ extern "C" {
         void *              abort_callback_data;
 
         // Keep the booleans together and at the end of the struct to avoid misalignment during copy-by-value.
-        bool embeddings;  // if true, extract embeddings (together with logits)
+        bool embeddings;  // if true, extract embeddings (together with logits, except for a pooling type other than NONE)
         bool offload_kqv; // offload the KQV ops (including the KV cache) to GPU
         bool no_perf;     // measure performance timings
         bool op_offload;  // offload host tensor operations to device

src/llama-context.cpp

@@ -1323,10 +1323,12 @@ void llama_context::set_abort_callback(bool (*abort_callback)(void * data), void
 void llama_context::set_embeddings(bool value) {
     LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value);
 
-    cparams.embeddings = value;
+    // a pooled context computes the logits only without embeddings, so the toggle changes the graph shape
+    if (cparams.embeddings != value && cparams.pooling_type != LLAMA_POOLING_TYPE_NONE) {
+        sched_need_reserve = true;
+    }
 
-    // TODO: not sure yet if we want to reserve here
-    //sched_need_reserve = true;
+    cparams.embeddings = value;
 }
 
 void llama_context::set_embeddings_nextn(bool value, bool masked) {
@@ -2233,7 +2235,6 @@ uint32_t llama_context::output_reserve(int32_t n_outputs) {
     const auto n_embd     = hparams.n_embd;
     const auto n_embd_out = hparams.n_embd_out();
 
-    // LFM2 builds logits only when embeddings are off, and set_embeddings() does not re-reserve
     bool has_logits     = !cparams.embeddings || graph_has_logits;
     bool has_embd       = cparams.embeddings;
     bool has_embd_nextn = cparams.embeddings_nextn;

src/llama-graph.cpp

@@ -4063,6 +4063,18 @@ void llm_graph_context::build_pooling(
     ggml_build_forward_expand(gf, cur);
 }
 
+void llm_graph_context::build_logits() const {
+    // a pooled embeddings context outputs one vector per sequence, so its graph skips the LM head
+    if (cparams.embeddings && pooling_type != LLAMA_POOLING_TYPE_NONE) {
+        res->t_logits = nullptr;
+    }
+
+    // the models declare t_logits without expanding it, so the LM head is computed only when it is expanded here
+    if (res->t_logits != nullptr) {
+        ggml_build_forward_expand(gf, res->t_logits);
+    }
+}
+
 void llm_graph_context::build_sampling() const {
     if (samplers.empty() || !res->t_logits) {
         return;

src/llama-graph.h

@@ -1429,6 +1429,12 @@ struct llm_graph_context {
             ggml_tensor * cls_out_b,
             ggml_tensor * cls_norm) const;
 
+    //
+    // logits
+    //
+
+    void build_logits() const;
+
     //
     // sampling (backend sampling)
     //

src/llama-model.cpp

@@ -3305,6 +3305,9 @@ ggml_cgraph * llama_model::build_graph(const llm_graph_params & params) const {
     // add on pooling layer
     llm->build_pooling(cls, cls_b, cls_out, cls_out_b, cls_norm);
 
+    // add the LM head (if the context outputs logits)
+    llm->build_logits();
+
     // add backend sampling layers (if any)
     llm->build_sampling();
 

src/models/afmoe.cpp

@@ -278,6 +278,4 @@ llama_model_afmoe::graph::graph(const llama_model & model, const llm_graph_param
     cur = build_lora_mm(model.output, cur, model.output_s);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/apertus.cpp

@@ -165,6 +165,4 @@ llama_model_apertus::graph::graph(const llama_model & model, const llm_graph_par
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/arcee.cpp

@@ -152,6 +152,4 @@ llama_model_arcee::graph::graph(const llama_model & model, const llm_graph_param
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/arctic.cpp

@@ -175,6 +175,4 @@ llama_model_arctic::graph::graph(const llama_model & model, const llm_graph_para
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/arwkv7.cpp

@@ -197,6 +197,4 @@ llama_model_arwkv7::graph::graph(const llama_model & model, const llm_graph_para
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/baichuan.cpp

@@ -150,6 +150,4 @@ llama_model_baichuan::graph::graph(const llama_model & model, const llm_graph_pa
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/bailingmoe.cpp

@@ -175,6 +175,4 @@ llama_model_bailingmoe::graph::graph(const llama_model & model, const llm_graph_
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/bailingmoe2.cpp

@@ -206,6 +206,4 @@ llama_model_bailingmoe2::graph::graph(const llama_model & model, const llm_graph
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/bailingmoe3.cpp

@@ -412,7 +412,6 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph
     cur = ggml_mul_mat(ctx0, model.output, cur);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-    ggml_build_forward_expand(gf, cur);
 }
 
 llama_model_bailingmoe3::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) :
@@ -555,5 +554,4 @@ llama_model_bailingmoe3::graph_mtp::graph_mtp(const llama_model & model, const l
     cur = ggml_mul_mat(ctx0, model.output, cur);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-    ggml_build_forward_expand(gf, cur);
 }

src/models/bitnet.cpp

@@ -165,6 +165,4 @@ llama_model_bitnet::graph::graph(const llama_model & model, const llm_graph_para
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/bloom.cpp

@@ -146,6 +146,4 @@ llama_model_bloom::graph::graph(const llama_model & model, const llm_graph_param
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/chameleon.cpp

@@ -199,6 +199,4 @@ llama_model_chameleon::graph::graph(const llama_model & model, const llm_graph_p
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/chatglm.cpp

@@ -156,6 +156,4 @@ llama_model_chatglm::graph::graph(const llama_model & model, const llm_graph_par
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/codeshell.cpp

@@ -148,6 +148,4 @@ llama_model_codeshell::graph::graph(const llama_model & model, const llm_graph_p
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/cogvlm.cpp

@@ -155,5 +155,4 @@ llama_model_cogvlm::graph::graph(const llama_model & model, const llm_graph_para
     cur = build_lora_mm(model.output, cur, model.output_s);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-    ggml_build_forward_expand(gf, cur);
 }

src/models/cohere2.cpp

@@ -154,6 +154,4 @@ llama_model_cohere2::graph::graph(const llama_model & model, const llm_graph_par
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/cohere2moe.cpp

@@ -275,8 +275,6 @@ llama_model_cohere2moe::graph::graph(const llama_model & model, const llm_graph_
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 llama_model_cohere2moe::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
@@ -426,6 +424,4 @@ llama_model_cohere2moe::graph_mtp::graph_mtp(const llama_model & model, const ll
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/command-r.cpp

@@ -139,6 +139,4 @@ llama_model_command_r::graph::graph(const llama_model & model, const llm_graph_p
     }
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/dbrx.cpp

@@ -149,6 +149,4 @@ llama_model_dbrx::graph::graph(const llama_model & model, const llm_graph_params
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/deci.cpp

@@ -186,6 +186,4 @@ llama_model_deci::graph::graph(const llama_model & model, const llm_graph_params
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/deepseek.cpp

@@ -189,6 +189,4 @@ llama_model_deepseek::graph::graph(const llama_model & model, const llm_graph_pa
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/deepseek2.cpp

@@ -406,7 +406,6 @@ llama_model_deepseek2::graph_mtp::graph_mtp(const llama_model & model, const llm
     cb(cur, "result_output", -1);
 
     res->t_logits = cur;
-    ggml_build_forward_expand(gf, cur);
 }
 
 llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_params & params) :
@@ -703,6 +702,4 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/deepseek32.cpp

@@ -475,8 +475,6 @@ llama_model_deepseek32::graph::graph(const llama_model & model, const llm_graph_
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 // LLM_GRAPH_TYPE_DECODER_MTP draft head for DeepSeek V3.2 (DEEPSEEK32).
@@ -716,6 +714,5 @@ llama_model_deepseek32::graph_mtp::graph_mtp(const llama_model & model, const ll
     cb(cur, "result_output", -1);
 
     res->t_logits = cur;
-    ggml_build_forward_expand(gf, cur);
 }
 

src/models/deepseek4.cpp

@@ -1353,8 +1353,6 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p
     cur = ggml_mul_mat(ctx0, model.output, cur);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 
@@ -1502,5 +1500,4 @@ llama_model_deepseek4::graph_mtp::graph_mtp(const llama_model & model, const llm
     cb(cur, "result_output", -1);
 
     res->t_logits = cur;
-    ggml_build_forward_expand(gf, cur);
 }

src/models/dflash.cpp

@@ -841,8 +841,6 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
     cb(cur, "result_output", -1);
     res->t_logits = cur;
 
-    ggml_build_forward_expand(gf, cur);
-
     // DSpark: bias the draft logits with the Markov head
     if (model.dspark_markov_w1) {
         build_dspark_markov_head(*this, model, inp_tokens);
@@ -1022,8 +1020,6 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
     cb(cur, "result_output", -1);
     res->t_logits = cur;
 
-    ggml_build_forward_expand(gf, cur);
-
     if (model.dspark_markov_w1) {
         build_dspark_markov_head(*this, model, inp_tokens);
     }

src/models/dots1.cpp

@@ -188,6 +188,4 @@ llama_model_dots1::graph::graph(const llama_model & model, const llm_graph_param
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/dots3note.cpp

@@ -471,6 +471,4 @@ llama_model_dots3note::graph::graph(const llama_model & model, const llm_graph_p
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/dream.cpp

@@ -133,6 +133,4 @@ llama_model_dream::graph::graph(const llama_model & model, const llm_graph_param
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/eagle3.cpp

@@ -321,8 +321,6 @@ llama_model_eagle3::graph<false>::graph(const llama_model & model, const llm_gra
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 std::unique_ptr<llm_graph_context> llama_model_eagle3::build_arch_graph(const llm_graph_params & params) const {

src/models/ernie4-5-moe.cpp

@@ -128,6 +128,4 @@ llama_model_ernie4_5_moe::graph::graph(const llama_model & model, const llm_grap
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/ernie4-5.cpp

@@ -159,6 +159,4 @@ llama_model_ernie4_5::graph::graph(const llama_model & model, const llm_graph_pa
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/exaone-moe.cpp

@@ -234,6 +234,4 @@ llama_model_exaone_moe::graph::graph(const llama_model & model, const llm_graph_
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/exaone.cpp

@@ -131,6 +131,4 @@ llama_model_exaone::graph::graph(const llama_model & model, const llm_graph_para
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/exaone4.cpp

@@ -181,8 +181,6 @@ llama_model_exaone4::graph<iswa>::graph(const llama_model & model, const llm_gra
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 // Explicit template instantiations

src/models/falcon-h1.cpp

@@ -204,6 +204,4 @@ llama_model_falcon_h1::graph::graph(const llama_model & model, const llm_graph_p
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/falcon.cpp

@@ -156,6 +156,4 @@ llama_model_falcon::graph::graph(const llama_model & model, const llm_graph_para
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/gemma.cpp

@@ -134,6 +134,4 @@ llama_model_gemma::graph::graph(const llama_model & model, const llm_graph_param
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/gemma2.cpp

@@ -170,6 +170,4 @@ llama_model_gemma2::graph::graph(const llama_model & model, const llm_graph_para
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/gemma3.cpp

@@ -213,8 +213,6 @@ llama_model_gemma3::graph<iswa>::graph(const llama_model & model, const llm_grap
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 template struct llama_model_gemma3::graph<false>;

src/models/gemma3n.cpp

@@ -309,8 +309,6 @@ llama_model_gemma3n::graph::graph(const llama_model & model, const llm_graph_par
     }
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 ggml_tensor * llama_model_gemma3n::graph::calc_magnitude(ggml_tensor * x) {

src/models/gemma4-assistant.cpp

@@ -201,6 +201,5 @@ llama_model_gemma4_assistant::graph::graph(const llama_model & model, const llm_
     cb(h_next, "h_nextn", -1);
     res->t_h_nextn = h_next;
 
-    ggml_build_forward_expand(gf, logits);
     ggml_build_forward_expand(gf, h_next);
 }

src/models/gemma4.cpp

@@ -439,8 +439,6 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 class llm_graph_input_gemma4_ple : public llm_graph_input_i {

src/models/glm-dsa.cpp

@@ -515,8 +515,6 @@ llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_par
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 // LLM_GRAPH_TYPE_DECODER_MTP draft head for GLM-5.2 (GLM_DSA).
@@ -756,5 +754,4 @@ llama_model_glm_dsa::graph_mtp::graph_mtp(const llama_model & model, const llm_g
     cb(cur, "result_output", -1);
 
     res->t_logits = cur;
-    ggml_build_forward_expand(gf, cur);
 }

src/models/glm4-moe.cpp

@@ -280,7 +280,6 @@ llama_model_glm4_moe::graph_mtp::graph_mtp(const llama_model & model, const llm_
     cb(cur, "result_output", -1);
 
     res->t_logits = cur;
-    ggml_build_forward_expand(gf, cur);
 }
 
 llama_model_glm4_moe::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
@@ -434,6 +433,4 @@ llama_model_glm4_moe::graph::graph(const llama_model & model, const llm_graph_pa
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/glm4.cpp

@@ -181,6 +181,4 @@ llama_model_glm4::graph::graph(const llama_model & model, const llm_graph_params
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/glm5-next.cpp

@@ -667,7 +667,6 @@ llama_model_glm5_next::graph_mtp::graph_mtp(const llama_model & model, const llm
     cb(cur, "result_output", -1);
 
     res->t_logits = cur;
-    ggml_build_forward_expand(gf, cur);
 }
 
 llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_params & params) :
@@ -800,8 +799,6 @@ llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_p
     cur = ggml_mul_mat(ctx0, model.output, cur);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 // KDA layer, g_a/g_b output gate

src/models/gpt2.cpp

@@ -143,6 +143,4 @@ llama_model_gpt2::graph::graph(const llama_model & model, const llm_graph_params
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/gptneox.cpp

@@ -214,6 +214,4 @@ llama_model_gptneox::graph::graph(const llama_model & model, const llm_graph_par
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/granite-hybrid.cpp

@@ -195,8 +195,6 @@ llama_model_granite_hybrid::graph::graph(const llama_model & model, const llm_gr
     }
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 ggml_tensor * llama_model_granite_hybrid::graph::build_attention_layer(ggml_tensor *             cur,

src/models/granite-swa.cpp

@@ -193,8 +193,6 @@ llama_model_granite_swa::graph::graph(
     cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_logit_scale);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 ggml_tensor * llama_model_granite_swa::graph::build_attention_layer(

src/models/granite-switch.cpp

@@ -326,8 +326,6 @@ llama_model_granite_switch::graph::graph(
     cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_logit_scale);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 ggml_tensor * llama_model_granite_switch::graph::build_attention_layer(

src/models/granite.cpp

@@ -198,8 +198,6 @@ llama_model_granite::graph::graph(
     cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_logit_scale);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 ggml_tensor * llama_model_granite::graph::build_attention_layer(

src/models/grok.cpp

@@ -218,6 +218,4 @@ llama_model_grok::graph::graph(const llama_model & model, const llm_graph_params
     }
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/grovemoe.cpp

@@ -188,6 +188,4 @@ llama_model_grovemoe::graph::graph(const llama_model & model, const llm_graph_pa
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/hrm-text.cpp

@@ -208,6 +208,4 @@ llama_model_hrm_text::graph::graph(const llama_model & model, const llm_graph_pa
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/hunyuan-moe.cpp

@@ -182,6 +182,4 @@ llama_model_hunyuan_moe::graph::graph(const llama_model & model, const llm_graph
     cur = build_lora_mm(model.output, cur, model.output_s);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/hunyuan-vl.cpp

@@ -186,6 +186,4 @@ llama_model_hunyuan_vl::graph::graph(const llama_model & model, const llm_graph_
     cur = build_lora_mm(model.output, cur, model.output_s);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/hy-v3.cpp

@@ -217,8 +217,6 @@ llama_model_hy_v3::graph::graph(const llama_model & model, const llm_graph_param
     cur = build_lora_mm(model.output, cur, model.output_s);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 // LLM_GRAPH_TYPE_DECODER_MTP draft head for HY V3 (MoE).
@@ -380,5 +378,4 @@ llama_model_hy_v3::graph_mtp::graph_mtp(const llama_model & model, const llm_gra
     cb(cur, "result_output", -1);
 
     res->t_logits = cur;
-    ggml_build_forward_expand(gf, cur);
 }

src/models/hy-v4.cpp

@@ -596,6 +596,4 @@ llama_model_hy_v4::graph::graph(const llama_model & model, const llm_graph_param
     cur = ggml_mul_mat(ctx0, model.output, cur);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/internlm2.cpp

@@ -134,6 +134,4 @@ llama_model_internlm2::graph::graph(const llama_model & model, const llm_graph_p
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/jais.cpp

@@ -127,6 +127,4 @@ llama_model_jais::graph::graph(const llama_model & model, const llm_graph_params
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/jais2.cpp

@@ -150,6 +150,4 @@ llama_model_jais2::graph::graph(const llama_model & model, const llm_graph_param
     cb(cur, "result_output", -1);
 
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/jamba.cpp

@@ -193,6 +193,4 @@ llama_model_jamba::graph::graph(const llama_model & model, const llm_graph_param
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/k2-horizon.cpp

@@ -357,6 +357,4 @@ llama_model_k2_horizon::graph::graph(const llama_model & model, const llm_graph_
     cur = build_lora_mm(model.output, cur, model.output_s);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/kimi-k3.cpp

@@ -346,8 +346,6 @@ llama_model_kimi_k3::graph::graph(const llama_model & model, const llm_graph_par
     cur = ggml_mul_mat(ctx0, model.output, cur);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 //

src/models/kimi-linear.cpp

@@ -557,6 +557,4 @@ llama_model_kimi_linear::graph::graph(const llama_model & model, const llm_graph
     cur = ggml_mul_mat(ctx0, model.output, cur);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/laguna.cpp

@@ -326,6 +326,4 @@ llama_model_laguna::graph::graph(const llama_model & model, const llm_graph_para
     cur = build_lora_mm(model.output, cur);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/lfm2.cpp

@@ -341,8 +341,6 @@ llama_model_lfm2::graph<iswa>::graph(const llama_model & model, const llm_graph_
 
         res->t_logits = cur;
     }
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 // media entries (an image or audio prefix) are embeddings, text entries are tokens

src/models/llada-moe.cpp

@@ -158,6 +158,4 @@ llama_model_llada_moe::graph::graph(const llama_model & model, const llm_graph_p
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/llada.cpp

@@ -148,6 +148,4 @@ llama_model_llada::graph::graph(const llama_model & model, const llm_graph_param
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/llama.cpp

@@ -242,8 +242,6 @@ llama_model_llama::graph<embed>::graph(const llama_model & model, const llm_grap
         cb(cur, "result_output", -1);
         res->t_logits = cur;
     }
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 template struct llama_model_llama::graph<false>;

src/models/llama4.cpp

@@ -263,8 +263,6 @@ llama_model_llama4::graph<iswa>::graph(const llama_model & model, const llm_grap
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 // Explicit template instantiations

src/models/maincoder.cpp

@@ -146,6 +146,4 @@ llama_model_maincoder::graph::graph(const llama_model & model, const llm_graph_p
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/mamba.cpp

@@ -132,6 +132,4 @@ llama_model_mamba::graph::graph(const llama_model & model, const llm_graph_param
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/maple.cpp

@@ -145,6 +145,4 @@ llama_model_maple::graph::graph(const llama_model & model, const llm_graph_param
     cur = build_lora_mm(model.output, cur, model.output_s);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/mellum.cpp

@@ -211,8 +211,6 @@ llama_model_mellum::graph<iswa>::graph(const llama_model & model, const llm_grap
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 template struct llama_model_mellum::graph<false>;

src/models/mimo2.cpp

@@ -253,8 +253,6 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 // Mirrors MiMo's appended NextN block: normalize and fuse token and hidden inputs, run the decoder block,
@@ -397,5 +395,4 @@ llama_model_mimo2::graph_mtp::graph_mtp(const llama_model & model, const llm_gra
     cb(cur, "result_output", -1);
 
     res->t_logits = cur;
-    ggml_build_forward_expand(gf, cur);
 }

src/models/minicpm3.cpp

@@ -247,6 +247,4 @@ llama_model_minicpm3::graph::graph(const llama_model & model, const llm_graph_pa
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/minimax-01.cpp

@@ -479,6 +479,4 @@ llama_model_minimax_01::graph::graph(const llama_model & model, const llm_graph_
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/minimax-m2.cpp

@@ -163,6 +163,4 @@ llama_model_minimax_m2::graph::graph(const llama_model & model, const llm_graph_
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/minimax-m3.cpp

@@ -603,6 +603,4 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_
     cur = build_lora_mm(model.output, cur, model.output_s);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/mistral3.cpp

@@ -230,6 +230,4 @@ llama_model_mistral3::graph::graph(const llama_model & model, const llm_graph_pa
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/mpt.cpp

@@ -165,6 +165,4 @@ llama_model_mpt::graph::graph(const llama_model & model, const llm_graph_params
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/muse-glimmer.cpp

@@ -194,8 +194,6 @@ llama_model_muse_glimmer::graph::graph(const llama_model & model, const llm_grap
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 std::unique_ptr<llm_graph_context> llama_model_muse_glimmer::build_arch_graph(const llm_graph_params & params) const {

src/models/nanbeige.cpp

@@ -180,6 +180,4 @@ llama_model_nanbeige::graph::graph(const llama_model & model, const llm_graph_pa
     cur = build_lora_mm(model.output, cur, model.output_s);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/nemotron-h-moe.cpp

@@ -156,5 +156,4 @@ llama_model_nemotron_h_moe::graph_mtp::graph_mtp(const llama_model & model, cons
     cb(cur, "result_output", -1);
 
     res->t_logits = cur;
-    ggml_build_forward_expand(gf, cur);
 }

src/models/nemotron-h.cpp

@@ -259,8 +259,6 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_
     cur = build_lora_mm(model.output, cur, model.output_s);
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 ggml_tensor * llama_model_nemotron_h::graph::build_attention_layer(ggml_tensor *             cur,

src/models/nemotron.cpp

@@ -145,6 +145,4 @@ llama_model_nemotron::graph::graph(const llama_model & model, const llm_graph_pa
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/olmo.cpp

@@ -137,6 +137,4 @@ llama_model_olmo::graph::graph(const llama_model & model, const llm_graph_params
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/olmo2.cpp

@@ -199,8 +199,6 @@ llama_model_olmo2::graph<iswa>::graph(const llama_model & model, const llm_graph
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }
 
 // Explicit template instantiations

src/models/olmoe.cpp

@@ -168,6 +168,4 @@ llama_model_olmoe::graph::graph(const llama_model & model, const llm_graph_param
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/openai-moe.cpp

@@ -170,6 +170,4 @@ llama_model_openai_moe::graph::graph(const llama_model & model, const llm_graph_
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/openelm.cpp

@@ -166,6 +166,4 @@ llama_model_openelm::graph::graph(const llama_model & model, const llm_graph_par
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/orion.cpp

@@ -136,6 +136,4 @@ llama_model_orion::graph::graph(const llama_model & model, const llm_graph_param
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }

src/models/paddleocr.cpp

@@ -102,6 +102,4 @@ llama_model_paddleocr::graph::graph(const llama_model & model, const llm_graph_p
 
     cb(cur, "result_output", -1);
     res->t_logits = cur;
-
-    ggml_build_forward_expand(gf, cur);
 }