Felipe97/llama-cpp-compiled
01.1k
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}541