Felipe97/llama-cpp-compiled

1#include "models.h"2 3void llama_model_cohere2moe::load_arch_hparams(llama_model_loader & ml) {4    const bool found_norm     = ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS,     hparams.f_norm_eps,     false);5    const bool found_norm_rms = ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps, false);6    if (!found_norm && !found_norm_rms) {7        throw std::runtime_error("missing Cohere2 MoE norm epsilon");8    }9    if (!found_norm_rms) {10        hparams.f_norm_rms_eps = 0.0f;11    }12 13    ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW,    hparams.n_swa);14    ml.get_key(LLM_KV_LOGIT_SCALE,                 hparams.f_logit_scale);15    ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT,   hparams.n_layer_dense_lead);16    ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);17    ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);18    ml.get_key(LLM_KV_EXPERT_SHARED_COUNT,         hparams.n_expert_shared, false);19    ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM,         hparams.expert_weights_norm, false);20    ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE,        hparams.expert_weights_scale, false);21    ml.get_key(LLM_KV_EXPERT_GATING_FUNC,          hparams.expert_gating_func, false);22 23    if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) {24        hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID;25    }26 27    hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;28    load_swa_pattern(ml, 4, true);29 30    hparams.rope_freq_base_train_swa  = hparams.rope_freq_base_train;31    hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;32    ml.get_key(LLM_KV_ROPE_FREQ_BASE_SWA, hparams.rope_freq_base_train_swa, false);33 34    switch (hparams.n_layer()) {35        case 49: type = LLM_TYPE_30B_A3B; break;36        default: type = LLM_TYPE_UNKNOWN;37    }38}39 40void llama_model_cohere2moe::load_arch_tensors(llama_model_loader & ml) {41    LLAMA_LOAD_LOCALS;42 43    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);44    // Trunk-only: the GGUF declares MTP layers in metadata but the actual MTP45    // tensors live in a separate file. Mark MTP tensors NOT_REQUIRED so the46    // trunk loads cleanly.47    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";48    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);49    const int trunk_flags = mtp_only  ? TENSOR_NOT_REQUIRED : 0;50    int mtp_flags         = trunk_only ? TENSOR_NOT_REQUIRED : 0;51 52    if (!ml.load_mtp) {53        mtp_flags |= TENSOR_SKIP;54    }55 56    tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);57 58    // output59    output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0);60    output      = create_tensor(tn(LLM_TENSOR_OUTPUT,      "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);61 62    // if output is NULL, init from the input tok embed63    if (output == NULL) {64        output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED);65    }66 67    if (n_expert == 0) {68        throw std::runtime_error("n_expert must be > 0 for Cohere2Moe");69    }70    if (n_expert_used == 0) {71        throw std::runtime_error("n_expert_used must be > 0 for Cohere2Moe");72    }73 74    auto load_block_trunk = [&](int i, int flags) {75        auto & layer = layers[i];76 77        layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, flags);78 79        create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_gqa, n_embd_gqa, flags);80        layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, flags);81 82        if (static_cast<uint32_t>(i) < hparams.n_layer_dense_lead) {83            layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, flags);84            layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, flags);85            layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), { n_embd, n_ff }, flags);86        } else {87            const int64_t n_ff_exp = hparams.n_ff_exp() ? hparams.n_ff_exp() : n_ff;88 89            layer.ffn_gate_inp  = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,  "weight", i), { n_embd, n_expert }, flags);90            layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), { n_ff_exp, n_embd, n_expert }, flags);91            create_tensor_gate_up_exps(layer, i, n_embd, n_ff_exp, n_expert, flags);92 93            if (hparams.n_expert_shared > 0) {94                const int64_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff_exp * hparams.n_expert_shared;95                layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), { n_embd, n_ff_shexp }, flags);96                layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_shexp, n_embd }, flags);97                layer.ffn_up_shexp   = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,   "weight", i), { n_embd, n_ff_shexp }, flags);98            }99        }100    };101 102    auto load_block_mtp = [&](int i, int flags) {103        auto & layer = layers[i];104 105        // MTP block looks like a full-attention Cohere2 MoE decoder block.106        layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, flags);107 108        create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_gqa, n_embd_gqa, flags);109        layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, flags);110 111        const int64_t n_ff_exp = hparams.n_ff_exp() ? hparams.n_ff_exp() : n_ff;112 113        // Routed experts114        layer.ffn_gate_inp  = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,  "weight", i), { n_embd, n_expert }, flags);115        layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), { n_ff_exp, n_embd, n_expert }, flags);116        create_tensor_gate_up_exps(layer, i, n_embd, n_ff_exp, n_expert, flags);117 118        if (hparams.n_expert_shared > 0) {119            const int64_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff_exp * hparams.n_expert_shared;120 121            // Shared experts122            layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), { n_embd, n_ff_shexp }, flags);123            layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_shexp, n_embd }, flags);124            layer.ffn_up_shexp   = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,   "weight", i), { n_embd, n_ff_shexp }, flags);125        }126 127        // NextN-specific tensors that define the MTP block.128        layer.nextn.eh_proj          = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ,          "weight", i), { 2 * n_embd, n_embd }, flags);129        layer.nextn.enorm            = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM,            "weight", i), { n_embd },              flags);130        layer.nextn.hnorm            = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM,            "weight", i), { n_embd },              flags);131        layer.nextn.embed_tokens     = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS,     "weight", i), { n_embd, n_vocab },     TENSOR_NOT_REQUIRED);132        layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), { n_embd, n_vocab },     TENSOR_NOT_REQUIRED);133        layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), { n_embd },              TENSOR_NOT_REQUIRED);134    };135 136    for (int i = 0; i < n_layer; ++i) {137        load_block_trunk(i, trunk_flags);138    }139    // MTP/NextN layers are loaded as extra decoder blocks.140    for (int i = n_layer; i < n_layer_all; ++i) {141        load_block_mtp(i, mtp_flags);142    }143}144 145std::unique_ptr<llm_graph_context> llama_model_cohere2moe::build_arch_graph(const llm_graph_params & params) const {146    if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {147        return std::make_unique<graph_mtp>(*this, params);148    }149    return std::make_unique<graph>(*this, params);150}151 152llama_model_cohere2moe::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {153    const int64_t n_embd_head = hparams.n_embd_head_v();154 155    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());156    GGML_ASSERT(n_embd_head == n_rot);157 158    const llm_norm_type cohere2moe_norm_type = hparams.f_norm_rms_eps == 0.0f ? LLM_NORM : LLM_NORM_RMS;159    const float f_logit_scale = hparams.f_logit_scale;160    ggml_tensor * cur;161    ggml_tensor * inpL = build_inp_embd(model.tok_embd);162    ggml_tensor * inp_pos = build_inp_pos();163 164    auto * inp_attn = build_attn_inp_kv_iswa();165    ggml_tensor * inp_out_ids = build_inp_out_ids();166 167    // MTP/NextN layers are loaded as extra decoder blocks but not executed in the main pass.168    for (int il = 0; il < n_layer; ++il) {169        const bool is_swa = hparams.is_swa(il);170        // Dense-prefix full-attention layers use RoPE; later layers follow the SWA pattern.171        const bool force_rope = static_cast<uint32_t>(il) < hparams.n_layer_dense_lead;172 173        cur = build_norm(inpL, model.layers[il].attn_norm, nullptr, cohere2moe_norm_type, il);174        cb(cur, "attn_norm", il);175 176        ggml_tensor * ffn_inp = cur;177 178        {179            const auto & layer = model.layers[il];180 181            auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur,182                    n_embd_head, n_head, n_head_kv, il);183 184            if (is_swa || force_rope) {185                ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);186 187                Qcur = ggml_rope_ext(188                        ctx0, Qcur, inp_pos, rope_factors,189                        n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,190                        ext_factor, attn_factor, beta_fast, beta_slow);191 192                Kcur = ggml_rope_ext(193                        ctx0, Kcur, inp_pos, rope_factors,194                        n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,195                        ext_factor, attn_factor, beta_fast, beta_slow);196            }197 198            cb(Qcur, "Qcur", il);199            cb(Kcur, "Kcur", il);200            cb(Vcur, "Vcur", il);201 202            cur = build_attn(inp_attn,203                    layer.wo, layer.wo_b, layer.wo_s,204                    Qcur, Kcur, Vcur, nullptr, nullptr, nullptr,205                    1.0f / sqrtf(float(n_embd_head)), il);206        }207 208        if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {209            cur     = ggml_get_rows(ctx0, cur, inp_out_ids);210            inpL    = ggml_get_rows(ctx0, inpL, inp_out_ids);211            ffn_inp = ggml_get_rows(ctx0, ffn_inp, inp_out_ids);212        }213 214        ggml_tensor * attn_out = cur;215 216        const auto & layer = model.layers[il];217 218        if (layer.ffn_gate_inp == nullptr) {219            cur = build_ffn(ffn_inp,220                    layer.ffn_up,   nullptr, layer.ffn_up_s,221                    layer.ffn_gate, nullptr, layer.ffn_gate_s,222                    layer.ffn_down, nullptr, layer.ffn_down_s,223                    nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);224            cb(cur, "ffn_out", il);225        } else {226            cur = build_moe_ffn(ffn_inp,227                    layer.ffn_gate_inp,228                    layer.ffn_up_exps,229                    layer.ffn_gate_exps,230                    layer.ffn_down_exps,231                    nullptr,232                    n_expert, n_expert_used,233                    LLM_FFN_SILU, hparams.expert_weights_norm,234                    hparams.expert_weights_scale,235                    (llama_expert_gating_func_type) hparams.expert_gating_func,236                    il,237                    nullptr, layer.ffn_gate_up_exps,238                    layer.ffn_up_exps_s,239                    layer.ffn_gate_exps_s,240                    layer.ffn_down_exps_s);241            cb(cur, "ffn_moe_out", il);242 243            if (layer.ffn_up_shexp) {244                ggml_tensor * ffn_shexp = build_ffn(ffn_inp,245                        layer.ffn_up_shexp,   nullptr, layer.ffn_up_shexp_s,246                        layer.ffn_gate_shexp, nullptr, layer.ffn_gate_shexp_s,247                        layer.ffn_down_shexp, nullptr, layer.ffn_down_shexp_s,248                        nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);249                cb(ffn_shexp, "ffn_shexp", il);250 251                cur = ggml_add(ctx0, cur, ffn_shexp);252                cur = ggml_scale(ctx0, cur, 0.5f);253                cb(cur, "ffn_out", il);254            }255        }256 257        cur = ggml_add(ctx0, cur, inpL);258        cur = ggml_add(ctx0, cur, attn_out);259 260        cur = build_cvec(cur, il);261        cb(cur, "l_out", il);262 263        inpL = cur;264    }265 266    cur = inpL;267    cur = build_norm(cur, model.output_norm, nullptr, cohere2moe_norm_type, -1);268 269    cb(cur, "h_nextn", -1);270    res->t_h_nextn = cur;271 272    if (!cparams.embeddings_nextn_masked && inp_out_ids) {273        cur = ggml_get_rows(ctx0, cur, inp_out_ids);274    }275 276    cb(cur, "result_norm", -1);277    res->t_embd = cur;278 279    cur = build_lora_mm(model.output, cur);280 281    if (f_logit_scale) {282        cur = ggml_scale(ctx0, cur, f_logit_scale);283    }284 285    cb(cur, "result_output", -1);286    res->t_logits = cur;287 288    ggml_build_forward_expand(gf, cur);289}290 291llama_model_cohere2moe::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {292    GGML_ASSERT(hparams.n_layer_nextn > 0 && "COHERE2MOE MTP requires n_layer_nextn > 0");293    GGML_ASSERT(hparams.n_layer_nextn == 1 && "COHERE2MOE MTP currently only supports a single MTP block");294 295    const int64_t n_embd_head = hparams.n_embd_head_v();296    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());297    GGML_ASSERT(n_embd_head == n_rot);298 299    const int il = hparams.n_layer();300    const auto & layer = model.layers[il];301    GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");302    GGML_ASSERT(layer.nextn.enorm   && "MTP block missing nextn.enorm");303    GGML_ASSERT(layer.nextn.hnorm   && "MTP block missing nextn.hnorm");304    GGML_ASSERT(layer.ffn_gate_inp  && "MTP block missing ffn_gate_inp");305 306    const llm_norm_type cohere2moe_norm_type = hparams.f_norm_rms_eps == 0.0f ? LLM_NORM : LLM_NORM_RMS;307 308    // TODO: extract in a common llm_graph_context::build_inp_embd_h()309    auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd);310 311    inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);312    ggml_set_input(inp->tokens);313 314    inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);315    ggml_set_input(inp->embd);316 317    // TODO: make static using `ggml_build_forward_select()`318    //       see llm_graph_context::build_inp_embd() for reference319    ggml_tensor * tok_embd;320    if (ubatch.token) {321        ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;322        tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);323    } else {324        tok_embd = inp->embd;325    }326    cb(tok_embd, "mtp_tok_embd", il);327 328    inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);329    ggml_set_input(inp->h);330    ggml_set_name(inp->h, "mtp_h_input");331 332    ggml_tensor * h_embd = inp->h;333 334    res->add_input(std::move(inp));335 336    ggml_tensor * inp_pos     = build_inp_pos();337    ggml_tensor * inp_out_ids = build_inp_out_ids();338    auto * inp_attn = build_attn_inp_kv_iswa();339 340    ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, cohere2moe_norm_type, il);341    cb(h_norm, "mtp_hnorm", il);342 343    ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, cohere2moe_norm_type, il);344    cb(e_norm, "mtp_enorm", il);345 346    ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);347    cb(concat, "mtp_concat", il);348 349    ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);350    cb(cur, "mtp_eh_proj", il);351 352    ggml_tensor * inpL = cur;353 354    cur = build_norm(cur, layer.attn_norm, nullptr, cohere2moe_norm_type, il);355    cb(cur, "mtp_attn_norm", il);356    ggml_tensor * ffn_inp = cur;357 358    auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, n_head, n_head_kv, il);359    ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);360    Qcur = ggml_rope_ext(361            ctx0, Qcur, inp_pos, rope_factors,362            n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,363            ext_factor, attn_factor, beta_fast, beta_slow);364    Kcur = ggml_rope_ext(365            ctx0, Kcur, inp_pos, rope_factors,366            n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,367            ext_factor, attn_factor, beta_fast, beta_slow);368 369    cb(Qcur, "mtp_Qcur", il);370    cb(Kcur, "mtp_Kcur", il);371    cb(Vcur, "mtp_Vcur", il);372 373    cur = build_attn(inp_attn,374            layer.wo, layer.wo_b, layer.wo_s,375            Qcur, Kcur, Vcur, nullptr, nullptr, nullptr,376            1.0f / sqrtf(float(n_embd_head)), il);377    cb(cur, "mtp_attn_out", il);378 379    ggml_tensor * attn_out = cur;380 381    cur = build_moe_ffn(ffn_inp,382            layer.ffn_gate_inp,383            layer.ffn_up_exps,384            layer.ffn_gate_exps,385            layer.ffn_down_exps,386            nullptr,387            n_expert, n_expert_used,388            LLM_FFN_SILU, hparams.expert_weights_norm,389            hparams.expert_weights_scale,390            (llama_expert_gating_func_type) hparams.expert_gating_func,391            il,392            nullptr, layer.ffn_gate_up_exps,393            layer.ffn_up_exps_s,394            layer.ffn_gate_exps_s,395            layer.ffn_down_exps_s);396    cb(cur, "mtp_ffn_moe_out", il);397 398    if (layer.ffn_up_shexp) {399        ggml_tensor * ffn_shexp = build_ffn(ffn_inp,400                layer.ffn_up_shexp,   nullptr, layer.ffn_up_shexp_s,401                layer.ffn_gate_shexp, nullptr, layer.ffn_gate_shexp_s,402                layer.ffn_down_shexp, nullptr, layer.ffn_down_shexp_s,403                nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);404        cb(ffn_shexp, "mtp_ffn_shexp", il);405 406        cur = ggml_add(ctx0, cur, ffn_shexp);407        cur = ggml_scale(ctx0, cur, 0.5f);408        cb(cur, "mtp_ffn_out", il);409    }410 411    cur = ggml_add(ctx0, cur, inpL);412    cur = ggml_add(ctx0, cur, attn_out);413    cb(cur, "mtp_post_ffn", il);414 415    ggml_tensor * head_norm_w = layer.nextn.shared_head_norm416            ? layer.nextn.shared_head_norm417            : model.output_norm;418    GGML_ASSERT(head_norm_w && "COHERE2MOE MTP: missing both nextn.shared_head_norm and output_norm");419    cur = build_norm(cur, head_norm_w, nullptr, cohere2moe_norm_type, -1);420 421    cb(cur, "h_nextn", -1);422    res->t_h_nextn = cur;423 424    cur = ggml_get_rows(ctx0, cur, inp_out_ids);425    cb(cur, "mtp_shared_head_norm", -1);426 427    ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;428    GGML_ASSERT(head_w && "COHERE2MOE MTP: missing LM head (nextn.shared_head_head or model.output)");429    cur = build_lora_mm(head_w, cur, layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : nullptr);430 431    if (hparams.f_logit_scale) {432        cur = ggml_scale(ctx0, cur, hparams.f_logit_scale);433    }434 435    cb(cur, "result_output", -1);436    res->t_logits = cur;437 438    ggml_build_forward_expand(gf, cur);439}440