echodict/llama.cpp

1#include "models.h"2 3template <bool iswa>4llm_build_gemma3<iswa>::llm_build_gemma3(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {5    const int64_t n_embd_head = hparams.n_embd_head_k();6 7    ggml_tensor * cur;8    ggml_tensor * inpL;9 10    inpL = build_inp_embd(model.tok_embd);11 12    // important: do not normalize weights for raw embeddings input (i.e. encoded image embeddings)13    inpL = ggml_scale(ctx0, inpL, ubatch.token ? sqrtf(n_embd) : 1.0f);14    cb(inpL, "inp_scaled", -1);15 16    // inp_pos - contains the positions17    ggml_tensor * inp_pos = build_inp_pos();18 19    // TODO: is causal == true correct? might need some changes20    using inp_attn_type = std::conditional_t<iswa, llm_graph_input_attn_kv_iswa, llm_graph_input_attn_kv>;21    inp_attn_type * inp_attn = nullptr;22 23    if constexpr (iswa) {24        inp_attn = build_attn_inp_kv_iswa();25    } else {26        inp_attn = build_attn_inp_kv();27    }28 29    ggml_tensor * inp_out_ids = build_inp_out_ids();30 31    for (int il = 0; il < n_layer; ++il) {32        float freq_base_l  = 0.0f;33        float freq_scale_l = 0.0f;34 35        if constexpr (iswa) {36            freq_base_l  = model.get_rope_freq_base (cparams, il);37            freq_scale_l = model.get_rope_freq_scale(cparams, il);38        } else {39            freq_base_l  = freq_base;40            freq_scale_l = freq_scale;41        }42 43        // norm44        cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);45        cb(cur, "attn_norm", il);46 47        // self-attention48        {49            // compute Q and K and RoPE them50            auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,51                    n_embd_head, n_head, n_head_kv, il);52 53            Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);54            cb(Qcur, "Qcur_normed", il);55 56            Qcur = ggml_rope_ext(57                    ctx0, Qcur, inp_pos, nullptr,58                    n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,59                    ext_factor, attn_factor, beta_fast, beta_slow);60 61            Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);62            cb(Kcur, "Kcur_normed", il);63 64            Kcur = ggml_rope_ext(65                    ctx0, Kcur, inp_pos, nullptr,66                    n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,67                    ext_factor, attn_factor, beta_fast, beta_slow);68 69            cb(Qcur, "Qcur", il);70            cb(Kcur, "Kcur", il);71            cb(Vcur, "Vcur", il);72 73            // ref: https://github.com/google/gemma_pytorch/blob/014acb7ac4563a5f77c76d7ff98f31b568c16508/gemma/model.py#L31574            Qcur = ggml_scale(ctx0, Qcur, hparams.f_attention_scale);75 76            cur = build_attn(inp_attn,77                    model.layers[il].wo, NULL, model.layers[il].wo_s,78                    Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f, il);79        }80        if (il == n_layer - 1 && inp_out_ids) {81            cur  = ggml_get_rows(ctx0,  cur, inp_out_ids);82            inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);83        }84        cur = build_norm(cur,85                model.layers[il].attn_post_norm, NULL,86                LLM_NORM_RMS, il);87        cb(cur, "attn_post_norm", il);88 89        ggml_tensor * sa_out = ggml_add(ctx0, cur, inpL);90        cb(sa_out, "sa_out", il);91 92        cur = build_norm(sa_out,93                model.layers[il].ffn_norm, NULL,94                LLM_NORM_RMS, il);95        cb(cur, "ffn_norm", il);96 97        // feed-forward network98        {99            cur = build_ffn(cur,100                    model.layers[il].ffn_up,   NULL, NULL,101                    model.layers[il].ffn_gate, NULL, NULL,102                    model.layers[il].ffn_down, NULL, NULL,103                    NULL,104                    LLM_FFN_GELU, LLM_FFN_PAR, il);105            cb(cur, "ffn_out", il);106        }107        cur = build_norm(cur,108                model.layers[il].ffn_post_norm, NULL,109                LLM_NORM_RMS, -1);110        cb(cur, "ffn_post_norm", il);111 112        cur = ggml_add(ctx0, cur, sa_out);113 114        cur = build_cvec(cur, il);115        cb(cur, "l_out", il);116 117        // input for next layer118        inpL = cur;119    }120    cur = inpL;121 122    cur = build_norm(cur,123            model.output_norm, NULL,124            LLM_NORM_RMS, -1);125 126    cb(cur, "result_norm", -1);127    res->t_embd = cur;128 129    // lm_head130    cur = build_lora_mm(model.output, cur);131 132    if (hparams.f_final_logit_softcapping) {133        cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping);134        cur = ggml_tanh(ctx0, cur);135        cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping);136    }137 138    cb(cur, "result_output", -1);139    res->t_logits = cur;140 141    ggml_build_forward_expand(gf, cur);142}143 144template struct llm_build_gemma3<false>;145template struct llm_build_gemma3<true>;146