From aa75b9f5914aa177580985ca0dfedba45807e25b Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 5 Aug 2026 16:57:07 +0300 Subject: [PATCH] cont : count input tensors before reserving the sched --- src/llama-context.cpp | 8 +++----- src/llama-context.h | 1 + 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 3fdead4422..02bd3d10b7 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -579,10 +579,7 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3 } } -// count the number of tensors in the graph that are flagged as inputs -// this includes the graph nodes themselves as well as their src tensors static int llama_graph_n_input_tensors(ggml_cgraph * gf) { - // map each input tensor to the graph nodes that use it std::unordered_map> users; for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { ggml_tensor * node = ggml_graph_node(gf, i); @@ -687,7 +684,7 @@ void llama_context::sched_reserve() { n_splits_pp = ggml_backend_sched_get_n_splits(sched.get()); n_nodes_pp = ggml_graph_n_nodes(gf); n_inputs_pp = get_gf_res_reserve()->inputs.size(); - n_input_tensors_pp = llama_graph_n_input_tensors(gf); + n_input_tensors_pp = this->n_intput_tensors; } // reserve with tg (token generation) graph to get the number of splits and nodes @@ -700,7 +697,7 @@ void llama_context::sched_reserve() { n_splits_tg = ggml_backend_sched_get_n_splits(sched.get()); n_nodes_tg = ggml_graph_n_nodes(gf); n_inputs_tg = get_gf_res_reserve()->inputs.size(); - n_input_tensors_tg = llama_graph_n_input_tensors(gf); + n_input_tensors_tg = this->n_intput_tensors; } // reserve again with pp graph to avoid ggml-alloc reallocations during inference @@ -2481,6 +2478,7 @@ ggml_cgraph * llama_context::graph_reserve( auto * gf = model.build_graph(gparams); + this->n_intput_tensors = llama_graph_n_input_tensors(gf); this->n_outputs = save_n_outputs; // initialize scheduler with the specified graph diff --git a/src/llama-context.h b/src/llama-context.h index bf91daa8b5..2e69a9d5db 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -330,6 +330,7 @@ private: // reuse the batch_allocr to avoid unnecessary memory allocations std::unique_ptr balloc; + uint32_t n_intput_tensors = 0; // number of tensors marked as input during the last graph reserve uint32_t n_outputs = 0; // number of actually-used outputs in the current ubatch or last logical batch std::vector output_ids; // map batch token positions to ids of the logits and embd buffers