diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 9f044de03..06865c981 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -789,6 +789,8 @@ void llama_context::sched_reserve() { n_nodes_pp = ggml_graph_n_nodes(gf); n_inputs_pp = get_gf_res_reserve()->inputs.size(); n_input_tensors_pp = this->n_input_tensors; + + graph_has_logits = get_gf_res_reserve()->get_logits() != nullptr; } // reserve with tg (token generation) graph to get the number of splits and nodes @@ -2231,7 +2233,8 @@ 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(); - bool has_logits = true; + // 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; diff --git a/src/llama-context.h b/src/llama-context.h index b89a519f7..6938fb103 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -368,6 +368,10 @@ private: bool sched_need_reserve = true; + // whether the reserve graph has a logits tensor, set in sched_reserve() + // true until then, because the first output_reserve() runs before any graph is reserved + bool graph_has_logits = true; + // state of sched_copy_experts, reset before each graph compute struct copy_experts_info { const ggml_tensor * ids = nullptr;