Fix Gemma4-E4B compute graph

This commit is contained in:
Kawrakow
2026-05-21 08:19:30 +00:00
parent 3dd282358b
commit f1e146859b
2 changed files with 10 additions and 8 deletions
+9 -7
View File
@@ -720,10 +720,18 @@ ggml_cgraph * llm_build_context::build_gemma4() {
auto inp_out_ids = n_tokens > 1 ? build_inp_out_ids() : nullptr;
if (model.split_mode == LLAMA_SPLIT_MODE_GRAPH) {
return build_gemma4_graph_parallel(*this, lctx, ctx0, inpL, inp_pos, inp_out_ids,
KQ_mask, KQ_mask_swa, n_tokens, cb);
}
auto gf = ggml_new_graph_custom(ctx0, model.max_nodes(n_tokens), false);
ggml_tensor * inp_per_layer = nullptr;
if (model.tok_embd_per_layer) {
if (batch.token) {
inp_per_layer = ggml_get_rows(ctx0, model.tok_embd_per_layer, lctx.inp_tokens);
ggml_build_forward_expand(gf, inp_per_layer);
inp_per_layer = ggml_reshape_3d(ctx0, inp_per_layer, hparams.n_embd_per_layer, n_layer, n_tokens);
inp_per_layer = ggml_scale(ctx0, inp_per_layer, sqrtf((float) hparams.n_embd_per_layer));
cb(inp_per_layer, "inp_per_layer_selected", -1);
@@ -735,6 +743,7 @@ ggml_cgraph * llm_build_context::build_gemma4() {
// Extract and dequantize padding token embedding (row 0)
auto padding = ggml_view_1d(ctx0, model.tok_embd_per_layer, embd_size, 0);
inp_per_layer = ggml_cast(ctx0, padding, GGML_TYPE_F32);
ggml_build_forward_expand(gf, inp_per_layer);
// Reshape to [n_embd_per_layer, n_layer, 1]
inp_per_layer = ggml_reshape_3d(ctx0, inp_per_layer, hparams.n_embd_per_layer, n_layer, 1);
@@ -745,13 +754,6 @@ ggml_cgraph * llm_build_context::build_gemma4() {
}
if (model.split_mode == LLAMA_SPLIT_MODE_GRAPH) {
return build_gemma4_graph_parallel(*this, lctx, ctx0, inpL, inp_pos, inp_out_ids,
KQ_mask, KQ_mask_swa, n_tokens, cb);
}
auto gf = ggml_new_graph_custom(ctx0, model.max_nodes(n_tokens), false);
// "5-to-1 interleaved attention"
// 5 layers of local attention followed by 1 layer of global attention
// we do this via swa_layers now
+1 -1
View File
@@ -2060,7 +2060,7 @@ bool create_tensors_helper::create_gemma4_tensors(const LLM_TN & tn) {
model.tok_embd = create_tensor(ctx_input, tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
if (n_embd_per_layer > 0) {
model.tok_embd_per_layer = create_tensor(ctx_output, tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), {n_embd_per_layer * n_layer, n_vocab}, 0);
model.tok_embd_per_layer = create_tensor(ctx_input, tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), {n_embd_per_layer * n_layer, n_vocab}, 0);
model.per_layer_model_proj = create_tensor(ctx_output, tn(LLM_TENSOR_PER_LAYER_MODEL_PROJ, "weight"), {n_embd, n_embd_per_layer * n_layer}, 0);
model.per_layer_proj_norm = create_tensor(ctx_output, tn(LLM_TENSOR_PER_LAYER_PROJ_NORM, "weight"), {n_embd_per_layer}, 0);
}