hey! this is a static git repo viewer. commit diffs, file history, older commits and other revisions won't load without javascript.

you are still able to view source files on branch/tag tips, but the experience is degraded.

char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Daniel Hanmodels : fix GDN normalization from `max` to `rsqrt` (#28068)5fdfa62822026-09-06

1#include "models.h"
2#include "llama-memory-recurrent.h"
3
4#include <algorithm>
5
6void llama_model_bailingmoe3::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_ATTENTION_KEY_LENGTH_MLA,         hparams.n_embd_head_k_mla_impl);
9    ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA,       hparams.n_embd_head_v_mla_impl);
10    ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK,           hparams.n_lora_kv);
11    ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK,            hparams.n_lora_q, false);
12    ml.get_key(LLM_KV_SSM_CONV_KERNEL,                  hparams.ssm_d_conv);
13    ml.get_key(LLM_KV_KDA_HEAD_DIM,                     hparams.n_embd_head_kda);
14    if (!ml.get_key(LLM_KV_KDA_SAFE_GATE, hparams.kda_safe_gate, false)) {
15        hparams.kda_safe_gate = true;
16    }
17    ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND,             hparams.kda_gate_lower_bound);
18    ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);
19    ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false);
20    ml.get_key(LLM_KV_EXPERT_SHARED_COUNT,              hparams.n_expert_shared);
21    ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT,        hparams.n_layer_dense_lead);
22    ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE,             hparams.expert_weights_scale, false);
23    ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM,              hparams.expert_weights_norm, false);
24    ml.get_key(LLM_KV_EXPERT_GATING_FUNC,               hparams.expert_gating_func);
25    ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP,           hparams.swiglu_clamp_exp,   hparams.n_layer_all, false);
26    ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP,         hparams.swiglu_clamp_shexp, hparams.n_layer_all, false);
27
28    if (hparams.n_ff_shexp == 0) {
29        hparams.n_ff_shexp = hparams.n_ff_exp() * std::max(1u, hparams.n_expert_shared);
30    }
31
32    GGML_ASSERT(hparams.kda_safe_gate);
33    GGML_ASSERT(hparams.kda_gate_lower_bound < 0.0f);
34
35    for (uint32_t il = 0; il < hparams.n_layer(); ++il) {
36        hparams.is_recr_impl[il] = hparams.n_head_kv(il) == 0;
37    }
38
39    switch (hparams.n_layer()) {
40        case 24: type = hparams.n_embd == 1536 && hparams.n_expert == 128 ? LLM_TYPE_7_9B_A1_3B : LLM_TYPE_UNKNOWN; break;
41        case 42: type = hparams.n_embd == 2560 && hparams.n_expert == 512 ? LLM_TYPE_124B_A5_1B : LLM_TYPE_UNKNOWN; break;
42        default: type = LLM_TYPE_UNKNOWN;
43    }
44}
45
46void llama_model_bailingmoe3::load_arch_tensors(llama_model_loader & ml) {
47    LLAMA_LOAD_LOCALS;
48
49    tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0);
50
51    output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0);
52    output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED);
53    if (output == nullptr) {
54        output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED);
55    }
56
57    const int64_t head_dim = hparams.n_embd_head_kda;
58    const int64_t d_inner = head_dim * n_head;
59    const int64_t d_conv = hparams.ssm_d_conv;
60    const int64_t kv_lora_rank = hparams.n_lora_kv;
61    const int64_t q_lora_rank  = hparams.n_lora_q;
62    const int64_t qk_rope_head_dim = hparams.n_rot();
63    const int64_t qk_head_dim = hparams.n_embd_head_k_mla();
64    const int64_t v_head_dim = hparams.n_embd_head_v_mla();
65
66    const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr);
67    const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight";
68    const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr);
69    const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0;
70    int       mtp_flags   = trunk_only ? TENSOR_NOT_REQUIRED : 0;
71
72    if (!ml.load_mtp) {
73        mtp_flags |= TENSOR_SKIP;
74    }
75
76    for (int il = 0; il < n_layer; ++il) {
77        auto & layer = layers[il];
78
79        layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, trunk_flags);
80
81        if (hparams.is_recr(il)) {
82            layer.ssm_q_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_Q, "weight", il), { d_conv, 1, d_inner, 1 }, trunk_flags);
83            layer.ssm_k_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_K, "weight", il), { d_conv, 1, d_inner, 1 }, trunk_flags);
84            layer.ssm_v_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_V, "weight", il), { d_conv, 1, d_inner, 1 }, trunk_flags);
85
86            create_tensor_qkv(layer, il, n_embd, d_inner, d_inner, d_inner, trunk_flags);
87            layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", il), { n_embd, d_inner }, trunk_flags);
88            layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_head }, trunk_flags);
89            layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, il), { 1, n_head }, trunk_flags);
90            layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { d_inner }, trunk_flags);
91            layer.ssm_g_a = create_tensor(tn(LLM_TENSOR_SSM_G_A, "weight", il), { n_embd, d_inner }, trunk_flags);
92            layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_dim }, trunk_flags);
93            layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { d_inner, n_embd }, trunk_flags);
94        } else {
95            if (q_lora_rank > 0) {
96                layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", il), { n_embd, q_lora_rank }, trunk_flags);
97                layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", il), { q_lora_rank }, trunk_flags);
98                layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", il), { q_lora_rank, n_head * qk_head_dim }, trunk_flags);
99            } else {
100                layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", il), { n_embd, n_head * qk_head_dim }, trunk_flags);
101            }
102            layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", il), { n_embd, kv_lora_rank + qk_rope_head_dim }, trunk_flags);
103            layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", il), { kv_lora_rank }, trunk_flags);
104            layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", il), { qk_head_dim - qk_rope_head_dim, kv_lora_rank, n_head }, trunk_flags);
105            layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", il), { kv_lora_rank, v_head_dim, n_head }, trunk_flags);
106            layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, n_head }, trunk_flags);
107            layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_head * v_head_dim, n_embd }, trunk_flags);
108        }
109
110        layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", il), { n_embd }, trunk_flags);
111        if ((uint32_t) il < hparams.n_layer_dense_lead) {
112            layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", il), { n_embd, n_ff }, trunk_flags);
113            layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", il), { n_embd, n_ff }, trunk_flags);
114            layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", il), { n_ff, n_embd }, trunk_flags);
115        } else {
116            layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, trunk_flags);
117            layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", il), { n_expert }, trunk_flags);
118            layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", il), { n_embd, hparams.n_ff_exp(), n_expert }, trunk_flags);
119            layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", il), { n_embd, hparams.n_ff_exp(), n_expert }, trunk_flags);
120            layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { hparams.n_ff_exp(), n_embd, n_expert }, trunk_flags);
121            layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, trunk_flags);
122            layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, trunk_flags);
123            layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { hparams.n_ff_shexp, n_embd }, trunk_flags);
124        }
125    }
126
127    for (int il = n_layer; il < n_layer_all; ++il) {
128        auto & layer = layers[il];
129        const int flags = mtp_flags;
130
131        layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, flags);
132        if (q_lora_rank > 0) {
133            layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", il), { n_embd, q_lora_rank }, flags);
134            layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", il), { q_lora_rank }, flags);
135            layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", il), { q_lora_rank, n_head * qk_head_dim }, flags);
136        } else {
137            layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", il), { n_embd, n_head * qk_head_dim }, flags);
138        }
139        layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", il), { n_embd, kv_lora_rank + qk_rope_head_dim }, flags);
140        layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", il), { kv_lora_rank }, flags);
141        layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", il), { qk_head_dim - qk_rope_head_dim, kv_lora_rank, n_head }, flags);
142        layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", il), { kv_lora_rank, v_head_dim, n_head }, flags);
143        layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, n_head }, flags);
144        layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_head * v_head_dim, n_embd }, flags);
145        layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", il), { n_embd }, flags);
146        layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, flags);
147        layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", il), { n_expert }, flags);
148        layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", il), { n_embd, hparams.n_ff_exp(), n_expert }, flags);
149        layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", il), { n_embd, hparams.n_ff_exp(), n_expert }, flags);
150        layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { hparams.n_ff_exp(), n_embd, n_expert }, flags);
151        layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, flags);
152        layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, flags);
153        layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { hparams.n_ff_shexp, n_embd }, flags);
154        layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, flags);
155        layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, flags);
156        layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, flags);
157        layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_LAYER_OUT_NORM, "weight", il), { n_embd }, flags);
158    }
159}
160
161std::unique_ptr<llm_graph_context> llama_model_bailingmoe3::build_arch_graph(const llm_graph_params & params) const {
162    if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
163        return std::make_unique<graph_mtp>(*this, params);
164    }
165    return std::make_unique<graph>(*this, params);
166}
167
168static ggml_tensor * bailingmoe3_causal_conv1d(
169        ggml_cgraph * gf,
170        ggml_context * ctx0,
171        ggml_tensor * conv_states_all,
172        ggml_tensor * conv_state_all,
173        int64_t qkv,
174        ggml_tensor * x,
175        ggml_tensor * proj_w,
176        ggml_tensor * conv_w,
177        int64_t d_conv,
178        int64_t head_dim,
179        int64_t n_head,
180        int64_t n_seq_tokens,
181        int64_t n_seqs,
182        int64_t n_tokens,
183        int64_t cache_head,
184        uint32_t mem_size,
185        uint32_t n_rs_seq) {
186    const int64_t d_inner = head_dim * n_head;
187    const int64_t conv_state_size = (d_conv - 1) * d_inner;
188    const int64_t total_state_size = 3 * conv_state_size;
189
190    ggml_tensor * conv_state = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs,
191            (d_conv - 1) * ggml_element_size(conv_state_all),
192            total_state_size * ggml_element_size(conv_state_all),
193            qkv * conv_state_size * ggml_element_size(conv_state_all));
194
195    ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x);
196    x_proj = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);
197    ggml_tensor * conv_x = ggml_concat(ctx0, conv_state, ggml_transpose(ctx0, x_proj), 0);
198
199    const int64_t K = (int64_t) n_rs_seq + 1;
200    const int64_t n_written = std::min<int64_t>(n_seq_tokens, K);
201
202    for (int64_t slot = 0; slot < n_written; ++slot) {
203        ggml_tensor * conv_snap = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,
204                conv_x->nb[1], conv_x->nb[2], (conv_x->ne[0] - (d_conv - 1) - slot) * conv_x->nb[0]);
205        ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_snap,
206                ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs,
207                    (d_conv - 1) * ggml_element_size(conv_states_all),
208                    total_state_size * ggml_element_size(conv_states_all),
209                    ((slot * mem_size + cache_head) * total_state_size + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
210    }
211
212    ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner);
213    ggml_tensor * out = ggml_ssm_conv(ctx0, conv_x, conv_weight);
214    out = ggml_silu(ctx0, ggml_reshape_2d(ctx0, out, d_inner, n_tokens));
215    return ggml_reshape_4d(ctx0, out, head_dim, n_head, n_seq_tokens, n_seqs);
216}
217
218llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph_params & params) :
219    llm_build_delta_net_base(params), model(model) {
220    ggml_tensor * inpL = build_inp_embd(model.tok_embd);
221    cb(inpL, "model.input_embed", -1);
222
223    auto * inp = build_inp_mem_hybrid_k();
224    auto * inp_rs = inp->get_recr();
225    auto * inp_attn = inp->get_attn();
226
227    ggml_tensor * inp_pos = build_inp_pos();
228    ggml_tensor * inp_out_ids = build_inp_out_ids();
229
230    const int64_t n_head = hparams.n_head();
231    const int64_t head_dim = hparams.n_embd_head_kda;
232    const int64_t d_inner = n_head * head_dim;
233    const int64_t d_conv = hparams.ssm_d_conv;
234    const int64_t n_seqs = ubatch.n_seqs;
235    const int64_t n_seq_tokens = ubatch.n_seq_tokens;
236    const int64_t qk_head_dim = hparams.n_embd_head_k_mla();
237    const int64_t v_head_dim = hparams.n_embd_head_v_mla();
238    const int64_t qk_rope_head_dim = hparams.n_rot();
239    const int64_t qk_nope_head_dim = qk_head_dim - qk_rope_head_dim;
240    const int64_t kv_lora_rank = hparams.n_lora_kv;
241    const float kq_scale = 1.0f / sqrtf((float) qk_head_dim);
242
243    GGML_ASSERT(n_seqs > 0);
244    GGML_ASSERT(ubatch.equal_seqs());
245    GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);
246
247    for (int il = 0; il < n_layer; ++il) {
248        res->t_layer_inp[il] = inpL;
249
250        const auto & layer = model.layers[il];
251        ggml_tensor * inpSA = inpL;
252        ggml_tensor * cur = build_norm(inpL, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
253        cb(cur, "attn_norm", il);
254
255        if (hparams.is_recr(il)) {
256            const auto * mctx_cur = inp_rs->mctx;
257            const auto cache_head = mctx_cur->get_head();
258            const auto mem_size = mctx_cur->get_size();
259            ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
260            ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);
261
262            ggml_tensor * q = bailingmoe3_causal_conv1d(
263                    gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv,
264                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
265            ggml_tensor * k = bailingmoe3_causal_conv1d(
266                    gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv,
267                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
268            ggml_tensor * v = bailingmoe3_causal_conv1d(
269                    gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv,
270                    d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq);
271
272            ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_f_a, cur);
273            gate = ggml_add(ctx0, gate, layer.ssm_dt_b);
274            gate = ggml_reshape_3d(ctx0, gate, head_dim, n_head, n_tokens);
275            ggml_tensor * a = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head, 1);
276            gate = ggml_scale(ctx0, ggml_sigmoid(ctx0, ggml_mul(ctx0, gate, a)), hparams.kda_gate_lower_bound);
277            gate = ggml_reshape_4d(ctx0, gate, head_dim, n_head, n_seq_tokens, n_seqs);
278            cb(gate, "kda_gate", il);
279
280            ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur);
281            beta = ggml_sigmoid(ctx0, ggml_reshape_4d(ctx0, beta, 1, n_head, n_seq_tokens, n_seqs));
282
283            q = build_gdn_l2_norm(ctx0, q, hparams.f_norm_rms_eps);
284            k = build_gdn_l2_norm(ctx0, k, hparams.f_norm_rms_eps);
285
286            ggml_tensor * states_all = mctx_cur->get_s_l(il);
287            ggml_tensor * state = build_rs(inp_rs, states_all, hparams.n_embd_s(), n_seqs);
288            state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head, n_seqs);
289
290            ggml_tensor * out = ggml_cont(ctx0, build_recurrent_attn(
291                    inp_rs, states_all, q, k, v, gate, beta, state, il));
292
293            ggml_tensor * out_gate = ggml_mul_mat(ctx0, layer.ssm_g_a, cur);
294            out_gate = ggml_reshape_3d(ctx0, out_gate, head_dim, n_head, n_tokens);
295            out = ggml_reshape_3d(ctx0, out, head_dim, n_head, n_tokens);
296            out = build_norm(out, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il);
297            out = ggml_mul(ctx0, out, ggml_sigmoid(ctx0, out_gate));
298            cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, out, d_inner, n_tokens));
299            cb(cur, "kda_out", il);
300        } else {
301            ggml_tensor * attn_input = cur;
302            ggml_tensor * q_all;
303            if (layer.wq_a) {
304                q_all = ggml_mul_mat(ctx0, layer.wq_a, cur);
305                cb(q_all, "q_a", il);
306                q_all = build_norm(q_all, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
307                cb(q_all, "q_a_norm", il);
308                q_all = ggml_mul_mat(ctx0, layer.wq_b, q_all);
309                cb(q_all, "q_b", il);
310            } else {
311                q_all = ggml_mul_mat(ctx0, layer.wq, cur);
312            }
313            ggml_tensor * q_nope = ggml_view_3d(ctx0, q_all, qk_nope_head_dim, n_head, n_tokens,
314                    ggml_row_size(q_all->type, qk_head_dim),
315                    ggml_row_size(q_all->type, qk_head_dim) * n_head, 0);
316            ggml_tensor * q_pe = ggml_view_3d(ctx0, q_all, qk_rope_head_dim, n_head, n_tokens,
317                    ggml_row_size(q_all->type, qk_head_dim),
318                    ggml_row_size(q_all->type, qk_head_dim) * n_head,
319                    ggml_row_size(q_all->type, qk_nope_head_dim));
320
321            ggml_tensor * kv_all = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);
322            ggml_tensor * kv = ggml_view_2d(ctx0, kv_all, kv_lora_rank, n_tokens,
323                    ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), 0);
324            ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_all, qk_rope_head_dim, 1, n_tokens,
325                    ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim),
326                    ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim),
327                    ggml_row_size(kv_all->type, kv_lora_rank));
328
329            q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
330                    ext_factor, attn_factor, beta_fast, beta_slow);
331            k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
332                    ext_factor, attn_factor, beta_fast, beta_slow);
333            kv = build_norm(kv, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
334
335            q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
336            q_nope = ggml_mul_mat(ctx0, layer.wk_b, q_nope);
337            q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
338
339            ggml_tensor * q = ggml_concat(ctx0, q_nope, q_pe, 0);
340            kv = ggml_reshape_3d(ctx0, kv, kv_lora_rank, 1, n_tokens);
341            ggml_tensor * k = ggml_concat(ctx0, kv, k_pe, 0);
342
343            cur = build_attn(inp_attn, nullptr, nullptr, nullptr,
344                    q, k, kv, nullptr, nullptr, layer.wv_b, kq_scale, il);
345
346            ggml_tensor * attn_gate = ggml_mul_mat(ctx0, layer.wqkv_gate, attn_input);
347            attn_gate = ggml_sigmoid(ctx0, ggml_reshape_3d(ctx0, attn_gate, 1, n_head, n_tokens));
348            cur = ggml_reshape_3d(ctx0, cur, v_head_dim, n_head, n_tokens);
349            cur = ggml_mul(ctx0, cur, attn_gate);
350            cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, cur, v_head_dim * n_head, n_tokens));
351            cb(cur, "mla_out", il);
352        }
353
354        if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) {
355            cur = ggml_get_rows(ctx0, cur, inp_out_ids);
356            inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
357        }
358
359        ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
360        cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
361
362        if ((uint32_t) il < hparams.n_layer_dense_lead) {
363            cur = build_ffn(cur,
364                    layer.ffn_up, nullptr, nullptr,
365                    layer.ffn_gate, nullptr, nullptr,
366                    layer.ffn_down, nullptr, nullptr,
367                    nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
368        } else {
369            ggml_tensor * moe = build_moe_ffn(cur,
370                    layer.ffn_gate_inp,
371                    layer.ffn_up_exps,
372                    layer.ffn_gate_exps,
373                    layer.ffn_down_exps,
374                    layer.ffn_exp_probs_b,
375                    n_expert, n_expert_used,
376                    LLM_FFN_SILU,
377                    hparams.expert_weights_norm,
378                    hparams.expert_weights_scale,
379                    (llama_expert_gating_func_type) hparams.expert_gating_func,
380                    il);
381            ggml_tensor * shared = build_ffn(cur,
382                    layer.ffn_up_shexp, nullptr, nullptr,
383                    layer.ffn_gate_shexp, nullptr, nullptr,
384                    layer.ffn_down_shexp, nullptr, nullptr,
385                    nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
386            cur = ggml_add(ctx0, moe, shared);
387        }
388
389        cur = ggml_add(ctx0, cur, ffn_inp);
390        cur = build_cvec(cur, il);
391        cb(cur, "l_out", il);
392        inpL = cur;
393    }
394
395    ggml_tensor * cur = build_norm(inpL, model.output_norm, nullptr, LLM_NORM_RMS, -1);
396    cb(cur, "h_nextn", -1);
397    res->t_h_nextn = cur;
398
399    if (!cparams.embeddings_nextn_masked && inp_out_ids) {
400        cur = ggml_get_rows(ctx0, cur, inp_out_ids);
401    }
402
403    cb(cur, "result_norm", -1);
404    res->t_embd = cur;
405
406    cur = ggml_mul_mat(ctx0, model.output, cur);
407    cb(cur, "result_output", -1);
408    res->t_logits = cur;
409    ggml_build_forward_expand(gf, cur);
410}
411
412llama_model_bailingmoe3::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) :
413    llm_graph_context(params) {
414    GGML_ASSERT(hparams.n_layer_nextn == 1 && "BailingMoE3 MTP requires one NextN layer");
415
416    const int il = hparams.n_layer() + cparams.nextn_layer_offset;
417    GGML_ASSERT(cparams.nextn_layer_offset >= 0 &&
418                cparams.nextn_layer_offset < (int) hparams.n_layer_nextn &&
419                "nextn_layer_offset out of range");
420    const auto & layer = model.layers[il];
421
422    GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj");
423    GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm");
424    GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm");
425    GGML_ASSERT(layer.nextn.shared_head_norm && "MTP block missing final norm");
426
427    const int64_t n_head = hparams.n_head();
428    const int64_t qk_head_dim = hparams.n_embd_head_k_mla();
429    const int64_t v_head_dim = hparams.n_embd_head_v_mla();
430    const int64_t qk_rope_head_dim = hparams.n_rot();
431    const int64_t qk_nope_head_dim = qk_head_dim - qk_rope_head_dim;
432    const int64_t kv_lora_rank = hparams.n_lora_kv;
433    const float kq_scale = 1.0f / sqrtf((float) qk_head_dim);
434
435    auto inp = std::make_unique<llm_graph_input_embd>(hparams.n_embd);
436    inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
437    ggml_set_input(inp->tokens);
438    inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens);
439    ggml_set_input(inp->embd);
440    ggml_set_name(inp->embd, "mtp_h_input");
441
442    ggml_tensor * tok_embd = ggml_get_rows(ctx0, model.tok_embd, inp->tokens);
443    ggml_tensor * h_norm = build_norm(inp->embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il);
444    ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il);
445    ggml_tensor * cur = ggml_mul_mat(ctx0, layer.nextn.eh_proj, ggml_concat(ctx0, e_norm, h_norm, 0));
446    cb(cur, "mtp_eh_proj", il);
447
448    res->add_input(std::move(inp));
449
450    ggml_tensor * inp_pos = build_inp_pos();
451    ggml_tensor * inp_out_ids = build_inp_out_ids();
452    auto * inp_attn = build_attn_inp_k();
453
454    ggml_tensor * inpSA = cur;
455    cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
456    ggml_tensor * attn_input = cur;
457
458    ggml_tensor * q_all;
459    if (layer.wq_a) {
460        q_all = ggml_mul_mat(ctx0, layer.wq_a, cur);
461        cb(q_all, "q_a", il);
462        q_all = build_norm(q_all, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il);
463        cb(q_all, "q_a_norm", il);
464        q_all = ggml_mul_mat(ctx0, layer.wq_b, q_all);
465        cb(q_all, "q_b", il);
466    } else {
467        q_all = ggml_mul_mat(ctx0, layer.wq, cur);
468    }
469    ggml_tensor * q_nope = ggml_view_3d(ctx0, q_all, qk_nope_head_dim, n_head, n_tokens,
470            ggml_row_size(q_all->type, qk_head_dim),
471            ggml_row_size(q_all->type, qk_head_dim) * n_head, 0);
472    ggml_tensor * q_pe = ggml_view_3d(ctx0, q_all, qk_rope_head_dim, n_head, n_tokens,
473            ggml_row_size(q_all->type, qk_head_dim),
474            ggml_row_size(q_all->type, qk_head_dim) * n_head,
475            ggml_row_size(q_all->type, qk_nope_head_dim));
476
477    ggml_tensor * kv_all = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);
478    ggml_tensor * kv = ggml_view_2d(ctx0, kv_all, kv_lora_rank, n_tokens,
479            ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), 0);
480    ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_all, qk_rope_head_dim, 1, n_tokens,
481            ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim),
482            ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim),
483            ggml_row_size(kv_all->type, kv_lora_rank));
484
485    q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
486            ext_factor, attn_factor, beta_fast, beta_slow);
487    k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
488            ext_factor, attn_factor, beta_fast, beta_slow);
489    kv = build_norm(kv, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
490
491    q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
492    q_nope = ggml_mul_mat(ctx0, layer.wk_b, q_nope);
493    q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
494
495    ggml_tensor * q = ggml_concat(ctx0, q_nope, q_pe, 0);
496    kv = ggml_reshape_3d(ctx0, kv, kv_lora_rank, 1, n_tokens);
497    ggml_tensor * k = ggml_concat(ctx0, kv, k_pe, 0);
498
499    cur = build_attn(inp_attn, nullptr, nullptr, nullptr,
500            q, k, kv, nullptr, nullptr, layer.wv_b, kq_scale, il);
501
502    ggml_tensor * attn_gate = ggml_mul_mat(ctx0, layer.wqkv_gate, attn_input);
503    attn_gate = ggml_sigmoid(ctx0, ggml_reshape_3d(ctx0, attn_gate, 1, n_head, n_tokens));
504    cur = ggml_reshape_3d(ctx0, cur, v_head_dim, n_head, n_tokens);
505    cur = ggml_mul(ctx0, cur, attn_gate);
506    cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, cur, v_head_dim * n_head, n_tokens));
507
508    ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
509    cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
510
511    ggml_tensor * moe = build_moe_ffn(cur,
512            layer.ffn_gate_inp,
513            layer.ffn_up_exps,
514            layer.ffn_gate_exps,
515            layer.ffn_down_exps,
516            layer.ffn_exp_probs_b,
517            n_expert, n_expert_used,
518            LLM_FFN_SILU,
519            hparams.expert_weights_norm,
520            hparams.expert_weights_scale,
521            (llama_expert_gating_func_type) hparams.expert_gating_func,
522            il);
523    ggml_tensor * shared = build_ffn(cur,
524            layer.ffn_up_shexp, nullptr, nullptr,
525            layer.ffn_gate_shexp, nullptr, nullptr,
526            layer.ffn_down_shexp, nullptr, nullptr,
527            nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
528    cur = ggml_add(ctx0, moe, shared);
529    cur = ggml_add(ctx0, cur, ffn_inp);
530    cur = build_norm(cur, layer.nextn.shared_head_norm, nullptr, LLM_NORM_RMS, -1);
531
532    cb(cur, "h_nextn", -1);
533    res->t_h_nextn = cur;
534
535    cur = ggml_get_rows(ctx0, cur, inp_out_ids);
536    cur = ggml_mul_mat(ctx0, model.output, cur);
537    cb(cur, "result_output", -1);
538    res->t_logits = cur;
539    ggml_build_forward_expand(gf, cur);
540}