echodict/llama.cpp

1#include "models.h"2 3template <bool iswa>4llm_build_smallthinker<iswa>::llm_build_smallthinker(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params){5    const int64_t n_embd_head = hparams.n_embd_head_v();6 7    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());8    GGML_ASSERT(n_embd_head == n_rot);9 10    ggml_tensor * cur;11    ggml_tensor * inpL;12 13    inpL = build_inp_embd(model.tok_embd);14 15    // inp_pos - contains the positions16    ggml_tensor * inp_pos = build_inp_pos();17 18    using inp_attn_type = std::conditional_t<iswa, llm_graph_input_attn_kv_iswa, llm_graph_input_attn_kv>;19    inp_attn_type * inp_attn = nullptr;20 21    if constexpr (iswa) {22        inp_attn = build_attn_inp_kv_iswa();23    } else {24        inp_attn = build_attn_inp_kv();25    }26    ggml_tensor * inp_out_ids = build_inp_out_ids();27 28    for (int il = 0; il < n_layer; ++il) {29        const float freq_base_l  = model.get_rope_freq_base (cparams, il);30        const float freq_scale_l = model.get_rope_freq_scale(cparams, il);31 32        ggml_tensor * inpSA  = inpL;33 34        // This overlaps with SWA layers in current models, so get_rope_freq_base/scale may be superfluous35        const bool use_rope = hparams.n_no_rope_layer_step == n_layer ||36                              il % hparams.n_no_rope_layer_step != 0;37 38        ggml_tensor * probs = build_lora_mm(model.layers[il].ffn_gate_inp, inpL);  // [n_expert, n_tokens]39        cb(probs, "ffn_moe_logits", il);40 41        // norm42        cur = build_norm(inpL,model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);43        cb(cur, "attn_norm", il);44 45        // self_attention46        {47            // compute Q and K and RoPE them48            auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,49                    n_embd_head, n_head, n_head_kv, il);50 51            if (use_rope) {52                Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,53                                    ext_factor, attn_factor, beta_fast, beta_slow);54 55                Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,56                                    ext_factor, attn_factor, beta_fast, beta_slow);57            }58            cb(Qcur, "Qcur", il);59            cb(Kcur, "Kcur", il);60 61            cur = build_attn(inp_attn,62                    model.layers[il].wo, model.layers[il].wo_b, model.layers[il].wo_s,63                    Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f / sqrtf(float(n_embd_head)), il);64        }65        if (il == n_layer - 1 && inp_out_ids) {66            cur = ggml_get_rows(ctx0, cur, inp_out_ids);67            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);68            probs = ggml_get_rows(ctx0, probs, inp_out_ids);69        }70        ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);71        cb(ffn_inp, "ffn_inp", il);72 73        // MoE branch74        cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il);75        cb(cur, "ffn_norm", il);76 77        ggml_tensor * ffn_out =78            build_moe_ffn(cur,79                    nullptr,80                    model.layers[il].ffn_up_exps,81                    model.layers[il].ffn_gate_exps,82                    model.layers[il].ffn_down_exps,83                    nullptr,84                    n_expert, n_expert_used,85                    LLM_FFN_RELU, true,86                    hparams.expert_weights_scale,87                    static_cast<llama_expert_gating_func_type>(hparams.expert_gating_func),88                    il, probs);89 90        cb(ffn_out, "ffn_out", il);91        cur = ffn_out;92 93        cur = ggml_add(ctx0, cur, ffn_inp);94 95        cur = build_cvec(cur, il);96        cb(cur, "l_out", il);97 98        // input for next layer99        inpL = cur;100    }101    cur = inpL;102 103    cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);104    cb(cur, "result_norm", -1);105    res->t_embd = cur;106 107    // lm_head108    cur = build_lora_mm(model.output, cur);109    cb(cur, "result_output", -1);110    res->t_logits = cur;111 112    ggml_build_forward_expand(gf, cur);113}114 115// Explicit template instantiations116template struct llm_build_smallthinker<false>;117template struct llm_build_smallthinker<true>;118