echodict/llama.cpp

1#include "models.h"2 3llm_build_nemotron_h::llm_build_nemotron_h(const llama_model & model, const llm_graph_params & params) :4    llm_build_mamba_base(params) {5    const int64_t n_embd_head = hparams.n_embd_head_v();6    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());7 8    ggml_tensor * cur;9    ggml_tensor * inpL;10 11    inpL = build_inp_embd(model.tok_embd);12    ggml_build_forward_expand(gf, inpL);13 14    auto * inp = build_inp_mem_hybrid();15 16    ggml_tensor * inp_out_ids = build_inp_out_ids();17 18    for (int il = 0; il < n_layer; ++il) {19        struct ggml_tensor * inpSA = inpL;20 21        // norm22        cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);23        cb(cur, "attn_norm", il);24 25        if (hparams.is_recurrent(il)) {26            // ssm layer //27            cur = build_mamba2_layer(inp->get_recr(), cur, model, ubatch, il);28        } else if (hparams.n_ff(il) == 0) {29            // attention layer //30            cur = build_attention_layer(cur, inp->get_attn(), model, n_embd_head, il);31        } else {32            cur = build_ffn_layer(cur, model, il);33        }34 35        if (il == n_layer - 1 && inp_out_ids) {36            cur   = ggml_get_rows(ctx0, cur, inp_out_ids);37            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);38        }39 40        // add residual41        cur = ggml_add(ctx0, cur, inpSA);42        cb(cur, "nemotron_h_block_out", il);43 44        // input for next layer45        inpL = cur;46    }47 48    cur = inpL;49 50    cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);51 52    cb(cur, "result_norm", -1);53    res->t_embd = cur;54 55    // lm_head56    cur = build_lora_mm(model.output, cur);57    cb(cur, "result_output", -1);58    res->t_logits = cur;59 60    ggml_build_forward_expand(gf, cur);61}62 63ggml_tensor * llm_build_nemotron_h::build_attention_layer(ggml_tensor *             cur,64                                                          llm_graph_input_attn_kv * inp_attn,65                                                          const llama_model &       model,66                                                                int64_t             n_embd_head,67                                                                int                 il) {68    auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il);69 70    const float kq_scale =71        hparams.f_attention_scale == 0.0f ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;72    cur = build_attn(inp_attn,73            model.layers[il].wo, model.layers[il].wo_b, model.layers[il].wo_s,74            Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);75    cb(cur, "attn_out", il);76    return cur;77}78 79ggml_tensor * llm_build_nemotron_h::build_ffn_layer(ggml_tensor * cur, const llama_model & model, int il) {80    if (model.layers[il].ffn_gate_inp == nullptr) {81        cur = build_ffn(cur,82                model.layers[il].ffn_up,   model.layers[il].ffn_up_b,   model.layers[il].ffn_up_s,83                NULL,                      NULL,                        NULL,84                model.layers[il].ffn_down, model.layers[il].ffn_down_b, model.layers[il].ffn_down_s,85                NULL,86                LLM_FFN_RELU_SQR, LLM_FFN_PAR, il);87        cb(cur, "ffn_out", il);88    } else {89        ggml_tensor * inp_emb    = cur;90        ggml_tensor * inp_latent = cur;91 92        if (model.layers[il].ffn_latent_down) {93            inp_latent = ggml_mul_mat(ctx0, model.layers[il].ffn_latent_down, cur);94        }95 96        ggml_tensor * router_logits = build_lora_mm(model.layers[il].ffn_gate_inp, cur);97        cb(router_logits, "ffn_moe_logits", il);98 99        ggml_tensor * moe_out =100            build_moe_ffn(inp_latent,101                    model.layers[il].ffn_gate_inp,102                    model.layers[il].ffn_up_exps,103                    nullptr, // no gate104                    model.layers[il].ffn_down_exps,105                    model.layers[il].ffn_exp_probs_b,106                    n_expert, n_expert_used,107                    LLM_FFN_RELU_SQR, hparams.expert_weights_norm,108                    hparams.expert_weights_scale,109                    LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID,110                    il,111                    router_logits, nullptr,112                    model.layers[il].ffn_up_exps_s,113                    nullptr, // no gate114                    model.layers[il].ffn_down_exps_s);115        cb(moe_out, "ffn_moe_out", il);116 117        if (model.layers[il].ffn_latent_up) {118            moe_out = ggml_mul_mat(ctx0, model.layers[il].ffn_latent_up, moe_out);119        }120 121        ggml_tensor * ffn_shexp = build_ffn(inp_emb,122                    model.layers[il].ffn_up_shexp,   NULL, model.layers[il].ffn_up_shexp_s,123                    NULL /* no gate */           ,   NULL, NULL,124                    model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s,125                    NULL,126                    LLM_FFN_RELU_SQR, LLM_FFN_PAR, il);127        cb(ffn_shexp, "ffn_shexp", il);128 129        cur = ggml_add(ctx0, moe_out, ffn_shexp);130        cb(cur, "ffn_out", il);131    }132 133    cur = build_cvec(cur, il);134    cb(cur, "l_out", il);135 136    return cur;137}138