Commit 8e2d31e0e for llama.cpp
commit 8e2d31e0eb658d0382d3106a61369dff81f0c080
Author: Georgi Gerganov <ggerganov@gmail.com>
Date: Fri Oct 9 20:43:05 2026 +0300
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
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index 5a6a2b320..1514b8aeb 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<llm_graph_input_embd>(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;
+
+ const int64_t n_tok_rows = has_mixed ? llm_graph_n_tok_rows(ubatch) : 0;
- 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);
+ // 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;
- // token embeddings with lora and padding
- auto build_tok = [&](ggml_tensor * ids) {
- ggml_tensor * cur = ggml_get_rows(ctx0, tok_embd, ids);
+ // 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<ggml_tensor *, 3> 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;
- // 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);
+ assert(ggml_are_same_shape (inps[0], inps[1]));
+ assert(ggml_are_same_stride(inps[0], inps[1]));
+ if (has_mixed) {
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 fbc4d2a6e..d88338437 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<llm_graph_input_gemma4_ple>(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 88e2e8891..ad01b5388 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),