echodict/llama.cpp

1#include "models.h"2 3llm_build_afmoe::llm_build_afmoe(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {4    const int64_t n_embd_head = hparams.n_embd_head_v();5    GGML_ASSERT(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    // MuP scaling: embeddings * sqrt(hidden_size)13    // mup_enabled = true, hidden_size = 1024, scale = 32.014    inpL = ggml_scale(ctx0, inpL, sqrtf(float(n_embd)));15    cb(inpL, "inp_embd_scaled", -1);16 17    // inp_pos - contains the positions18    ggml_tensor * inp_pos = build_inp_pos();19    auto * inp_attn = build_attn_inp_kv_iswa();20    ggml_tensor * inp_out_ids = build_inp_out_ids();21 22    const float kq_scale = 1.0f/sqrtf(float(n_embd_head));23 24    for (int il = 0; il < n_layer; ++il) {25        const float freq_base_l  = model.get_rope_freq_base (cparams, il);26        const float freq_scale_l = model.get_rope_freq_scale(cparams, il);27 28        ggml_tensor * inpSA = inpL;29 30        // This overlaps with SWA layers in current models, so get_rope_freq_base/scale may be superfluous31        const bool use_rope = hparams.n_no_rope_layer_step > 0 &&32                              (il + 1) % hparams.n_no_rope_layer_step != 0;33 34        // dual attention normalization (pre)35        cur = build_norm(inpL,36                model.layers[il].attn_norm, NULL,37                LLM_NORM_RMS, il);38        cb(cur, "attn_norm", il);39 40        // self-attention41        {42            ggml_tensor * attn_inp = cur;  // save input for gate computation43 44            auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,45                    n_embd_head, n_head, n_head_kv, il);46 47            // compute gate from input48            ggml_tensor * gate = build_lora_mm(model.layers[il].wqkv_gate, attn_inp);49            cb(gate, "attn_gate_proj", il);50 51            // Q/K normalization52            Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);53            Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);54            cb(Qcur, "Qcur_normed", il);55            cb(Kcur, "Kcur_normed", il);56 57            if (use_rope) {58                Qcur = ggml_rope_ext(59                        ctx0, Qcur, inp_pos, nullptr,60                        n_rot, rope_type, n_ctx_orig, freq_base_l, freq_scale_l,61                        ext_factor, attn_factor, beta_fast, beta_slow);62                cb(Qcur, "Qcur_rope", 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                cb(Kcur, "Kcur_rope", il);69            }70 71            cur = build_attn(inp_attn,72                    NULL, NULL, NULL,  // wo will be applied after gating73                    Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);74            cb(cur, "attn_out", il);75 76            // attention gating: attn_out * sigmoid(gate) BEFORE o_proj77            gate = ggml_sigmoid(ctx0, gate);78            cb(gate, "attn_gate_sig", il);79            cur = ggml_mul(ctx0, cur, gate);80            cb(cur, "attn_gated", il);81 82            // now apply output projection83            cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s);84            cb(cur, "attn_o_proj", il);85        }86 87        // dual attention normalization (post)88        cur = build_norm(cur,89                model.layers[il].attn_post_norm, NULL,90                LLM_NORM_RMS, il);91        cb(cur, "attn_post_norm", il);92 93        if (il == n_layer - 1 && inp_out_ids) {94            cur   = ggml_get_rows(ctx0,   cur, inp_out_ids);95            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);96        }97 98        ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);99        cb(ffn_inp, "ffn_inp", il);100 101        // dual ffn normalization (pre)102        cur = build_norm(ffn_inp,103                model.layers[il].ffn_norm, NULL,104                LLM_NORM_RMS, il);105        cb(cur, "ffn_norm", il);106 107        // MoE or dense FFN108        if ((uint32_t)il >= hparams.n_layer_dense_lead) {109            // MoE layer with sigmoid routing, normalization, and scaling110            ggml_tensor * moe_out = build_moe_ffn(cur,111                    model.layers[il].ffn_gate_inp,112                    model.layers[il].ffn_up_exps,113                    model.layers[il].ffn_gate_exps,114                    model.layers[il].ffn_down_exps,115                    model.layers[il].ffn_exp_probs_b,116                    n_expert, n_expert_used,117                    LLM_FFN_SILU,118                    hparams.expert_weights_norm,           // norm_w (route_norm=True)119                    hparams.expert_weights_scale,          // w_scale (route_scale=2.826)120                    (llama_expert_gating_func_type) hparams.expert_gating_func,121                    il);122            cb(moe_out, "ffn_moe_out", il);123 124            // shared expert125            if (hparams.n_expert_shared > 0) {126                ggml_tensor * ffn_shexp = build_ffn(cur,127                        model.layers[il].ffn_up_shexp,   NULL, NULL,128                        model.layers[il].ffn_gate_shexp, NULL, NULL,129                        model.layers[il].ffn_down_shexp, NULL, NULL,130                        NULL,131                        LLM_FFN_SILU, LLM_FFN_PAR, il);132                cb(ffn_shexp, "ffn_shexp", il);133 134                cur = ggml_add(ctx0, moe_out, ffn_shexp);135                cb(cur, "ffn_out", il);136            } else {137                cur = moe_out;138            }139        } else {140            // dense layer141            cur = build_ffn(cur,142                    model.layers[il].ffn_up,   NULL, NULL,143                    model.layers[il].ffn_gate, NULL, NULL,144                    model.layers[il].ffn_down, NULL, NULL,145                    NULL,146                    LLM_FFN_SILU, LLM_FFN_PAR, il);147            cb(cur, "ffn_out", il);148        }149 150        // dual ffn normalization (post)151        cur = build_norm(cur,152                model.layers[il].ffn_post_norm, NULL,153                LLM_NORM_RMS, il);154        cb(cur, "ffn_post_norm", il);155 156        cur = ggml_add(ctx0, cur, ffn_inp);157        cur = build_cvec(cur, il);158        cb(cur, "l_out", il);159 160        // input for next layer161        inpL = cur;162    }163 164    cur = inpL;165 166    cur = build_norm(cur,167            model.output_norm, NULL,168            LLM_NORM_RMS, -1);169    cb(cur, "result_norm", -1);170 171    res->t_embd = cur;172 173    // lm_head174    cur = build_lora_mm(model.output, cur);175    cb(cur, "result_output", -1);176    res->t_logits = cur;177 178    ggml_build_forward_expand(gf, cur);179}180