echodict/llama.cpp

1#include "models.h"2 3// RND1 is a Qwen3Moe AR model converted to diffusion model.4llm_build_rnd1::llm_build_rnd1(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {5    const int64_t n_embd_head = hparams.n_embd_head_v();6 7    GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());8    GGML_ASSERT(n_embd_head == n_rot);9 10    ggml_tensor * cur;11    ggml_tensor * inpL;12 13    inpL = build_inp_embd(model.tok_embd);14 15    // inp_pos - contains the positions16    ggml_tensor * inp_pos = build_inp_pos();17 18    // Non-causal attention for diffusion19    auto * inp_attn = build_attn_inp_no_cache();20 21    ggml_tensor * inp_out_ids = build_inp_out_ids();22 23    for (int il = 0; il < n_layer; ++il) {24        ggml_tensor * inpSA = inpL;25 26        // norm27        cur = build_norm(inpL,28                model.layers[il].attn_norm, NULL,29                LLM_NORM_RMS, il);30        cb(cur, "attn_norm", il);31 32        // self_attention33        {34            // compute Q and K and RoPE them35            auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,36                    n_embd_head, n_head, n_head_kv, il);37 38            Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il);39            cb(Qcur, "Qcur_normed", il);40 41            Qcur = ggml_rope_ext(42                    ctx0, Qcur, inp_pos, nullptr,43                    n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,44                    ext_factor, attn_factor, beta_fast, beta_slow45                    );46 47            Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il);48            cb(Kcur, "Kcur_normed", il);49 50            Kcur = ggml_rope_ext(51                    ctx0, Kcur, inp_pos, nullptr,52                    n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,53                    ext_factor, attn_factor, beta_fast, beta_slow54                    );55 56            cb(Qcur, "Qcur", il);57            cb(Kcur, "Kcur", il);58            cb(Vcur, "Vcur", il);59 60            cur = build_attn(inp_attn,61                    model.layers[il].wo, model.layers[il].wo_b, model.layers[il].wo_s,62                    Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);63        }64        if (il == n_layer - 1 && inp_out_ids) {65            cur   = ggml_get_rows(ctx0,   cur, inp_out_ids);66            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);67        }68        ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);69        cb(ffn_inp, "ffn_inp", il);70 71        // MoE branch72        cur = build_norm(ffn_inp,73                model.layers[il].ffn_norm, NULL,74                LLM_NORM_RMS, il);75        cb(cur, "ffn_norm", il);76 77        ggml_tensor * moe_out =78            build_moe_ffn(cur,79                    model.layers[il].ffn_gate_inp,80                    model.layers[il].ffn_up_exps,81                    model.layers[il].ffn_gate_exps,82                    model.layers[il].ffn_down_exps,83                    nullptr,84                    n_expert, n_expert_used,85                    LLM_FFN_SILU, true,86                    hparams.expert_weights_scale,87                    LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX,88                    il);89        cb(moe_out, "ffn_moe_out", il);90        cur = moe_out;91 92        cur = ggml_add(ctx0, cur, ffn_inp);93 94        cur = build_cvec(cur, il);95        cb(cur, "l_out", il);96 97        // input for next layer98        inpL = cur;99    }100    cur = inpL;101 102    cur = build_norm(cur,103            model.output_norm, NULL,104            LLM_NORM_RMS, -1);105 106    cb(cur, "result_norm", -1);107    res->t_embd = cur;108 109    // lm_head110    cur = build_lora_mm(model.output, cur);111 112    cb(cur, "result_output", -1);113    res->t_logits = cur;114 115    ggml_build_forward_expand(gf, cur);116}117