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

Better routing for Gemma4-MoE (#1615)

This commit is contained in:
Kawrakow 2026-04-11 15:19:02 +02:00 committed by GitHub
parent b0750b5d43
commit 08ae48c667
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 62 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
src/llama-build-context.cpp
View File

@ -6033,7 +6033,7 @@ ggml_cgraph * llm_build_context::build_gemma3() {
return gf;
}
static ggml_cgraph * build_gemma4_graph_paralle(llm_build_context & llm, llama_context & lctx, ggml_context * ctx0,
static ggml_cgraph * build_gemma4_graph_parallel(llm_build_context & llm, llama_context & lctx, ggml_context * ctx0,
ggml_tensor * inpL, ggml_tensor * inp_pos, ggml_tensor * inp_out_ids,
ggml_tensor * KQ_mask, ggml_tensor * KQ_mask_swa, int n_tokens, const llm_build_cb & cb) {
auto & model = lctx.model;
@ -6307,14 +6307,12 @@ static ggml_cgraph * build_gemma4_graph_paralle(llm_build_context & llm, llama_c
if (is_moe) {
cur = do_split_norm(ctx0, ffn_inp[id], model.layers[il].ffn_pre_norm_2, hparams, cb, id, il_cb, false);
cb(cur, "ffn_moe_inp", il_cb);
auto tmp = ggml_rms_norm(ctx0, ffn_inp[id], hparams.f_norm_rms_eps);
auto tmp = ggml_fused_rms_norm(ctx0, ffn_inp[id],
((const ggml_split_tensor_t *)model.layers[il].ffn_gate_inp_s->extra)->splits[id], hparams.f_norm_rms_eps);
cb(tmp, "tmp", il_cb);
tmp = ggml_scale(ctx0, tmp, 1.0f / sqrtf((float) hparams.n_embd));
cb(tmp, "tmp_scaled", il_cb);
tmp = ggml_mul(ctx0, tmp, ((const ggml_split_tensor_t *)model.layers[il].ffn_gate_inp_s->extra)->splits[id]);
cb(tmp, "tmp_mul", il_cb);
auto logits = llm.llm_build_lora_mm(lctx, ctx0, ((const ggml_split_tensor_t *)model.layers[il].ffn_gate_inp->extra)->splits[id], tmp);
cb(logits, "logits", il_cb);
ggml_build_forward_expand(gf, logits);
auto moe = llm. llm_build_moe_ffn(ctx0, lctx, cur,
nullptr, nullptr, nullptr,
@ -6432,6 +6430,7 @@ ggml_cgraph * llm_build_context::build_gemma4() {
struct ggml_tensor * inpL;
inpL = llm_build_inp_embd(ctx0, lctx, hparams, batch, model.tok_embd, cb);
cb(inpL, "tok_embd", -1);
// important: do not normalize weights for raw embeddings input (i.e. encoded image emdeddings)
if (batch.token) {
@ -6475,7 +6474,7 @@ ggml_cgraph * llm_build_context::build_gemma4() {
}
if (model.split_mode == LLAMA_SPLIT_MODE_GRAPH) {
return build_gemma4_graph_paralle(*this, lctx, ctx0, inpL, inp_pos, inp_out_ids,
return build_gemma4_graph_parallel(*this, lctx, ctx0, inpL, inp_pos, inp_out_ids,
KQ_mask, KQ_mask_swa, n_tokens, cb);
}
@ -6577,9 +6576,8 @@ ggml_cgraph * llm_build_context::build_gemma4() {
cb(cur_moe, "ffn_norm_2", il);
// custom MoE logits calculation (router operates on attn_out, not cur)
auto tmp = ggml_rms_norm(ctx0, attn_out, hparams.f_norm_rms_eps);
tmp = ggml_scale(ctx0, tmp, 1.0f / sqrtf((float) n_embd));
tmp = ggml_mul(ctx0, tmp, model.layers[il].ffn_gate_inp_s);
auto tmp = ggml_fused_rms_norm(ctx0, attn_out, model.layers[il].ffn_gate_inp_s, hparams.f_norm_rms_eps);
cb(tmp, "tmp", il);
auto logits = llm_build_lora_mm(lctx, ctx0, model.layers[il].ffn_gate_inp, tmp); // [n_expert, n_tokens]
cb(logits, "ffn_moe_logits", il);

1
src/llama-model.h
View File

@ -426,6 +426,7 @@ struct llama_model {
std::vector<rpc_device> rpc_servers;
std::vector<int32_t> devices;
std::vector<int32_t> default_layer_device;
std::vector<float> aux_buffer;
// gguf metadata
std::unordered_map<std::string, std::string> gguf_kv;

53
src/llama.cpp
View File

@ -1934,6 +1934,56 @@ static void llm_prepare_mla(llama_model & model, int mla) {
ggml_free(ctx);
}
static void llm_scale_gate_inp_s(llama_model & model, bool uses_mmap) {
auto & hparams = model.hparams;
printf("%s: n_embd = %d\n", __func__, hparams.n_embd);
std::vector<float> values(hparams.n_embd);
float scale = 1.0f/sqrtf((float)hparams.n_embd);
int n_host = 0;
for (auto & l : model.layers) {
auto gis = l.ffn_gate_inp_s;
if (!gis) continue;
bool is_host = ggml_backend_buffer_is_host(gis->buffer);
if (is_host) {
if (uses_mmap) {
++n_host;
} else {
auto val = (float *)gis->data;
for (int j = 0; j < hparams.n_embd; ++j) val[j] *= scale;
}
} else {
auto extra = (ggml_split_tensor_t *)gis->extra;
if (extra) {
ggml_backend_tensor_get(extra->splits[0], values.data(), 0, values.size()*sizeof(float));
} else {
ggml_backend_tensor_get(gis, values.data(), 0, values.size()*sizeof(float));
}
for (int j = 0; j < hparams.n_embd; ++j) values[j] *= scale;
if (extra) {
for (int id = 0; id < extra->n_device; ++id) {
ggml_backend_tensor_set(extra->splits[id], values.data(), 0, values.size()*sizeof(float));
}
} else {
ggml_backend_tensor_set(gis, values.data(), 0, values.size()*sizeof(float));
}
}
}
if (n_host > 0) {
model.aux_buffer.resize(n_host * hparams.n_embd * sizeof(float));
auto ptr = model.aux_buffer.data();
for (auto & l : model.layers) {
auto gis = l.ffn_gate_inp_s;
if (!gis) continue;
if (ggml_backend_buffer_is_host(gis->buffer)) {
auto val = (const float *)gis->data;
for (int j = 0; j < hparams.n_embd; ++j) ptr[j] = val[j] * scale;
gis->data = ptr;
ptr += hparams.n_embd;
}
}
}
}
// Backend (reg) enumeration
static bool striequals(const char* a, const char* b) {
for (; *a && *b; a++, b++) {
@ -2759,6 +2809,9 @@ static bool llm_load_tensors(
if (model.arch == LLM_ARCH_DEEPSEEK2 || model.arch == LLM_ARCH_GLM_DSA || model.arch == LLM_ARCH_MISTRAL4) {
llm_prepare_mla(model, mla_attn);
}
if (model.arch == LLM_ARCH_GEMMA4) {
llm_scale_gate_inp_s(model, use_mmap_buffer);
}
if (use_mmap_buffer) {
for (auto & mapping : ml.mappings) {