Felipe97/llama-cpp-compiled

1#include "models.h"2 3std::unique_ptr<llm_graph_context> llama_model_nemotron_h_moe::build_arch_graph(const llm_graph_params & params) const {4    if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {5        return std::make_unique<graph_mtp>(*this, params);6    }7    return std::make_unique<graph>(*this, params);8}9 10// MTP draft head for Nemotron-H MoE11llama_model_nemotron_h_moe::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params)12    : llm_graph_context(params) {13    GGML_ASSERT(hparams.n_layer_nextn == 1 && "NEMOTRON_H_MOE MTP currently supports a single MTP block");14 15    const int64_t n_embd_head = hparams.n_embd_head_v();16    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());17 18    const int il = hparams.n_layer();19    const auto & layer = model.layers[il];20 21    GGML_ASSERT(layer.nextn.eh_proj && layer.nextn.enorm && layer.nextn.hnorm);22    GGML_ASSERT(layer.ffn_gate_inp);23 24    // token embedding weights25    ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd;26    GGML_ASSERT(tok_embd_w != nullptr && "NEMOTRON_H_MOE MTP requires token embeddings");27 28    auto inp = std::make_unique<llm_graph_input_embd_h>(hparams.n_embd);29 30    inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);31    ggml_set_input(inp->tokens);32 33    inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens);34    ggml_set_input(inp->embd);35 36    ggml_tensor * tok_embd;37    if (ubatch.token) {38        tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens);39    } else {40        tok_embd = inp->embd;41    }42    cb(tok_embd, "mtp_tok_embd", il);43 44    inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);45    ggml_set_input(inp->h);46    ggml_set_name(inp->h, "mtp_h_input");47 48    ggml_tensor * h_embd = inp->h;49 50    res->add_input(std::move(inp));51 52    ggml_tensor * inp_out_ids = build_inp_out_ids();53 54    // attention fills KV over all tokens, but the MoE is position-wise: gather output rows before55    // it to save FFN compute (unless unmasked embeddings_nextn needs the full-length hidden state)56    const bool emit_h_nextn    = cparams.embeddings_nextn;57    const bool crop_before_ffn = inp_out_ids && (!emit_h_nextn || cparams.embeddings_nextn_masked);58 59    auto * inp_attn = build_attn_inp_kv();60 61    ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);62    cb(h_norm, "mtp_hnorm", il);63    ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);64    cb(e_norm, "mtp_enorm", il);65 66    ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0);67    cb(concat, "mtp_concat", il);68 69    ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s);70    cb(cur, "mtp_eh_proj", il);71 72    // dense NoPE attention sub-layer (mtp.layers.0)73    ggml_tensor * inpSA = cur;74    cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);75    cb(cur, "mtp_attn_norm", il);76 77    {78        auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il);79        const float kq_scale = hparams.f_attention_scale == 0.0f80            ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale;81        cur = build_attn(inp_attn, layer.wo, layer.wo_b, layer.wo_s,82                Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);83        cb(cur, "mtp_attn_out", il);84    }85 86    cur = ggml_add(ctx0, cur, inpSA);87    cb(cur, "mtp_attn_residual", il);88 89    // gather the output rows here so the MoE FFN below only runs on the positions we keep90    if (crop_before_ffn) {91        cur = ggml_get_rows(ctx0, cur, inp_out_ids);92    }93 94    // MoE FFN sub-layer (mtp.layers.1)95    ggml_tensor * ffn_residual = cur;96    cur = build_norm(cur, layer.attn_post_norm, nullptr, LLM_NORM_RMS, il);97    cb(cur, "mtp_attn_post_norm", il);98 99    {100        ggml_tensor * router_logits = build_lora_mm(layer.ffn_gate_inp, cur);101        cb(router_logits, "mtp_ffn_moe_logits", il);102 103        ggml_tensor * ffn_shexp = build_ffn(cur,104                layer.ffn_up_shexp,   NULL, layer.ffn_up_shexp_s,105                NULL,                 NULL, NULL,106                layer.ffn_down_shexp, NULL, layer.ffn_down_shexp_s,107                NULL,108                LLM_FFN_RELU_SQR, LLM_FFN_PAR, il);109        cb(ffn_shexp, "mtp_ffn_shexp", il);110 111        if (layer.ffn_latent_down) {112            cur = ggml_mul_mat(ctx0, layer.ffn_latent_down, cur);113        }114 115        ggml_tensor * moe_out =116            build_moe_ffn(cur,117                layer.ffn_gate_inp,118                layer.ffn_up_exps,119                nullptr, // no gate120                layer.ffn_down_exps,121                layer.ffn_exp_probs_b,122                n_expert, n_expert_used,123                LLM_FFN_RELU_SQR, hparams.expert_weights_norm,124                hparams.expert_weights_scale,125                LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID,126                il,127                router_logits, nullptr,128                layer.ffn_up_exps_s,129                nullptr, // no gate130                layer.ffn_down_exps_s);131        cb(moe_out, "mtp_ffn_moe_out", il);132 133        if (layer.ffn_latent_up) {134            moe_out = ggml_mul_mat(ctx0, layer.ffn_latent_up, moe_out);135        }136 137        cur = ggml_add(ctx0, moe_out, ffn_shexp);138        cb(cur, "mtp_ffn_out", il);139    }140 141    cur = ggml_add(ctx0, cur, ffn_residual);142    cb(cur, "mtp_post_ffn", il);143 144    // final head norm: the MTP head has its own LayerNorm145    GGML_ASSERT(layer.nextn.shared_head_norm && "NEMOTRON_H_MOE MTP: missing final head norm");146    cur = build_norm(cur, layer.nextn.shared_head_norm, nullptr, LLM_NORM, -1);147 148    cb(cur, "h_nextn", -1);149    res->t_h_nextn = cur;150 151    if (!crop_before_ffn && inp_out_ids) {152        cur = ggml_get_rows(ctx0, cur, inp_out_ids);153    }154 155    // LM head156    ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output;157    ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s;158    GGML_ASSERT(head_w != nullptr && "NEMOTRON_H_MOE MTP requires an output projection");159    cur = build_lora_mm(head_w, cur, head_s);160    cb(cur, "result_output", -1);161 162    res->t_logits = cur;163    ggml_build_forward_expand(gf, cur);164}165