Felipe97/llama-cpp-compiled

1#include "models.h"2 3// HRM-Text: alternating low/high transformer stacks over the same token stream.4// Reference: HrmTextModel in transformers, DFM Mimir 1B.5 6void llama_model_hrm_text::load_arch_hparams(llama_model_loader & ml) {7    ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);8    ml.get_key(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale, false);9 10    ml.get_key(LLM_KV_HRM_LAYERS_PER_STACK, hparams.n_hrm_layers_per_stack);11    ml.get_key(LLM_KV_HRM_H_CYCLES, hparams.n_hrm_h_cycles);12    ml.get_key(LLM_KV_HRM_L_CYCLES, hparams.n_hrm_l_cycles);13 14    // prefix-LM prefill is not implemented (causal attention only); kept for round-trip15    ml.get_key(LLM_KV_HRM_PREFIX_LM, hparams.hrm_prefix_lm, false);16 17    GGML_ASSERT(hparams.n_hrm_layers_per_stack > 0);18    GGML_ASSERT(hparams.n_hrm_h_cycles > 0);19    GGML_ASSERT(hparams.n_hrm_l_cycles > 0);20 21    // the GGUF block count is the expanded cache-slot count22    const uint32_t n_slot = hparams.n_hrm_layers_per_stack * hparams.n_hrm_h_cycles * (hparams.n_hrm_l_cycles + 1);23    GGML_ASSERT(hparams.n_layer() == n_slot);24 25    switch (hparams.n_embd) {26        case 1536:27            type = LLM_TYPE_1B;28            break;29        default:30            type = LLM_TYPE_UNKNOWN;31    }32}33 34void llama_model_hrm_text::load_arch_tensors(llama_model_loader &) {35    LLAMA_LOAD_LOCALS;36 37    tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);38 39    // output40    output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);41    // if output is NULL, init from the input tok embed42    if (output == NULL) {43        output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED);44    }45 46    hrm_z_l_init = create_tensor(tn(LLM_TENSOR_HRM_Z_L_INIT), { n_embd }, 0);47 48    const int lps = hparams.n_hrm_layers_per_stack;49 50    // blocks [0, lps) hold the low stack, blocks [lps, 2*lps) hold the high stack.51    // the first low and high passes create the layers; later passes alias them.52    const int l_first = 0;53    const int h_first = hparams.n_hrm_l_cycles * lps;54 55    for (int h = 0; h < (int) hparams.n_hrm_h_cycles; ++h) {56        for (int l = 0; l < (int) hparams.n_hrm_l_cycles + 1; ++l) {57            const int slot_base = (h * (hparams.n_hrm_l_cycles + 1) + l) * lps;58            const int blk_base  = l == (int) hparams.n_hrm_l_cycles ? lps : 0;59 60            if (h > 0 || (l > 0 && l < (int) hparams.n_hrm_l_cycles)) {61                // alias pass: these cache slots hold the same layers as the first passes62                const int src_base = l == (int) hparams.n_hrm_l_cycles ? h_first : l_first;63                for (int il = 0; il < lps; ++il) {64                    layers[slot_base + il] = layers[src_base + il];65                }66                continue;67            }68 69            for (int il = 0; il < lps; ++il) {70                auto &    layer = layers[slot_base + il];71                const int bid   = blk_base + il;72 73                create_tensor_qkv(layer, bid, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0);74 75                // sigmoid attention gate, applied to the attention output before o_proj76                layer.wqkv_gate =77                    create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", bid), { n_embd, n_embd_head_k * n_head }, 0);78                layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", bid), { n_embd_head_k * n_head, n_embd }, 0);79 80                layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", bid), { n_embd, n_ff }, 0);81                layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", bid), { n_ff, n_embd }, 0);82                layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", bid), { n_embd, n_ff }, 0);83            }84        }85    }86}87 88std::unique_ptr<llm_graph_context> llama_model_hrm_text::build_arch_graph(const llm_graph_params & params) const {89    return std::make_unique<graph>(*this, params);90}91 92// one stack invocation: lps pre-norm decoder layers, then the parameterless final norm93ggml_tensor * llama_model_hrm_text::graph::build_stack(llm_graph_input_attn_kv * inp_attn,94                                                       ggml_tensor *             inp_pos,95                                                       ggml_tensor *             cur,96                                                       int                       slot_base) const {97    const float kq_scale = 1.0f / sqrtf(float(n_embd_head_k));98 99    const int lps = model.hparams.n_hrm_layers_per_stack;100 101    for (int il = 0; il < lps; ++il) {102        const int    s     = slot_base + il;103        const auto & layer = model.layers[s];104 105        ggml_tensor * inpSA = cur;106 107        cur = build_norm(cur, nullptr, nullptr, LLM_NORM_RMS, s);108        cb(cur, "attn_norm", s);109 110        // sigmoid-gated self-attention (same shape as qwen3next attention layers)111        {112            ggml_tensor * gate = build_lora_mm(layer.wqkv_gate, cur);113            cb(gate, "attn_gate_proj", s);114 115            auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head_k, n_head, n_head_kv, s);116 117            Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,118                                 ext_factor, attn_factor, beta_fast, beta_slow);119            cb(Qcur, "Qcur", s);120 121            Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,122                                 ext_factor, attn_factor, beta_fast, beta_slow);123            cb(Kcur, "Kcur", s);124 125            cur = build_attn(inp_attn,126                nullptr, nullptr, nullptr,127                Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, s);128            cb(cur, "attn_pregate", s);129 130            gate = ggml_sigmoid(ctx0, gate);131            cb(gate, "attn_gate_sigmoid", s);132 133            cur = ggml_mul(ctx0, cur, gate);134            cb(cur, "attn_gated", s);135 136            cur = build_lora_mm(layer.wo, cur, layer.wo_s);137            cb(cur, "attn_out", s);138        }139 140        cur = ggml_add(ctx0, cur, inpSA);141        cb(cur, "attn_add", s);142 143        inpSA = cur;144        cur   = build_norm(cur, nullptr, nullptr, LLM_NORM_RMS, s);145        cb(cur, "ffn_norm", s);146 147        cur = build_ffn(cur,148            layer.ffn_up, nullptr, nullptr,149            layer.ffn_gate, nullptr, nullptr,150            layer.ffn_down, nullptr, nullptr,151            nullptr,152            LLM_FFN_SILU, LLM_FFN_PAR, s);153        cb(cur, "ffn_out", s);154 155        cur = ggml_add(ctx0, cur, inpSA);156        cb(cur, "ffn_add", s);157 158        cur = build_cvec(cur, s);159        cb(cur, "l_out", s);160    }161 162    cur = build_norm(cur, nullptr, nullptr, LLM_NORM_RMS, slot_base);163    cb(cur, "stack_norm", slot_base);164 165    return cur;166}167 168llama_model_hrm_text::graph::graph(const llama_model & model, const llm_graph_params & params) :169    llm_graph_context(params),170    model(model) {171    ggml_tensor * cur;172 173    // {n_embd, n_tokens}, scaled by hparams.f_embedding_scale inside build_inp_embd174    ggml_tensor * zH = build_inp_embd(model.tok_embd);175 176    ggml_tensor * inp_pos = build_inp_pos();177 178    auto * inp_attn = build_attn_inp_kv();179 180    ggml_tensor * inp_out_ids = build_inp_out_ids();181 182    // the learned low-cycle state is [n_embd]; binary ops broadcast it over [n_embd, n_tokens]183    ggml_tensor * zL = model.hrm_z_l_init;184 185    for (uint32_t h = 0; h < model.hparams.n_hrm_h_cycles; ++h) {186        for (uint32_t l = 0; l < model.hparams.n_hrm_l_cycles; ++l) {187            const int slot_base = (h * (model.hparams.n_hrm_l_cycles + 1) + l) * model.hparams.n_hrm_layers_per_stack;188 189            zL = build_stack(inp_attn, inp_pos, ggml_add(ctx0, zH, zL), slot_base);190        }191 192        const int slot_base = (h * (model.hparams.n_hrm_l_cycles + 1) + model.hparams.n_hrm_l_cycles) *193                              model.hparams.n_hrm_layers_per_stack;194 195        zH = build_stack(inp_attn, inp_pos, ggml_add(ctx0, zH, zL), slot_base);196    }197 198    cur = zH;199 200    if (inp_out_ids) {201        cur = ggml_get_rows(ctx0, cur, inp_out_ids);202    }203 204    cb(cur, "result_norm", -1);205    res->t_embd = cur;206 207    cur = build_lora_mm(model.output, cur, model.output_s);208 209    cb(cur, "result_output", -1);210    res->t_logits = cur;211 212    ggml_build_forward_expand(gf, cur);213}214