echodict/llama.cpp

1#include "models.h"2 3llm_build_grovemoe::llm_build_grovemoe(const llama_model & model, const llm_graph_params & params) :4    llm_graph_context(params) {5    const int64_t n_embd_head    = hparams.n_embd_head_v();6    const int64_t n_chunk_expert = n_expert / hparams.n_group_experts;7 8    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());9    GGML_ASSERT(n_embd_head == n_rot);10 11    ggml_tensor * cur;12    ggml_tensor * inpL;13 14    inpL = build_inp_embd(model.tok_embd);15 16    // inp_pos - contains the positions17    ggml_tensor * inp_pos = build_inp_pos();18 19    auto * inp_attn = build_attn_inp_kv();20 21    ggml_tensor * inp_out_ids = build_inp_out_ids();22 23    for (int il = 0; il < n_layer; ++il) {24        ggml_tensor * inpSA = inpL;25 26        // norm27        cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);28        cb(cur, "attn_norm", il);29 30        // self_attention31        {32            // compute Q and K and RoPE them33            auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,34                    n_embd_head, n_head, n_head_kv, il);35 36            Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);37            cb(Qcur, "Qcur_normed", il);38 39            Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,40                                 ext_factor, attn_factor, beta_fast, beta_slow);41 42            Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);43            cb(Kcur, "Kcur_normed", il);44 45            Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,46                                 ext_factor, attn_factor, beta_fast, beta_slow);47 48            cb(Qcur, "Qcur", il);49            cb(Kcur, "Kcur", il);50            cb(Vcur, "Vcur", il);51 52            cur = build_attn(inp_attn,53                    model.layers[il].wo, model.layers[il].wo_b, model.layers[il].wo_s,54                    Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f / sqrtf(float(n_embd_head)), il);55        }56 57        if (il == n_layer - 1 && inp_out_ids) {58            cur   = ggml_get_rows(ctx0, cur, inp_out_ids);59            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);60        }61 62        ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);63        cb(ffn_inp, "ffn_inp", il);64 65        // MoE branch66        cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il);67        cb(cur, "ffn_norm", il);68 69        ggml_tensor * probs = build_lora_mm(model.layers[il].ffn_gate_inp, cur);  // [n_expert, n_tokens]70        cb(probs, "ffn_moe_logits", il);71 72        ggml_tensor * moe_out =73            build_moe_ffn(cur,74                nullptr,75                model.layers[il].ffn_up_exps,76                model.layers[il].ffn_gate_exps,77                model.layers[il].ffn_down_exps,78                nullptr,79                n_expert, n_expert_used,80                LLM_FFN_SILU, true,81                hparams.expert_weights_scale,82                LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX,83                il,84                probs);85        cb(moe_out, "ffn_moe_out", il);86        cur = moe_out;87 88        // TODO: Only do the expert selection and weights once89        moe_out = build_moe_ffn(cur,90                    nullptr,91                    model.layers[il].ffn_up_chexps,92                    model.layers[il].ffn_gate_chexps,93                    model.layers[il].ffn_down_chexps,94                    nullptr,95                    n_chunk_expert, n_expert_used > n_chunk_expert ? n_chunk_expert : n_expert_used,96                    LLM_FFN_SILU, true,97                    hparams.expert_weights_scale,98                    LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX,99                    il,100                    probs);101        cb(moe_out, "ffn_adj_moe_out", il);102 103        cur = ggml_add(ctx0, cur, ggml_scale(ctx0, moe_out, hparams.expert_group_scale));104        cb(cur, "ffn_final_moe_out", il);105 106        cur = ggml_add(ctx0, cur, ffn_inp);107 108        cur = build_cvec(cur, il);109        cb(cur, "l_out", il);110 111        // input for next layer112        inpL = cur;113    }114 115    cur = inpL;116 117    cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);118 119    cb(cur, "result_norm", -1);120    res->t_embd = cur;121 122    // lm_head123    cur = build_lora_mm(model.output, cur);124 125    cb(cur, "result_output", -1);126    res->t_logits = cur;127 128    ggml_build_forward_expand(gf, cur);129}130