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
This commit is contained in:
Pascal authored and GitHub committed 2026-10-13 16:16:14 +02:00
1 parent a17e94d229
commit cb55ffe6e2
136 files changed
+27 -282

No files matched your search

-2
View File
@@ -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);
}