Felipe97/llama-cpp-compiled

1#include "models.h"2 3#include <sstream>4 5void llama_model_granite::load_arch_hparams(llama_model_loader & ml) {6    ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);7    ml.get_key(LLM_KV_LOGIT_SCALE,                 hparams.f_logit_scale);8    ml.get_key(LLM_KV_RESIDUAL_SCALE,              hparams.f_residual_scale, false);9    ml.get_key(LLM_KV_EMBEDDING_SCALE,             hparams.f_embedding_scale, false);10    ml.get_key(LLM_KV_ATTENTION_SCALE,             hparams.f_attention_scale, false);11 12    // Granite4 Vision uses array deepstack_mapping13    ml.get_arr(LLM_KV_DEEPSTACK_MAPPING, hparams.deepstack_mapping_arr, false);14 15    // Count the unique deepstack input indices16    std::unordered_set<uint32_t> unique_deepstack_idxs;17    for (const auto val : hparams.deepstack_mapping_arr) {18        if (val >= 0) {19            unique_deepstack_idxs.insert(val);20        }21    }22    hparams.n_deepstack_layers = unique_deepstack_idxs.size();23 24    // Ensure all values are valid (avoid overflow attacks)25    for (const auto val : unique_deepstack_idxs) {26        if (val > hparams.n_deepstack_layers) {27            std::stringstream ss;28            ss << "Invalid deepstack index: " << val << " > " << hparams.n_deepstack_layers;29            throw std::runtime_error(ss.str());30        }31    }32 33    // Granite uses rope_finetuned as a switch for rope, so default to true34    bool rope_finetuned = true;35    ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false);36    hparams.rope_finetuned = rope_finetuned; // needed for round trip save37    std::fill(hparams.rope_pattern.begin(), hparams.rope_pattern.end(), rope_finetuned);38 39    switch (hparams.n_layer()) {40        case 32: type = LLM_TYPE_3B; break;41        case 40: {42            switch (hparams.n_embd) {43                case 2048: type = LLM_TYPE_2B; break;44                case 2560: type = LLM_TYPE_3B; break;45                case 4096: type = LLM_TYPE_8B; break;46                default: type = LLM_TYPE_UNKNOWN;47            }48            break;49        }50        case 64: type = LLM_TYPE_30B; break;51        // Add additional layer/vocab/etc checks here for other model sizes52        default: type = LLM_TYPE_UNKNOWN;53    }54 55    // For Granite MoE Shared56    ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, /* required */ false);57}58 59void llama_model_granite::load_arch_tensors(llama_model_loader &) {60    LLAMA_LOAD_LOCALS;61 62    tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);63 64    // output65    output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);66    output      = create_tensor(tn(LLM_TENSOR_OUTPUT,      "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);67 68    // if output is NULL, init from the input tok embed69    if (output == NULL) {70        output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);71    }72 73    for (int i = 0; i < n_layer; ++i) {74        auto & layer = layers[i];75 76        layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);77 78        create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0);79        layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0);80 81        // optional bias tensors82        layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED);83 84        layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);85 86        if (hparams.rope_scaling_type_train == LLAMA_ROPE_SCALING_TYPE_LONGROPE) {87            layer.rope_long  = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_LONG,  "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));88            layer.rope_short = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_SHORT, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));89        }90        else {91            layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0));92        }93 94        if (n_expert == 0) {95            layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd,   n_ff}, 0);96            layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {  n_ff, n_embd}, 0);97            layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd,   n_ff}, 0);98 99            // optional MLP bias100            layer.ffn_gate_b = create_tensor(tn(LLM_TENSOR_FFN_GATE, "bias", i), {n_ff}, TENSOR_NOT_REQUIRED);101            layer.ffn_down_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED);102            layer.ffn_up_b   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "bias", i), {n_ff}, TENSOR_NOT_REQUIRED);103        } else {104            layer.ffn_gate_inp  = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP,  "weight", i), {n_embd, n_expert}, 0);105            layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd,   n_ff, n_expert}, TENSOR_NOT_REQUIRED);106            layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {  n_ff, n_embd, n_expert}, 0);107            layer.ffn_up_exps   = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS,   "weight", i), {n_embd,   n_ff, n_expert}, 0);108 109            // For Granite MoE Shared110            if (hparams.n_ff_shexp > 0) {111                layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, hparams.n_ff_shexp}, 0);112                layer.ffn_up_shexp   = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP,   "weight", i), {n_embd, hparams.n_ff_shexp}, 0);113                layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {hparams.n_ff_shexp, n_embd}, 0);114            }115        }116    }117}118 119std::unique_ptr<llm_graph_context> llama_model_granite::build_arch_graph(const llm_graph_params & params) const {120    return std::make_unique<graph>(*this, params);121}122 123llama_model_granite::graph::graph(124    const llama_model & model,125    const llm_graph_params & params)126    : llm_graph_context(params) {127 128    const int64_t n_embd_head = hparams.n_embd_head_v();129 130    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());131    GGML_ASSERT(n_embd_head == n_rot);132 133    ggml_tensor * cur;134    ggml_tensor * inpL;135 136    inpL = build_inp_embd(model.tok_embd);137 138    // inp_pos - built only if rope enabled139    ggml_tensor * inp_pos = nullptr;140    if (hparams.has_rope(0)) {141        inp_pos = build_inp_pos();142    }143    auto * inp_attn = build_attn_inp_kv();144 145    ggml_tensor * inp_out_ids = build_inp_out_ids();146 147    for (int il = 0; il < n_layer; ++il) {148 149        // Granite Vision 4.1 deepstack: inject the projector stream that150        // targets decoder layer `il` before the decoder runs.151        // NOTE: skip the first deepstack layer since that's inpL152        const auto & deepstack_emb_idx = hparams.deepstack_mapping_arr[il];153        if (il > 0 && deepstack_emb_idx >= 0) {154            ggml_tensor * ds = ggml_view_2d(ctx0,155                res->t_inp_embd, n_embd, n_tokens,156                res->t_inp_embd->nb[1],157                deepstack_emb_idx * n_embd * sizeof(float));158            inpL = ggml_add(ctx0, inpL, ds);159            cb(inpL, "deepstack_in", il);160        }161 162        ggml_tensor * inpSA = inpL;163 164        // norm165        cur = build_norm(inpL,166                model.layers[il].attn_norm, NULL,167                LLM_NORM_RMS, il);168        cb(cur, "attn_norm", il);169 170        // self-attention171        cur = build_attention_layer(172            cur, inp_pos, inp_attn,173            model, n_embd_head, il);174 175        if (il == n_layer - 1 && inp_out_ids) {176            cur   = ggml_get_rows(ctx0,   cur, inp_out_ids);177            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);178        }179        // ffn180        cur = build_layer_ffn(cur, inpSA, model, il);181 182        // input for next layer183        inpL = cur;184    }185    cur = inpL;186 187    cur = build_norm(cur,188            model.output_norm, NULL,189            LLM_NORM_RMS, -1);190 191    cb(cur, "result_norm", -1);192    res->t_embd = cur;193 194    // lm_head195    cur = build_lora_mm(model.output, cur, model.output_s);196 197    // For Granite architectures - scale logits198    cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_logit_scale);199    cb(cur, "result_output", -1);200    res->t_logits = cur;201 202    ggml_build_forward_expand(gf, cur);203}204 205ggml_tensor * llama_model_granite::graph::build_attention_layer(206          ggml_tensor             * cur,207          ggml_tensor             * inp_pos,208          llm_graph_input_attn_kv * inp_attn,209    const llama_model             & model,210    const int64_t                 n_embd_head,211    const int                     il) {212 213    auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,214            n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il);215 216    if (hparams.has_rope(il)) {217        ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);218        Qcur = ggml_rope_ext(219                ctx0, Qcur, inp_pos, rope_factors,220                n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,221                ext_factor, attn_factor, beta_fast, beta_slow222                );223 224        Kcur = ggml_rope_ext(225                ctx0, Kcur, inp_pos, rope_factors,226                n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,227                ext_factor, attn_factor, beta_fast, beta_slow228                );229    }230 231    cb(Qcur, "Qcur", il);232    cb(Kcur, "Kcur", il);233    cb(Vcur, "Vcur", il);234 235    const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale;236    cur = build_attn(inp_attn,237            model.layers[il].wo, model.layers[il].wo_b, model.layers[il].wo_s,238            Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);239            cb(cur, "attn_out", il);240    return cur;241}242 243ggml_tensor * llama_model_granite::graph::build_layer_ffn(244          ggml_tensor       * cur,245          ggml_tensor       * inpSA,246    const llama_model       & model,247    const int                 il) {248 249    // For Granite architectures - scale residual250    if (hparams.f_residual_scale) {251        cur = ggml_scale(ctx0, cur, hparams.f_residual_scale);252    }253    ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);254    cb(ffn_inp, "ffn_inp", il);255 256    // feed-forward network (non-MoE)257    if (model.layers[il].ffn_gate_inp == nullptr) {258 259        cur = build_norm(ffn_inp,260                model.layers[il].ffn_norm, NULL,261                LLM_NORM_RMS, il);262                cb(cur, "ffn_norm", il);263 264        cur = build_ffn(cur,265                model.layers[il].ffn_up,   model.layers[il].ffn_up_b,   NULL,266                model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, NULL,267                model.layers[il].ffn_down, model.layers[il].ffn_down_b, NULL,268                NULL,269                LLM_FFN_SILU, LLM_FFN_PAR, il);270                cb(cur, "ffn_out", il);271 272    } else {273        // MoE branch274        cur = build_norm(ffn_inp,275                model.layers[il].ffn_norm, NULL,276                LLM_NORM_RMS, il);277                cb(cur, "ffn_norm", il);278 279        ggml_tensor * moe_out = build_moe_ffn(cur,280                model.layers[il].ffn_gate_inp,281                model.layers[il].ffn_up_exps,282                model.layers[il].ffn_gate_exps,283                model.layers[il].ffn_down_exps,284                nullptr,285                n_expert, n_expert_used,286                LLM_FFN_SILU, true,287                hparams.expert_weights_scale,288                LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX,289                il);290        cb(moe_out, "ffn_moe_out", il);291 292        // For Granite MoE Shared293        if (hparams.n_ff_shexp > 0) {294            ggml_tensor * ffn_shexp = build_ffn(cur,295                model.layers[il].ffn_up_shexp,   NULL, NULL,296                model.layers[il].ffn_gate_shexp, NULL, NULL,297                model.layers[il].ffn_down_shexp, NULL, NULL,298                NULL,299                LLM_FFN_SILU, LLM_FFN_PAR, il);300            cb(ffn_shexp, "ffn_shexp", il);301 302            cur = ggml_add(ctx0, moe_out, ffn_shexp);303            cb(cur, "ffn_out", il);304        } else {305            cur = moe_out;306        }307    }308 309    // For Granite architectures - scale residual310    if (hparams.f_residual_scale) {311        cur = ggml_scale(ctx0, cur, hparams.f_residual_scale);312    }313    cur = ggml_add(ctx0, cur, ffn_inp);314    cb(cur, "ffn_out", il);315 316    cur = build_cvec(cur, il);317    cb(cur, "l_out", il);318 319    return cur;320}321