nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Apply callback to extract features in spec (#2348)

This commit is contained in:
Samuel Oliveira Alves 2026-08-24 07:17:23 -03:00 committed by GitHub
parent 477852c1c9
commit ad26e68bee
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/graphs/build_gemma4.cpp
View File

@ -215,6 +215,7 @@ static ggml_cgraph * build_gemma4_graph_parallel(llm_build_context & llm, llama_
int nhave = 0; int nhave = 0;
ggml_tensor * sa_last = nullptr; ggml_tensor * sa_last = nullptr;
bool emitted_prev_l_out = false;
for (int id = 0; id < n_device; ++id) { for (int id = 0; id < n_device; ++id) {
GGML_ASSERT((wq->splits[id] && wk->splits[id] && (!wv || wv->splits[id]) && wo->splits[id]) || GGML_ASSERT((wq->splits[id] && wk->splits[id] && (!wv || wv->splits[id]) && wo->splits[id]) ||
(!wq->splits[id] && !wk->splits[id] && (!wv || !wv->splits[id]) && !wo->splits[id])); (!wq->splits[id] && !wk->splits[id] && (!wv || !wv->splits[id]) && !wo->splits[id]));
@ -257,6 +258,10 @@ static ggml_cgraph * build_gemma4_graph_parallel(llm_build_context & llm, llama_
sa_inp[id] = ggml_mul(ctx0, sa_inp[id], scale->splits[id]); sa_inp[id] = ggml_mul(ctx0, sa_inp[id], scale->splits[id]);
cb(sa_inp[id], "sa_inp_scaled", il_cb); cb(sa_inp[id], "sa_inp_scaled", il_cb);
} }
if (!emitted_prev_l_out) {
cb(sa_inp[id], "l_out", il - 1);
emitted_prev_l_out = true;
}
} }
auto cur = llm_build_context::do_split_norm(ctx0, sa_inp[id], model.layers[il].attn_norm, hparams, cb, id, il_cb, false); auto cur = llm_build_context::do_split_norm(ctx0, sa_inp[id], model.layers[il].attn_norm, hparams, cb, id, il_cb, false);
cb(cur, "sa_inp_normed", il_cb); cb(cur, "sa_inp_normed", il_cb);
@ -517,6 +522,7 @@ static ggml_cgraph * build_gemma4_graph_parallel(llm_build_context & llm, llama_
cur = ggml_mul(ctx0, cur, scale->splits[idx]); cur = ggml_mul(ctx0, cur, scale->splits[idx]);
cb(cur, "ffn_out_scaled", hparams.n_layer-1); cb(cur, "ffn_out_scaled", hparams.n_layer-1);
} }
cb(cur, "l_out", hparams.n_layer-1);
cur = llm_build_context::build_output(lctx, ctx0, cur, model.output, model.output_norm, cb); cur = llm_build_context::build_output(lctx, ctx0, cur, model.output, model.output_norm, cb);
if (hparams.f_final_logit_softcapping > 0) { if (hparams.f_final_logit_softcapping > 0) {