echodict/llama.cpp

1#include "models.h"2 3llm_build_deepseek2::llm_build_deepseek2(const llama_model & model, const llm_graph_params & params) :4    llm_graph_context(params) {5    // lite variants include DeepSeek-V2-Lite, GigaChat3-10B-A1.8B6    bool is_ocr = model.arch == LLM_ARCH_DEEPSEEK2OCR;7 8    const bool is_mla = hparams.is_mla();9 10    // note: these are the actual head sizes you get when treating as MHA or after "decompression" using wv_b for MLA11    const int64_t n_embd_head_k = hparams.n_embd_head_k_mla();12    const int64_t n_embd_head_v = hparams.n_embd_head_v_mla();13 14    const int64_t n_embd_head_qk_rope = hparams.n_rot();15    const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope;16 17    const uint32_t kv_lora_rank = hparams.n_lora_kv;18 19    // We have to pre-scale kq_scale and attn_factor to make the YaRN RoPE work correctly.20    // See https://github.com/ggml-org/llama.cpp/discussions/7416 for detailed explanation.21    // And also: https://github.com/ggml-org/llama.cpp/pull/17945 [TAG_DEEPSEEK2_YARN_LOG_MUL_FIX]22 23    // first cancel the adjustment from llama_hparams::yarn_attn_factor_adjust to get the original attn_factor24    GGML_ASSERT(ext_factor >= 0.0f);25    const float attn_factor_org = attn_factor * (1.0f + 0.1f * logf(1.0f / freq_scale));26 27    // use the original attn_factor to pre-scale the kq_scale28    const float mscale   = attn_factor_org * (1.0f + 0.1f * hparams.rope_yarn_log_mul * logf(1.0f / freq_scale));29    const float kq_scale = 1.0f * mscale * mscale / sqrtf(float(n_embd_head_k));30 31    ggml_tensor * cur;32    ggml_tensor * inpL;33 34    // {n_embd, n_tokens}35    inpL = build_inp_embd(model.tok_embd);36 37    // (optional) temperature tuning - used by mistral-large38    ggml_tensor * inp_attn_scale = nullptr;39    if (hparams.f_attn_temp_scale != 0.0f) {40        inp_attn_scale = build_inp_attn_scale();41    }42 43    // inp_pos - contains the positions44    ggml_tensor * inp_pos = build_inp_pos();45 46    auto * inp_attn_kv = !is_mla ? build_attn_inp_kv() : nullptr;47    auto * inp_attn_k  =  is_mla ? build_attn_inp_k()  : nullptr;48 49    ggml_tensor * inp_out_ids = build_inp_out_ids();50 51    int effective_n_layers = hparams.n_layer - hparams.nextn_predict_layers;52    for (int il = 0; il < effective_n_layers; ++il) {53        ggml_tensor * inpSA = inpL;54 55        // norm56        cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il);57        cb(cur, "attn_norm", il);58 59        // self_attention60        if (is_ocr) {61            const int n_embed_head = hparams.n_embd / hparams.n_head();62            const int ocr_rope_type = GGML_ROPE_TYPE_NEOX;63            GGML_ASSERT(n_embed_head == n_embd_head_k && n_embed_head == n_embd_head_v);64 65            ggml_tensor * Qcur = NULL;66            ggml_tensor * Kcur = NULL;67            ggml_tensor * Vcur = NULL;68 69            Qcur = ggml_mul_mat(ctx0, model.layers[il].wq, cur);70            Kcur = ggml_mul_mat(ctx0, model.layers[il].wk, cur);71            Vcur = ggml_mul_mat(ctx0, model.layers[il].wv, cur);72            cb(Qcur, "q", il);73            cb(Kcur, "k", il);74            cb(Vcur, "v", il);75 76            Qcur = ggml_reshape_3d(ctx0, Qcur, n_embed_head, n_head, n_tokens);77            Kcur = ggml_reshape_3d(ctx0, Kcur, n_embed_head, n_head, n_tokens);78            Vcur = ggml_reshape_3d(ctx0, Vcur, n_embed_head, n_head, n_tokens);79 80            GGML_ASSERT(fabs(freq_base - 10000.0) < 1e-4);81            Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_embed_head, ocr_rope_type, 0, freq_base, 1, 0, 1, 0, 0);82            Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_embed_head, ocr_rope_type, 0, freq_base, 1, 0, 1, 0, 0);83            cb(Qcur, "q_pe", il);84            cb(Kcur, "k_pe", il);85 86            cur = build_attn(inp_attn_kv,87                        model.layers[il].wo, NULL, model.layers[il].wo_s,88                        Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);89            cb(cur, "attn_out", il);90        }91        else {92            ggml_tensor * q = NULL;93 94            const bool is_lite = model.layers[il].wq;95 96            if (!is_lite) {97                q = ggml_mul_mat(ctx0, model.layers[il].wq_a, cur);98                cb(q, "q", il);99 100                q = build_norm(q, model.layers[il].attn_q_a_norm, nullptr, LLM_NORM_RMS, il);101                cb(q, "q", il);102 103                q = ggml_mul_mat(ctx0, model.layers[il].wq_b, q);104                cb(q, "q", il);105            } else {106                q = ggml_mul_mat(ctx0, model.layers[il].wq, cur);107                cb(q, "q", il);108            }109            // split into {n_embd_head_qk_nope, n_head, n_tokens}110            ggml_tensor * q_nope =111                ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k),112                             ggml_row_size(q->type, n_embd_head_k) * n_head, 0);113            cb(q_nope, "q_nope", il);114 115            // and {n_embd_head_qk_rope, n_head, n_tokens}116            ggml_tensor * q_pe = ggml_view_3d(117                ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k),118                ggml_row_size(q->type, n_embd_head_k) * n_head, ggml_row_size(q->type, n_embd_head_qk_nope));119            cb(q_pe, "q_pe", il);120 121            ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, model.layers[il].wkv_a_mqa, cur);122            cb(kv_cmpr_pe, "kv_cmpr_pe", il);123 124            // split into {kv_lora_rank, n_tokens}125            ggml_tensor * kv_cmpr =126                ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,127                             ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0);128            cb(kv_cmpr, "kv_cmpr", il);129 130            // and {n_embd_head_qk_rope, 1, n_tokens}131            ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens,132                                              ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),133                                              ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope),134                                              ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));135            cb(k_pe, "k_pe", il);136 137            q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,138                                 ext_factor, attn_factor, beta_fast, beta_slow);139            cb(q_pe, "q_pe", il);140 141            k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,142                                 ext_factor, attn_factor, beta_fast, beta_slow);143            cb(k_pe, "k_pe", il);144 145            kv_cmpr = build_norm(kv_cmpr, model.layers[il].attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);146            cb(kv_cmpr, "kv_cmpr", il);147 148            if (is_mla) {149                // {n_embd_head_qk_nope, n_tokens, n_head}150                q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);151                cb(q_nope, "q_nope_perm", il);152 153                // {n_embd_head_qk_nope, kv_lora_rank, n_head} x {n_embd_head_qk_nope, n_tokens, n_head}154                ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, model.layers[il].wk_b, q_nope);155                cb(q_nope_absorbed, "q_nope_absorbed", il);156 157                // {kv_lora_rank, n_head, n_tokens}158                q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);159                cb(q_nope_absorbed, "q_nope_absorbed_perm", il);160 161                // {n_embd_head_qk_rope + kv_lora_rank, n_head, n_tokens}162                // note: rope must go first for in-place context shifting in build_rope_shift()163                ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);164                cb(Qcur, "Qcur", il);165 166                kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);167                cb(kv_cmpr, "kv_cmpr_reshape", il);168 169                // {n_embd_head_qk_rope + kv_lora_rank, 1, n_tokens}170                ggml_tensor * Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0);171                cb(Kcur, "Kcur", il);172 173                // {kv_lora_rank, 1, n_tokens}174                ggml_tensor * Vcur = kv_cmpr;175                cb(Vcur, "Vcur", il);176 177                if (inp_attn_scale) {178                    // apply llama 4 temperature scaling179                    Qcur = ggml_mul(ctx0, Qcur, inp_attn_scale);180                    cb(Qcur, "Qcur_attn_temp_scaled", il);181                }182 183                // note: MLA with the absorption optimization converts into MQA (ie: GQA with 1 group)184                cur = build_attn(inp_attn_k,185                        model.layers[il].wo, NULL, model.layers[il].wo_s,186                        Qcur, Kcur, Vcur, nullptr, nullptr, model.layers[il].wv_b, kq_scale, il);187            } else {188                ggml_tensor * kv = ggml_mul_mat(ctx0, model.layers[il].wkv_b, kv_cmpr);189                cb(kv, "kv", il);190 191                // split into {n_embd_head_qk_nope, n_head, n_tokens}192                ggml_tensor * k_nope =193                    ggml_view_3d(ctx0, kv, n_embd_head_qk_nope, n_head, n_tokens,194                                 ggml_row_size(kv->type, n_embd_head_qk_nope + n_embd_head_v),195                                 ggml_row_size(kv->type, n_embd_head_qk_nope + n_embd_head_v) * n_head, 0);196                cb(k_nope, "k_nope_view", il);197 198                // and {n_embd_head_v, n_head, n_tokens}199                ggml_tensor * Vcur = ggml_view_3d(ctx0, kv, n_embd_head_v, n_head, n_tokens,200                                                  ggml_row_size(kv->type, n_embd_head_qk_nope + n_embd_head_v),201                                                  ggml_row_size(kv->type, n_embd_head_qk_nope + n_embd_head_v) * n_head,202                                                  ggml_row_size(kv->type, n_embd_head_qk_nope));203                cb(Vcur, "Vcur_view", il);204 205                Vcur = ggml_cont(ctx0, Vcur);206                cb(Vcur, "Vcur_cont", il);207 208                ggml_tensor * Qcur = ggml_concat(ctx0, q_nope, q_pe, 0);209                cb(Qcur, "Qcur", il);210 211                ggml_tensor * Kcur = ggml_concat(ctx0, k_nope, ggml_repeat(ctx0, k_pe, q_pe), 0);212                cb(Kcur, "Kcur", il);213 214                if (inp_attn_scale) {215                    // apply llama 4 temperature scaling216                    Qcur = ggml_mul(ctx0, Qcur, inp_attn_scale);217                    cb(Qcur, "Qcur_attn_temp_scaled", il);218                }219 220                // note: MLA without the absorption optimization converts into MHA (ie: GQA with full n_head groups)221                cur = build_attn(inp_attn_kv,222                            model.layers[il].wo, NULL, model.layers[il].wo_s,223                            Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);224            }225        }226        if (il == effective_n_layers - 1 && inp_out_ids) {227            cur   = ggml_get_rows(ctx0, cur, inp_out_ids);228            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);229        }230        ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);231        cb(ffn_inp, "ffn_inp", il);232 233        cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il);234        cb(cur, "ffn_norm", il);235 236        if ((uint32_t) il < hparams.n_layer_dense_lead) {237            cur = build_ffn(cur,238                model.layers[il].ffn_up, NULL, NULL,239                model.layers[il].ffn_gate, NULL, NULL,240                model.layers[il].ffn_down, NULL, NULL,241                NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);242            cb(cur, "ffn_out", il);243        } else {244            // MoE branch245            ggml_tensor * moe_out = build_moe_ffn(cur,246                model.layers[il].ffn_gate_inp,247                model.layers[il].ffn_up_exps,248                model.layers[il].ffn_gate_exps,249                model.layers[il].ffn_down_exps,250                model.layers[il].ffn_exp_probs_b,251                n_expert, n_expert_used,252                LLM_FFN_SILU, hparams.expert_weights_norm,253                hparams.expert_weights_scale,254                (llama_expert_gating_func_type) hparams.expert_gating_func,255                il,256                nullptr,257                model.layers[il].ffn_gate_up_exps);258            cb(moe_out, "ffn_moe_out", il);259 260            // FFN shared expert261            {262                ggml_tensor * ffn_shexp =263                    build_ffn(cur,264                        model.layers[il].ffn_up_shexp, NULL, NULL,265                        model.layers[il].ffn_gate_shexp, NULL, NULL,266                        model.layers[il].ffn_down_shexp, NULL, NULL,267                        NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);268                cb(ffn_shexp, "ffn_shexp", il);269 270                cur = ggml_add(ctx0, moe_out, ffn_shexp);271                cb(cur, "ffn_out", il);272            }273        }274        cur = ggml_add(ctx0, cur, ffn_inp);275 276        cur = build_cvec(cur, il);277        cb(cur, "l_out", il);278 279        // input for next layer280        inpL = cur;281    }282    cur = inpL;283 284    cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1);285 286    cb(cur, "result_norm", -1);287    res->t_embd = cur;288 289    // lm_head290    cur = ggml_mul_mat(ctx0, model.output, cur);291 292    cb(cur, "result_output", -1);293    res->t_logits = cur;294 295    ggml_build_forward_expand(gf, cur);296}297