LLM-Inference/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-08-12 22:29:39 +04:00
Code Issues Packages Projects Releases Wiki Activity

Muse-glimmer: Slightly better split mode graph (+2% TG)

Browse Source
This commit is contained in:
Kawrakow
2026-08-12 08:21:32 +00:00
parent 89a978b545
commit 140c72c735
3 changed files with 73 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/graphs/build_muse_glimmer.cpp
+45 -35
View File
@@ -33,55 +33,39 @@ ggml_cgraph * llm_build_context::build_muse_glimmer() {
ggml_tensor * ffn_inp = nullptr;
post_norm_data pnd;
pnd.f_rms_eps = post_norm_eps;
post_norm_data * pnd_ptr = nullptr;
bool add_input = model.split_mode == LLAMA_SPLIT_MODE_GRAPH ? false : true;
int n_active_layer = hparams.n_layer - hparams.nextn_predict_layers;
std::vector<ggml_tensor *> pn_tensors;
auto do_post_norm = [&] (ggml_tensor * cur, ggml_tensor * post_norm, ggml_tensor * inp, const std::string & tag, int il, bool get_rows) {
GGML_ASSERT(cur->op == GGML_OP_REDUCE);
int n = cur->op_params[1];
if ((int)pn_tensors.size() != n) pn_tensors.resize(n);
for (int id = 0; id < n; ++id) {
if (!cur->src[id]) {
pn_tensors[id] = nullptr;
continue;
}
auto pn_extra = (ggml_split_tensor_t *)post_norm->extra;
GGML_ASSERT(pn_extra && pn_extra->splits[id]);
auto normed = ggml_fused_rms_norm(ctx0, cur->src[id], pn_extra->splits[id], post_norm_eps);
cb(normed, (tag + "_pn").c_str(), 1000*(il+1) + id);
auto add = get_input_tensor_sm_graph(ctx0, inp, id);
if (get_rows && il == n_active_layer - 1 && inp_out_ids) {
add = ggml_get_rows(ctx0, add, inp_out_ids);
}
auto added = ggml_add(ctx0, normed, add);
cb(added, (tag + "_pn_add").c_str(), 1000*(il+1) + id);
pn_tensors[id] = added;
}
cur = ggml_reduce(ctx0, pn_tensors.data(), n, GGML_OP_ADD);
cb(cur, (tag + "_final").c_str(), il);
cur->op_params[3] = 1;
ggml_build_forward_expand(gf, cur);
return cur;
};
for (int il = 0; il < n_active_layer; ++il) {
bool use_rope = hparams.swa_layers[il];
auto this_KQ_mask = use_rope ? KQ_mask_swa : KQ_mask;
int this_n_swa = use_rope ? hparams.n_swa : 0;
if (model.split_mode == LLAMA_SPLIT_MODE_GRAPH && il > 0) {
GGML_ASSERT(pnd.next_input.size() == model.devices.size());
pnd.norm = model.layers[il-1].ffn_post_norm;
pnd_ptr = &pnd;
}
cur = build_std_attention(gf, model.layers[il].attn_norm, inpL,
inp_pos, il == n_active_layer - 1 ? inp_out_ids : nullptr, nullptr,
this_KQ_mask, nullptr, nullptr, kq_scale, 0.0f, this_n_swa, il, use_rope, false, add_input, false, false,
model.layers[il].attn_post_norm, -1, post_norm_eps, pnd_ptr);
if (model.split_mode == LLAMA_SPLIT_MODE_GRAPH) {
cur = do_post_norm(cur, model.layers[il].attn_post_norm, inpL, "attn", il, true);
pnd_ptr = &pnd;
if (il == 0) {
pnd.next_input.resize(model.devices.size(), inpL);
} else {
GGML_ASSERT(pnd.next_input.size() == model.devices.size());
}
pnd.norm = model.layers[il].attn_post_norm;
}
ffn_inp = cur;
@@ -95,10 +79,6 @@ ggml_cgraph * llm_build_context::build_muse_glimmer() {
model.layers[il].ffn_post_norm, post_norm_eps, pnd_ptr);
cb(cur, "ffn_out", il);
if (model.split_mode == LLAMA_SPLIT_MODE_GRAPH) {
cur = do_post_norm(cur, model.layers[il].ffn_post_norm, ffn_inp, "ffn", il, false);
}
cur = lctx.cvec.apply_to(ctx0, cur, il);
cb(cur, "l_out", il);
@@ -107,6 +87,36 @@ ggml_cgraph * llm_build_context::build_muse_glimmer() {
}
cur = inpL;
if (model.split_mode == LLAMA_SPLIT_MODE_GRAPH) {
GGML_ASSERT(inpL->op == GGML_OP_REDUCE);
int idx = model.default_layer_device[n_active_layer];
cur = inpL->src[idx];
if (!cur) {
for (idx = 0; idx < int(model.devices.size()); ++idx) {
if (inpL->src[idx]) {
cur = inpL->src[idx]; break;
}
}
GGML_ASSERT(cur);
}
auto pn_extra = (ggml_split_tensor_t *)model.layers[n_active_layer-1].ffn_post_norm->extra;
GGML_ASSERT(pn_extra && pn_extra->splits[idx]);
cur = ggml_fused_rms_norm(ctx0, cur, pn_extra->splits[idx], pnd.f_rms_eps);
cb(cur, "ffn_post_norm", n_active_layer-1);
GGML_ASSERT(idx < (int)pnd.next_input.size());
auto add = pnd.next_input[idx];
if (!add) {
for (int j = 0; j < int(pnd.next_input.size()); ++j) {
if (pnd.next_input[j]) {
add = pnd.next_input[j]; break;
}
}
GGML_ASSERT(add);
}
cur = ggml_add(ctx0, cur, add);
cb(cur, "ffn_final", -1);
}
// lm_head
cur = build_output(lctx, ctx0, cur, model.output, model.output_norm, cb);
cur = ggml_scale(ctx0, cur, hparams.f_logit_scale);
src/llama-build-context.cpp
+27 -18
View File
@@ -1093,6 +1093,28 @@ ggml_tensor * llm_build_context::do_split_norm(ggml_context * ctx, ggml_tensor *
return cur;
}
static ggml_tensor * llm_do_split_post_norm(ggml_context * ctx, ggml_tensor * cur, post_norm_data * pnd, int id, int n_device, const char * tag, int il_cb, const llm_build_cb & cb) {
auto pn_extra = (ggml_split_tensor_t *)pnd->norm->extra;
GGML_ASSERT(pn_extra && pn_extra->splits[id]);
GGML_ASSERT((int)pnd->next_input.size() == n_device);
cur = ggml_fused_rms_norm(ctx, cur, pn_extra->splits[id], pnd->f_rms_eps);
cb(cur, tag, il_cb);
auto add = pnd->next_input[id];
if (!add) {
for (int j = 0; j < n_device; ++j) {
if (pnd->next_input[j]) {
add = pnd->next_input[j];
break;
}
}
GGML_ASSERT(add);
}
cur = ggml_add(ctx, cur, add);
cb(cur, "inp_added", il_cb);
pnd->next_input[id] = cur;
return cur;
}
ggml_tensor * llm_build_context::llm_build_ffn(
ggml_context * ctx,
llama_context & lctx,
@@ -1136,15 +1158,7 @@ ggml_tensor * llm_build_context::llm_build_ffn(
if (!split_u) continue;
auto cur = get_input_tensor_sm_graph(ctx, input, id);
if (pnd) {
auto pn_extra = (ggml_split_tensor_t *)pnd->norm->extra;
GGML_ASSERT(pn_extra && pn_extra->splits[id]);
cur = ggml_fused_rms_norm(ctx, cur, pn_extra->splits[id], pnd->f_rms_eps);
cb(cur, "ffn_post_norm", il_cb);
if (pnd->add) {
auto add_id = get_input_tensor_sm_graph(ctx, pnd->add, id);
cur = ggml_add(ctx, cur, add_id);
cb(cur, "inp_added", il_cb);
}
cur = llm_do_split_post_norm(ctx, cur, pnd, id, u->n_device, "attn_post_norm", il_cb, cb);
}
cur = do_split_norm(ctx, cur, ffn_norm, lctx.model.hparams, cb, id, il_cb, is_norm);
if (input->op != GGML_OP_REDUCE) {
@@ -3111,15 +3125,7 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens
if (!split_wq) continue;
auto cur = get_input_tensor_sm_graph(ctx0, input, id);
if (pnd) {
auto pn_extra = (ggml_split_tensor_t *)pnd->norm->extra;
GGML_ASSERT(pn_extra && pn_extra->splits[id]);
cur = ggml_fused_rms_norm(ctx0, cur, pn_extra->splits[id], pnd->f_rms_eps);
cb(cur, "att_post_norm", il_cb);
if (pnd->add) {
auto add_id = get_input_tensor_sm_graph(ctx0, pnd->add, id);
cur = ggml_add(ctx0, cur, add_id);
cb(cur, "inp_added", il_cb);
}
cur = llm_do_split_post_norm(ctx0, cur, pnd, id, wq->n_device, "ffn_post_norm", il_cb, cb);
}
cur = do_split_norm(ctx0, cur, the_attn_norm, lctx.model.hparams, cb, id, il_cb, is_norm);
auto input_normed = cur;
@@ -3323,6 +3329,9 @@ ggml_tensor * llm_build_context::build_std_attention(ggml_cgraph * gf, ggml_tens
if (inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
cb(cur, "fa_get_rows", il_cb);
if (pnd) {
pnd->next_input[id] = ggml_get_rows(ctx0, pnd->next_input[id], inp_out_ids);
}
}
cur = llm_build_lora_mm(lctx, ctx0, split_wo, cur);
src/llama-build-context.h
+1 -1
View File
@@ -38,8 +38,8 @@ enum llm_norm_type {
};
struct post_norm_data {
std::vector<ggml_tensor *> next_input;
ggml_tensor * norm;
ggml_tensor * add;
float f_rms_eps;
};