From 8e2d31e0eb658d0382d3106a61369dff81f0c080 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Fri, 9 Oct 2026 20:43:05 +0300 Subject: [PATCH] graph : reorder get_rows for embeddings (#30160) * graph : reorder get_rows for embeddings * cont : fix gemma4 and improve input embedding construction logic * cont : add TODO for lora * cont : fix raw embeddings path * gemma4 : avoid ple cast in embeddings path --- src/llama-graph.cpp | 83 ++++++++++++++++++++++++++++------------- src/models/gemma4.cpp | 44 +++++++++++++--------- src/models/qwen4exp.cpp | 22 +++++------ 3 files changed, 95 insertions(+), 54 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 5a6a2b3208..1514b8aebe 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2471,19 +2471,57 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd, float to auto inp = std::make_unique(n_embd_inp); - inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ubatch.n_tokens); - cb(inp->tokens, "inp_tokens", -1); - ggml_set_input(inp->tokens); - res->t_inp_tokens = inp->tokens; + // mixed path (ubatch.is_mixed()): set_rows the token rows into a copy of the embd rows, with its own inputs as select branches must not share tensors + // TODO: use inp->tokens and inp->embd once ggml_build_forward_select allows it + const bool has_mixed = llm_arch_supports_mixed_batch(arch) && cparams.ctx_type == LLAMA_CONTEXT_TYPE_DEFAULT; - inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, ubatch.n_tokens); - cb(inp->embd, "inp_embd", -1); - ggml_set_input(inp->embd); + const int64_t n_tok_rows = has_mixed ? llm_graph_n_tok_rows(ubatch) : 0; - // token embeddings with lora and padding - auto build_tok = [&](ggml_tensor * ids) { - ggml_tensor * cur = ggml_get_rows(ctx0, tok_embd, ids); + // we have 3 standard paths to produce the input embeddings for the first layer: + // - embd0: extract from the token embeddings weight (`tok_embd`) using the input token ids + // - embd1: pass raw embeddings, skipping the `tok_embd` + // - embd2: mixed path of both tokens ids + raw embeddings (if supported) + ggml_tensor * embd0 = nullptr; + ggml_tensor * embd1 = nullptr; + ggml_tensor * embd2 = nullptr; + // construct the input tensors + { + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ubatch.n_tokens); + cb(inp->tokens, "inp_tokens", -1); + ggml_set_input(inp->tokens); + res->t_inp_tokens = inp->tokens; + + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, ubatch.n_tokens); + cb(inp->embd, "inp_embd", -1); + ggml_set_input(inp->embd); + + if (has_mixed) { + inp->mixed_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tok_rows); + cb(inp->mixed_tokens, "inp_mixed_tokens", -1); + ggml_set_input(inp->mixed_tokens); + } + } + + // the embeddings placeholders for the 3 paths + // we use ggml_build_forward_order to make the GET_ROWS ops stick at the beginning of the compute graph + // this way the embeddings remain in the host buffer, and the GET_ROWS run before any other computations + { + embd0 = ggml_get_rows(ctx0, tok_embd, inp->tokens); + ggml_build_forward_order(gf, embd0); + + embd1 = inp->embd; + + if (has_mixed) { + embd2 = ggml_get_rows(ctx0, tok_embd, inp->mixed_tokens); + ggml_build_forward_order(gf, embd2); + } + } + + // helper for extracting token embeddings with lora and padding + // TODO: when lora is active, this is likely going to cause issues similar to https://github.com/ggml-org/llama.cpp/pull/30160 + // need to add lora tests and refactor the logic to make the lora GET_ROWS go at the front of the graph + auto build_tok = [&](ggml_tensor * cur, ggml_tensor * ids) { // apply lora for embedding tokens if needed for (const auto & lora : *loras) { llama_adapter_lora_weight * lw = lora.first->get_weight(tok_embd); @@ -2514,21 +2552,15 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd, float to std::array inps = {}; // token embeddings path (ubatch.token != nullptr) - inps[0] = build_tok(inp->tokens); + inps[0] = build_tok(embd0, inp->tokens); // vector embeddings path (ubatch.embd != nullptr) - inps[1] = inp->embd; + inps[1] = embd1; + + assert(ggml_are_same_shape (inps[0], inps[1])); + assert(ggml_are_same_stride(inps[0], inps[1])); - // mixed path (ubatch.is_mixed()): set_rows the token rows into a copy of the embd rows, with its own inputs as select branches must not share tensors - // TODO: use inp->tokens and inp->embd once ggml_build_forward_select allows it - const bool has_mixed = llm_arch_supports_mixed_batch(arch) && cparams.ctx_type == LLAMA_CONTEXT_TYPE_DEFAULT; if (has_mixed) { - const int64_t n_tok_rows = llm_graph_n_tok_rows(ubatch); - - inp->mixed_tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tok_rows); - cb(inp->mixed_tokens, "inp_mixed_tokens", -1); - ggml_set_input(inp->mixed_tokens); - inp->mixed_slots = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_tok_rows); cb(inp->mixed_slots, "inp_mixed_slots", -1); ggml_set_input(inp->mixed_slots); @@ -2538,11 +2570,12 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd, float to ggml_set_input(inp->mixed_embd); // note: set_rows writes into its destination, so it gets a copy of the input - inps[2] = ggml_set_rows(ctx0, ggml_dup(ctx0, inp->mixed_embd), build_tok(inp->mixed_tokens), inp->mixed_slots); - } + ggml_tensor * embd_mixed = build_tok(embd2, inp->mixed_tokens); + inps[2] = ggml_set_rows(ctx0, ggml_dup(ctx0, inp->mixed_embd), embd_mixed, inp->mixed_slots); - assert(ggml_are_same_shape (inps[0], inps[1])); - assert(ggml_are_same_stride(inps[0], inps[1])); + assert(ggml_are_same_shape (inps[0], inps[2])); + assert(ggml_are_same_stride(inps[0], inps[2])); + } const int idx = ubatch.is_mixed() ? 2 : ubatch.token ? 0 : 1; diff --git a/src/models/gemma4.cpp b/src/models/gemma4.cpp index fbc4d2a6e2..d883384375 100644 --- a/src/models/gemma4.cpp +++ b/src/models/gemma4.cpp @@ -159,6 +159,14 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para ggml_tensor * cur; ggml_tensor * inpL; + // do the PLE first to guarantee it is done in the host buffer + // ref: https://github.com/ggml-org/llama.cpp/pull/30160 + ggml_tensor * inp_per_layer = nullptr; + if (model.per_layer_tok_embd) { + inp_per_layer = build_inp_per_layer(); + ggml_build_forward_expand(gf, inp_per_layer); + } + // important: do not normalize weights for raw embeddings input (i.e. encoded image emdeddings) inpL = build_inp_embd(model.tok_embd, sqrtf(n_embd)); cb(inpL, "inp_scaled", -1); @@ -171,10 +179,11 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para ggml_tensor * inp_out_ids = build_inp_out_ids(); - ggml_tensor * inp_per_layer = nullptr; if (model.per_layer_tok_embd) { - inp_per_layer = build_inp_per_layer(); - ggml_build_forward_expand(gf, inp_per_layer); + const float tok_embd_scale = sqrtf((float) n_embd_per_layer); + + inp_per_layer = ggml_scale (ctx0, inp_per_layer, tok_embd_scale); + inp_per_layer = ggml_reshape_3d(ctx0, inp_per_layer, n_embd_per_layer, n_layer, inp_per_layer->ne[1]); // inp_per_layer shape: [n_embd_per_layer, n_tokens, n_layer] inp_per_layer = project_per_layer_inputs(inpL, inp_per_layer); @@ -448,10 +457,13 @@ public: llama_prefetch_rows(ple, ubatch->token, ubatch->n_tokens); } ggml_backend_tensor_set(tokens, ubatch->token, 0, ubatch->n_tokens * ggml_element_size(tokens)); - } else if (prefetch) { - // [TAG_GEMMA4_IMG_PADDING] + } else { const int32_t padding = 0; - llama_prefetch_rows(ple, &padding, 1); + if (prefetch) { + // [TAG_GEMMA4_IMG_PADDING] + llama_prefetch_rows(ple, &padding, 1); + } + ggml_backend_tensor_set(token0, &padding, 0, ggml_element_size(token0)); } } @@ -460,6 +472,7 @@ public: } ggml_tensor * tokens = nullptr; + ggml_tensor * token0 = nullptr; const llama_model & model; }; @@ -470,30 +483,25 @@ ggml_tensor * llama_model_gemma4::graph::build_inp_per_layer() { auto inp = std::make_unique(model); ggml_tensor * inp_per_layer; - float tok_embd_scale = sqrtf((float) n_embd_per_layer); // mixed ubatch: embd rows have token id 0, same padding row as below + // TODO: use ggml_build_forward_select if (ubatch.token) { inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, ubatch.n_tokens); ggml_set_input(inp->tokens); res->t_inp_tokens = inp->tokens; - inp_per_layer = ggml_get_rows (ctx0, model.per_layer_tok_embd, inp->tokens); - inp_per_layer = ggml_reshape_3d(ctx0, inp_per_layer, n_embd_per_layer, n_layer, n_tokens); - inp_per_layer = ggml_scale (ctx0, inp_per_layer, tok_embd_scale); + inp_per_layer = ggml_get_rows(ctx0, model.per_layer_tok_embd, inp->tokens); cb(inp_per_layer, "inp_per_layer_selected", -1); } else { - // [TAG_GEMMA4_IMG_PADDING] // Multimodal embedding path: use padding token (ID=0) embedding // TODO: verify if this is the correct behavior in transformers implementation - const int64_t embd_size = model.per_layer_tok_embd->ne[0]; // n_embd_per_layer * n_layer - // Extract and dequantize padding token embedding (row 0) - ggml_tensor * padding = ggml_view_1d(ctx0, model.per_layer_tok_embd, embd_size, 0); - inp_per_layer = ggml_cast (ctx0, padding, GGML_TYPE_F32); - inp_per_layer = ggml_scale(ctx0, inp_per_layer, tok_embd_scale); + // [TAG_GEMMA4_IMG_PADDING] + inp->token0 = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 1); + ggml_set_input(inp->token0); + res->t_inp_tokens = inp->token0; - // Reshape to [n_embd_per_layer, n_layer, 1] - inp_per_layer = ggml_reshape_3d(ctx0, inp_per_layer, n_embd_per_layer, n_layer, 1); + inp_per_layer = ggml_get_rows(ctx0, model.per_layer_tok_embd, inp->token0); cb(inp_per_layer, "inp_per_layer_multimodal", -1); } res->add_input(std::move(inp)); diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index 88e2e88914..ad01b5388c 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -405,10 +405,6 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa int sections[4]; std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections); - ggml_tensor * inpL = build_inp_embd(model.tok_embd); - cb(inpL, "model.input_embed", -1); - ggml_build_forward_expand(gf, inpL); - auto * inp = build_inp_mem_hybrid(); // qwen4exp always builds llama_memory_hybrid_idx, so this downcast is safe @@ -421,6 +417,17 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa "the indexer cache must track the attention cache cell for cell"); } + ggml_tensor * ple_emb = nullptr; + if (hparams.ple_n_heads > 0) { + ple_emb = build_inp_ple(mctx_hyb); + // make sure ple_emb and build_inp_embd are in the same graph split + ggml_build_forward_expand(gf, ple_emb); + } + + ggml_tensor * inpL = build_inp_embd(model.tok_embd); + cb(inpL, "model.input_embed", -1); + ggml_build_forward_expand(gf, inpL); + // the QSA layers share one set of k-pool inputs // the CUDA lightning indexer takes 32 or 64 heads, QSA has a few, so it scores with plain ops llm_graph_input_kpool * inp_kpool = nullptr; @@ -431,13 +438,6 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa ggml_tensor * inp_pos = build_inp_pos(); ggml_tensor * inp_out_ids = build_inp_out_ids(); - ggml_tensor * ple_emb = nullptr; - if (hparams.ple_n_heads > 0) { - ple_emb = build_inp_ple(mctx_hyb); - // make sure ple_emb and build_inp_embd are in the same graph split - ggml_build_forward_expand(gf, ple_emb); - } - // the wide residual starts as hc identical copies of the embedding ggml_tensor * res_hc = ggml_repeat_4d(ctx0, ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens),