diff --git a/src/llama-build-context.cpp b/src/llama-build-context.cpp index 335794dd..5c809c57 100644 --- a/src/llama-build-context.cpp +++ b/src/llama-build-context.cpp @@ -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); diff --git a/src/llama-model.h b/src/llama-model.h index 04797ba5..c3dce07f 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -426,6 +426,7 @@ struct llama_model { std::vector rpc_servers; std::vector devices; std::vector default_layer_device; + std::vector aux_buffer; // gguf metadata std::unordered_map gguf_kv; diff --git a/src/llama.cpp b/src/llama.cpp index 6e604aae..069ede9d 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -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 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) {