From 96550613656e7f024df65f91cf8b2d80a83cf09e Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Mon, 21 Sep 2026 19:13:04 +0300 Subject: [PATCH] llama-context : report graph inputs and input tensors during sched reserve (#26625) * llama-context : report graph inputs and input tensors during sched reserve - fix the tg (token generation) graph bs label to use n_seqs instead of a hardcoded 1 - report the number of graph inputs from llm_graph_result::inputs for both the pp and tg graphs - report the number of input tensors (nodes and their src tensors flagged with GGML_TENSOR_FLAG_INPUT) - log a warning when an input tensor has an op other than GGML_OP_NONE - log a trace line for each input tensor and the nodes (name and op) that use it Assisted-by: llama.cpp:DeepSeek-v4-Flash-0731 * cont : count input tensors before reserving the sched * wip * llama-graph : name the unnamed graph input tensors - name the kv-cache idxs input tensors (attn_inp_k_idxs, attn_inp_v_idxs) - name the recurrent state copy idxs input tensor (rs_s_copy) - report the input tensor shape in the sched_reserve trace Assisted-by: pi:llama.cpp/Qwen3.8-27B * llama-context : rename "graph inputs" to "graph input objects" Assisted-by: pi:llama.cpp/Qwen3.8-27B * llama-context : report the sched reserve graph stats on a single line - print nodes, splits, input objects and input tensors in one line - when the pp and tg graphs differ, print each value as 'pp / tg' and annotate the line with the batch sizes used for each graph Assisted-by: pi:llama.cpp/Qwen3.8-27B * cont : pad logs --- src/llama-context.cpp | 83 +++++++++++++++++++++++++++++++++--------- src/llama-context.h | 1 + src/llama-graph.cpp | 11 +++++- src/llama-kv-cache.cpp | 2 + 4 files changed, 79 insertions(+), 18 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index ef53728d1d..fcd4dfb13b 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -19,6 +19,7 @@ #include #include #include +#include // // llama_context @@ -579,6 +580,40 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3 } } +static int llama_graph_n_input_tensors(ggml_cgraph * gf) { + std::unordered_map> users; + for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { + ggml_tensor * node = ggml_graph_node(gf, i); + if (node->flags & GGML_TENSOR_FLAG_INPUT) { + users[node].push_back(node); + } + for (int j = 0; j < GGML_MAX_SRC; ++j) { + ggml_tensor * src = node->src[j]; + if (!src) { + break; + } + if (src->flags & GGML_TENSOR_FLAG_INPUT) { + users[src].push_back(node); + } + } + } + + for (const auto & [tensor, nodes] : users) { + if (tensor->op != GGML_OP_NONE) { + LLAMA_LOG_WARN("%s: input tensor '%32s' has op %s, expected GGML_OP_NONE\n", + __func__, tensor->name, ggml_op_name(tensor->op)); + } + for (const ggml_tensor * node : nodes) { + LLAMA_LOG_DEBUG("%s: input tensor '%32s' [%s, ne = { %5" PRId64 ", %5" PRId64 ", %5" PRId64 ", %5" PRId64 " }] is used by node '%s' (%s)\n", + __func__, tensor->name, ggml_type_name(tensor->type), + tensor->ne[0], tensor->ne[1], tensor->ne[2], tensor->ne[3], + node->name, ggml_op_name(node->op)); + } + } + + return (int) users.size(); +} + void llama_context::sched_reserve() { if (!sched_need_reserve) { return; @@ -624,11 +659,15 @@ void llama_context::sched_reserve() { resolve_fused_ops(mctx.get(), n_seqs); // reserve worst-case graph - int n_splits_pp = -1; - int n_nodes_pp = -1; + int n_splits_pp = -1; + int n_nodes_pp = -1; + int n_inputs_pp = -1; + int n_input_tensors_pp = -1; - int n_splits_tg = -1; - int n_nodes_tg = -1; + int n_splits_tg = -1; + int n_nodes_tg = -1; + int n_inputs_tg = -1; + int n_input_tensors_tg = -1; const uint32_t n_outputs_pp = std::min(n_tokens, cparams.n_outputs_max); @@ -648,8 +687,10 @@ 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_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 = this->n_input_tensors; } // reserve with tg (token generation) graph to get the number of splits and nodes @@ -659,8 +700,10 @@ void llama_context::sched_reserve() { throw std::runtime_error("failed to allocate compute tg buffers"); } - n_splits_tg = ggml_backend_sched_get_n_splits(sched.get()); - n_nodes_tg = ggml_graph_n_nodes(gf); + 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 = this->n_input_tensors; } // reserve again with pp graph to avoid ggml-alloc reallocations during inference @@ -698,16 +741,21 @@ void llama_context::sched_reserve() { } } - if (n_nodes_pp == n_nodes_tg) { - LLAMA_LOG_INFO("%s: graph nodes = %d\n", __func__, n_nodes_pp); - } else { - LLAMA_LOG_INFO("%s: graph nodes = %d (with bs=%d), %d (with bs=1)\n", __func__, n_nodes_pp, n_tokens, n_nodes_tg); - } + { + const bool diff = n_nodes_pp != n_nodes_tg || n_splits_pp != n_splits_tg || + n_inputs_pp != n_inputs_tg || n_input_tensors_pp != n_input_tensors_tg; - if (n_splits_pp == n_splits_tg) { - LLAMA_LOG_INFO("%s: graph splits = %d\n", __func__, n_splits_pp); - } else { - LLAMA_LOG_INFO("%s: graph splits = %d (with bs=%d), %d (with bs=1)\n", __func__, n_splits_pp, n_tokens, n_splits_tg); + const auto val = [diff](int v_pp, int v_tg) -> std::string { + return diff ? format("%d / %d", v_pp, v_tg) : format("%d", v_pp); + }; + + LLAMA_LOG_INFO("%s: graph%s: nodes = %s, splits = %s, input objects = %s, input tensors = %s\n", + __func__, + diff ? format(" (pp bs=%d, tg bs=%d)", n_tokens, n_seqs).c_str() : "", + val(n_nodes_pp, n_nodes_tg).c_str(), + val(n_splits_pp, n_splits_tg).c_str(), + val(n_inputs_pp, n_inputs_tg).c_str(), + val(n_input_tensors_pp, n_input_tensors_tg).c_str()); } const int64_t t_end_us = ggml_time_us(); @@ -2475,6 +2523,7 @@ ggml_cgraph * llama_context::graph_reserve( auto * gf = model.build_graph(gparams); + this->n_input_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 b7a9db5913..77ef92fc6a 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -333,6 +333,7 @@ private: // reuse the batch_allocr to avoid unnecessary memory allocations std::unique_ptr balloc; + uint32_t n_input_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 diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index fd4290cf05..02ae8bd92a 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2451,6 +2451,7 @@ ggml_tensor * llm_graph_context::build_inp_pos() const { cur = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, (int64_t)n_tokens*hparams.n_pos_per_embd()); ggml_set_input(cur); + cb(cur, "inp_pos", -1); res->add_input(std::move(inp)); @@ -2465,7 +2466,7 @@ ggml_tensor * llm_graph_context::build_inp_attn_scale() const { // this need to be 1x1xN for broadcasting cur = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, 1, n_tokens); ggml_set_input(cur); - ggml_set_name(cur, "attn_scale"); + cb(cur, "inp_attn_scale", -1); res->add_input(std::move(inp)); @@ -2487,6 +2488,7 @@ ggml_tensor * llm_graph_context::build_inp_out_ids() const { cur = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_outputs); ggml_set_input(cur); + ggml_set_name(cur, "out_ids"); res->add_input(std::move(inp)); @@ -2500,6 +2502,7 @@ ggml_tensor * llm_graph_context::build_inp_mean() const { cur = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_tokens, ubatch.n_seqs_unq); ggml_set_input(cur); + ggml_set_name(cur, "mean"); res->add_input(std::move(inp)); @@ -2513,6 +2516,7 @@ ggml_tensor * llm_graph_context::build_inp_cls() const { cur = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ubatch.n_seqs_unq); ggml_set_input(cur); + ggml_set_name(cur, "cls"); res->add_input(std::move(inp)); @@ -2537,6 +2541,7 @@ ggml_tensor * llm_graph_context::build_inp_cross_embd() const { cur = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd, n_enc); ggml_set_input(cur); + ggml_set_name(cur, "cross_embd"); res->add_input(std::move(inp)); @@ -2550,6 +2555,7 @@ ggml_tensor * llm_graph_context::build_inp_pos_bucket_enc() const { cur = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_tokens, n_tokens); ggml_set_input(cur); + ggml_set_name(cur, "pos_bucket_enc"); res->add_input(std::move(inp)); @@ -2567,6 +2573,7 @@ ggml_tensor * llm_graph_context::build_inp_pos_bucket_dec() const { cur = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_tokens); ggml_set_input(cur); + ggml_set_name(cur, "pos_bucket_dec"); res->add_input(std::move(inp)); @@ -2735,6 +2742,7 @@ llm_graph_input_attn_no_cache * llm_graph_context::build_attn_inp_no_cache() con // note: there is no KV cache, so the number of KV values is equal to the number of tokens in the batch inp->self_kq_mask = ggml_new_tensor_4d(ctx0, type_mask, n_tokens, n_tokens, 1, 1); ggml_set_input(inp->self_kq_mask); + cb(inp->self_kq_mask, "self_kq_mask", -1); inp->self_kq_mask_cnv = inp->self_kq_mask; @@ -3511,6 +3519,7 @@ static std::unique_ptr build_rs_inp_impl( inp->s_copy = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_rs); ggml_set_input(inp->s_copy); + ggml_set_name(inp->s_copy, "rs_s_copy"); inp->s_copy_main = ggml_view_1d(ctx0, inp->s_copy, n_seqs, 0); inp->s_copy_extra = ggml_view_1d(ctx0, inp->s_copy, n_rs - n_seqs, n_seqs * inp->s_copy->nb[0]); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index a342ee1191..332d1abe02 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1412,6 +1412,7 @@ ggml_tensor * llama_kv_cache::build_input_k_idxs(ggml_context * ctx, const llama ggml_tensor * k_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_tokens); ggml_set_input(k_idxs); + ggml_set_name(k_idxs, "attn_inp_k_idxs"); return k_idxs; } @@ -1428,6 +1429,7 @@ ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama } ggml_set_input(v_idxs); + ggml_set_name(v_idxs, "attn_inp_v_idxs"); return v_idxs; }