diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 72968607db8..232d9f45e35 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -105,6 +105,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_PLM, "plm" }, { LLM_ARCH_BAILINGMOE, "bailingmoe" }, { LLM_ARCH_BAILINGMOE2, "bailingmoe2" }, + { LLM_ARCH_BAILINGMOE3, "bailingmoe3" }, { LLM_ARCH_DOTS1, "dots1" }, { LLM_ARCH_ARCEE, "arcee" }, { LLM_ARCH_AFMOE, "afmoe" }, @@ -294,7 +295,9 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_SSM_GROUP_COUNT, "%s.ssm.group_count" }, { LLM_KV_SSM_DT_B_C_RMS, "%s.ssm.dt_b_c_rms" }, - { LLM_KV_KDA_HEAD_DIM, "%s.kda.head_dim" }, + { LLM_KV_KDA_HEAD_DIM, "%s.kda.head_dim" }, + { LLM_KV_KDA_SAFE_GATE, "%s.kda.safe_gate" }, + { LLM_KV_KDA_GATE_LOWER_BOUND, "%s.kda.gate_lower_bound" }, { LLM_KV_WKV_HEAD_SIZE, "%s.wkv.head_size" }, @@ -944,6 +947,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_QWEN3NEXT: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_BAILINGMOE3: case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: return true; @@ -1001,6 +1005,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_MINIMAX_M2: case LLM_ARCH_MISTRAL4: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_BAILINGMOE3: return false; default: return true; diff --git a/src/llama-arch.h b/src/llama-arch.h index b74d53af4a3..f4a4e5dde6d 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -110,6 +110,7 @@ enum llm_arch { LLM_ARCH_PLM, LLM_ARCH_BAILINGMOE, LLM_ARCH_BAILINGMOE2, + LLM_ARCH_BAILINGMOE3, LLM_ARCH_DOTS1, LLM_ARCH_ARCEE, LLM_ARCH_AFMOE, @@ -300,6 +301,8 @@ enum llm_kv { LLM_KV_SSM_DT_B_C_RMS, LLM_KV_KDA_HEAD_DIM, + LLM_KV_KDA_SAFE_GATE, + LLM_KV_KDA_GATE_LOWER_BOUND, LLM_KV_WKV_HEAD_SIZE, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 8794dc1a1da..92dc0de0b74 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2318,6 +2318,7 @@ void llama_context::output_reorder() { uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || + model.arch == LLM_ARCH_BAILINGMOE3 || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_DEEPSEEK4) { diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 8be5f28f39e..b1a32da7815 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -161,6 +161,8 @@ struct llama_hparams { // for Kimi Linear KDA uint32_t n_embd_head_kda = 0; + bool kda_safe_gate = false; + float kda_gate_lower_bound = 0.0f; bool ssm_dt_b_c_rms = false; diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index a3928523ba8..1a3ff4c2ad5 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -119,6 +119,7 @@ void llama_model_saver::add_kv(const enum llm_kv key, const Container & value, c } // instantiate for external usage: template void llama_model_saver::add_kv>(const enum llm_kv, const std::vector &, const bool); +template void llama_model_saver::add_kv>(const enum llm_kv, const std::vector &, const bool); void llama_model_saver::add_kv(const enum llm_kv key, const std::vector & value) { std::vector tmp(value.size()); @@ -213,8 +214,10 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp); add_kv(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_chexp); - add_kv(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp); - add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp); + add_kv(LLM_KV_SWIGLU_CLAMP_EXP, std::vector( + hparams.swiglu_clamp_exp.begin(), hparams.swiglu_clamp_exp.begin() + hparams.n_layer_all)); + add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, std::vector( + hparams.swiglu_clamp_shexp.begin(), hparams.swiglu_clamp_shexp.begin() + hparams.n_layer_all)); add_kv(LLM_KV_USE_PARALLEL_RESIDUAL, hparams.use_par_res); // add_kv(LLM_KV_TENSOR_DATA_LAYOUT, ???); add_kv(LLM_KV_EXPERT_COUNT, hparams.n_expert); @@ -314,6 +317,8 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_SSM_DT_B_C_RMS, hparams.ssm_dt_b_c_rms); add_kv(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda); + add_kv(LLM_KV_KDA_SAFE_GATE, hparams.kda_safe_gate); + add_kv(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound); add_kv(LLM_KV_WKV_HEAD_SIZE, hparams.wkv_head_size); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index eaf3f35d2d8..947b9efd3b8 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -242,6 +242,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_bailingmoe(params); case LLM_ARCH_BAILINGMOE2: return new llama_model_bailingmoe2(params); + case LLM_ARCH_BAILINGMOE3: + return new llama_model_bailingmoe3(params); case LLM_ARCH_SEED_OSS: return new llama_model_seed_oss(params); case LLM_ARCH_DOTS1: @@ -787,6 +789,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_A13B: return "A13B"; case LLM_TYPE_7B_A1B: return "7B.A1B"; case LLM_TYPE_8B_A1B: return "8B.A1B"; + case LLM_TYPE_7_9B_A1_3B: return "7.9B.A1.3B"; case LLM_TYPE_12B_A2_5B: return "12B.A2.5B"; case LLM_TYPE_16B_A1B: return "16B.A1B"; case LLM_TYPE_21B_A3B: return "21B.A3B"; @@ -802,6 +805,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_106B_A12B: return "106B.A12B"; case LLM_TYPE_120B_A12B: return "120B.A12B"; case LLM_TYPE_122B_A10B: return "122B.A10B"; + case LLM_TYPE_124B_A5_1B: return "124B.A5.1B"; case LLM_TYPE_196B_A11B: return "196B.A11B"; case LLM_TYPE_230B_A10B: return "230B.A10B"; case LLM_TYPE_235B_A22B: return "235B.A22B"; @@ -1902,7 +1906,7 @@ void llama_model::print_info() const { LLAMA_LOG_INFO("%s: expert_weights_norm = %d\n", __func__, hparams.expert_weights_norm); } - if (arch == LLM_ARCH_BAILINGMOE2) { + if (arch == LLM_ARCH_BAILINGMOE2 || arch == LLM_ARCH_BAILINGMOE3) { LLAMA_LOG_INFO("%s: n_layer_dense_lead = %d\n", __func__, hparams.n_layer_dense_lead); LLAMA_LOG_INFO("%s: n_ff_exp = %d\n", __func__, hparams.n_ff_exp); LLAMA_LOG_INFO("%s: n_ff_shexp = %d\n", __func__, hparams.n_ff_shexp); @@ -2072,9 +2076,11 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, { // The MTP head is dense-attention only on hybrid Qwen3.5/3.6, so use a plain // attention KV cache for the MTP context instead of the hybrid wrapper. + // Dense MTP heads use a plain attention KV cache instead of the hybrid wrapper. const bool mtp_on_hybrid_qwen35 = params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && - (arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE); + (arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || + arch == LLM_ARCH_BAILINGMOE3); if (llm_arch_is_recurrent(arch)) { res = new llama_memory_recurrent( @@ -2462,6 +2468,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_GRANITE_HYBRID: case LLM_ARCH_CHAMELEON: case LLM_ARCH_BAILINGMOE: + case LLM_ARCH_BAILINGMOE3: case LLM_ARCH_NEO_BERT: case LLM_ARCH_SMOLLM3: case LLM_ARCH_ARCEE: diff --git a/src/llama-model.h b/src/llama-model.h index 45b054cedf1..c2fc4a2609d 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -117,6 +117,7 @@ enum llm_type { LLM_TYPE_A13B, LLM_TYPE_7B_A1B, LLM_TYPE_8B_A1B, // lfm2moe + LLM_TYPE_7_9B_A1_3B, // Ling-3.0-tiny LLM_TYPE_12B_A2_5B, LLM_TYPE_16B_A1B, LLM_TYPE_21B_A3B, // Ernie MoE small @@ -132,6 +133,7 @@ enum llm_type { LLM_TYPE_106B_A12B, // GLM-4.5-Air LLM_TYPE_120B_A12B, // Nemotron 3 Super LLM_TYPE_122B_A10B, // Qwen3.5 + LLM_TYPE_124B_A5_1B, // Ling-3.0-flash LLM_TYPE_196B_A11B, // Step3.5-Flash LLM_TYPE_230B_A10B, // Minimax M2 LLM_TYPE_235B_A22B, diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index b66759b2776..0051ad88ca9 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -462,7 +462,12 @@ static ggml_type llama_tensor_get_type_impl(quantize_state_impl & qs, ggml_type } else if (ftype == LLAMA_FTYPE_MOSTLY_MXFP4_MOE) { // MoE tensors -> MXFP4 // other tensors -> Q8_0 - if (tensor->ne[2] > 1) { + // MLA projection tensors are also 3D, so match expert tensor roles explicitly. + const bool is_bailingmoe3_expert = arch == LLM_ARCH_BAILINGMOE3 && + (category == tensor_category::FFN_UP || + category == tensor_category::FFN_GATE || + category == tensor_category::FFN_DOWN); + if (tensor->ne[2] > 1 && (arch != LLM_ARCH_BAILINGMOE3 || is_bailingmoe3_expert)) { new_type = GGML_TYPE_MXFP4; } else { new_type = GGML_TYPE_Q8_0; diff --git a/src/models/bailingmoe3.cpp b/src/models/bailingmoe3.cpp new file mode 100644 index 00000000000..f0a5437af83 --- /dev/null +++ b/src/models/bailingmoe3.cpp @@ -0,0 +1,530 @@ +#include "models.h" +#include "llama-memory-recurrent.h" + +void llama_model_bailingmoe3::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl); + ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv); + ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q, false); + ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); + ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda); + if (!ml.get_key(LLM_KV_KDA_SAFE_GATE, hparams.kda_safe_gate, false)) { + hparams.kda_safe_gate = true; + } + ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound); + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func); + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); + ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all, false); + ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, false); + + if (hparams.n_ff_shexp == 0) { + hparams.n_ff_shexp = hparams.n_ff_exp * std::max(1u, hparams.n_expert_shared); + } + + GGML_ASSERT(hparams.kda_safe_gate); + GGML_ASSERT(hparams.kda_gate_lower_bound < 0.0f); + + for (uint32_t il = 0; il < hparams.n_layer(); ++il) { + hparams.is_recr_impl[il] = hparams.n_head_kv(il) == 0; + } + + switch (hparams.n_layer()) { + case 24: type = hparams.n_embd == 1536 && hparams.n_expert == 128 ? LLM_TYPE_7_9B_A1_3B : LLM_TYPE_UNKNOWN; break; + case 42: type = hparams.n_embd == 2560 && hparams.n_expert == 512 ? LLM_TYPE_124B_A5_1B : LLM_TYPE_UNKNOWN; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_bailingmoe3::load_arch_tensors(llama_model_loader & ml) { + LLAMA_LOAD_LOCALS; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0); + + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED); + if (output == nullptr) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED); + } + + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_inner = head_dim * n_head; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t kv_lora_rank = hparams.n_lora_kv; + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t qk_rope_head_dim = hparams.n_rot(); + const int64_t qk_head_dim = hparams.n_embd_head_k_mla(); + const int64_t v_head_dim = hparams.n_embd_head_v_mla(); + + const bool mtp_only = (hparams.n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr); + const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight"; + const bool trunk_only = (hparams.n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr); + const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; + int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0; + + // In this fork ml.load_mtp is absent; trunk_only already sets TENSOR_SKIP via mtp_flags. + + for (int il = 0; il < n_layer; ++il) { + auto & layer = layers[il]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, trunk_flags); + + if (hparams.is_recr(il)) { + layer.ssm_q_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_Q, "weight", il), { d_conv, 1, d_inner, 1 }, trunk_flags); + layer.ssm_k_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_K, "weight", il), { d_conv, 1, d_inner, 1 }, trunk_flags); + layer.ssm_v_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_V, "weight", il), { d_conv, 1, d_inner, 1 }, trunk_flags); + + create_tensor_qkv(layer, il, n_embd, d_inner, d_inner, d_inner, trunk_flags); + layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", il), { n_embd, d_inner }, trunk_flags); + layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_head }, trunk_flags); + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, il), { 1, n_head }, trunk_flags); + layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { d_inner }, trunk_flags); + layer.ssm_g_a = create_tensor(tn(LLM_TENSOR_SSM_G_A, "weight", il), { n_embd, d_inner }, trunk_flags); + layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_dim }, trunk_flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { d_inner, n_embd }, trunk_flags); + } else { + if (q_lora_rank > 0) { + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", il), { n_embd, q_lora_rank }, trunk_flags); + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", il), { q_lora_rank }, trunk_flags); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", il), { q_lora_rank, n_head * qk_head_dim }, trunk_flags); + } else { + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", il), { n_embd, n_head * qk_head_dim }, trunk_flags); + } + 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); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", il), { kv_lora_rank }, trunk_flags); + 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); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", il), { kv_lora_rank, v_head_dim, n_head }, trunk_flags); + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, n_head }, trunk_flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_head * v_head_dim, n_embd }, trunk_flags); + } + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", il), { n_embd }, trunk_flags); + if ((uint32_t) il < hparams.n_layer_dense_lead) { + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", il), { n_embd, n_ff }, trunk_flags); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", il), { n_embd, n_ff }, trunk_flags); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", il), { n_ff, n_embd }, trunk_flags); + } else { + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, trunk_flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", il), { n_expert }, trunk_flags); + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", il), { n_embd, hparams.n_ff_exp, n_expert }, trunk_flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", il), { n_embd, hparams.n_ff_exp, n_expert }, trunk_flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { hparams.n_ff_exp, n_embd, n_expert }, trunk_flags); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, trunk_flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, trunk_flags); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { hparams.n_ff_shexp, n_embd }, trunk_flags); + } + } + + for (int il = n_layer; il < n_layer_all; ++il) { + auto & layer = layers[il]; + const int flags = mtp_flags; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), { n_embd }, flags); + if (q_lora_rank > 0) { + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", il), { n_embd, q_lora_rank }, flags); + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", il), { q_lora_rank }, flags); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", il), { q_lora_rank, n_head * qk_head_dim }, flags); + } else { + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", il), { n_embd, n_head * qk_head_dim }, flags); + } + 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); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", il), { kv_lora_rank }, flags); + 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); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", il), { kv_lora_rank, v_head_dim, n_head }, flags); + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, n_head }, flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_head * v_head_dim, n_embd }, flags); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", il), { n_embd }, flags); + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", il), { n_expert }, flags); + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", il), { n_embd, hparams.n_ff_exp, n_expert }, flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", il), { n_embd, hparams.n_ff_exp, n_expert }, flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { hparams.n_ff_exp, n_embd, n_expert }, flags); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, hparams.n_ff_shexp }, flags); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { hparams.n_ff_shexp, n_embd }, flags); + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, flags); + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, flags); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { n_embd }, flags); + layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_LAYER_OUT_NORM, "weight", il), { n_embd }, flags); + } +} + +std::unique_ptr llama_model_bailingmoe3::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } + return std::make_unique(*this, params); +} + +static ggml_tensor * bailingmoe3_causal_conv1d( + ggml_cgraph * gf, + ggml_context * ctx0, + ggml_tensor * conv_states_all, + ggml_tensor * conv_state_all, + int64_t qkv, + ggml_tensor * x, + ggml_tensor * proj_w, + ggml_tensor * conv_w, + int64_t d_conv, + int64_t head_dim, + int64_t n_head, + int64_t n_seq_tokens, + int64_t n_seqs, + int64_t n_tokens, + int64_t cache_head) { + const int64_t d_inner = head_dim * n_head; + const int64_t conv_state_size = (d_conv - 1) * d_inner; + const int64_t total_state_size = 3 * conv_state_size; + + ggml_tensor * conv_state = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_state_all), + total_state_size * ggml_element_size(conv_state_all), + qkv * conv_state_size * ggml_element_size(conv_state_all)); + + ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x); + x_proj = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs); + ggml_tensor * conv_x = ggml_concat(ctx0, conv_state, ggml_transpose(ctx0, x_proj), 0); + + ggml_tensor * last_conv_x = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs, + conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, last_conv_x, + ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_states_all), + total_state_size * ggml_element_size(conv_states_all), + (cache_head * total_state_size + qkv * conv_state_size) * ggml_element_size(conv_states_all)))); + + ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner); + ggml_tensor * out = ggml_ssm_conv(ctx0, conv_x, conv_weight); + out = ggml_silu(ctx0, ggml_reshape_2d(ctx0, out, d_inner, n_tokens)); + return ggml_reshape_4d(ctx0, out, head_dim, n_head, n_seq_tokens, n_seqs); +} + +llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_build_delta_net_base(params), model(model) { + ggml_tensor * inpL = build_inp_embd(model.tok_embd); + cb(inpL, "model.input_embed", -1); + + auto * inp = build_inp_mem_hybrid_k(); + auto * inp_rs = inp->get_recr(); + auto * inp_attn = inp->get_attn(); + + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + const int64_t n_head = hparams.n_head(); + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_inner = n_head * head_dim; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + const int64_t qk_head_dim = hparams.n_embd_head_k_mla(); + const int64_t v_head_dim = hparams.n_embd_head_v_mla(); + const int64_t qk_rope_head_dim = hparams.n_rot(); + const int64_t qk_nope_head_dim = qk_head_dim - qk_rope_head_dim; + const int64_t kv_lora_rank = hparams.n_lora_kv; + const float kq_scale = 1.0f / sqrtf((float) qk_head_dim); + + GGML_ASSERT(n_seqs > 0); + GGML_ASSERT(ubatch.equal_seqs()); + GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + ggml_tensor * inpSA = inpL; + ggml_tensor * cur = build_norm(inpL, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + if (hparams.is_recr(il)) { + const auto * mctx_cur = inp_rs->mctx; + const auto cache_head = mctx_cur->get_head(); + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs); + + ggml_tensor * q = bailingmoe3_causal_conv1d( + gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head); + ggml_tensor * k = bailingmoe3_causal_conv1d( + gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head); + ggml_tensor * v = bailingmoe3_causal_conv1d( + gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head); + + ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_f_a, cur); + gate = ggml_add(ctx0, gate, layer.ssm_dt_b); + gate = ggml_reshape_3d(ctx0, gate, head_dim, n_head, n_tokens); + ggml_tensor * a = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head, 1); + gate = ggml_scale(ctx0, ggml_sigmoid(ctx0, ggml_mul(ctx0, gate, a)), hparams.kda_gate_lower_bound); + gate = ggml_reshape_4d(ctx0, gate, head_dim, n_head, n_seq_tokens, n_seqs); + cb(gate, "kda_gate", il); + + ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur); + beta = ggml_sigmoid(ctx0, ggml_reshape_4d(ctx0, beta, 1, n_head, n_seq_tokens, n_seqs)); + + q = ggml_l2_norm(ctx0, q, hparams.f_norm_rms_eps); + k = ggml_l2_norm(ctx0, k, hparams.f_norm_rms_eps); + + ggml_tensor * states_all = mctx_cur->get_s_l(il); + ggml_tensor * state = build_rs(inp_rs, states_all, hparams.n_embd_s(), n_seqs); + state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head, n_seqs); + + auto result = build_delta_net(q, k, v, gate, beta, state, il); + ggml_tensor * out = ggml_cont(ctx0, result.first); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, result.second, + ggml_view_1d(ctx0, states_all, hparams.n_embd_s() * n_seqs, + cache_head * hparams.n_embd_s() * ggml_element_size(states_all)))); + + ggml_tensor * out_gate = ggml_mul_mat(ctx0, layer.ssm_g_a, cur); + out_gate = ggml_reshape_3d(ctx0, out_gate, head_dim, n_head, n_tokens); + out = ggml_reshape_3d(ctx0, out, head_dim, n_head, n_tokens); + out = build_norm(out, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il); + out = ggml_mul(ctx0, out, ggml_sigmoid(ctx0, out_gate)); + cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, out, d_inner, n_tokens)); + cb(cur, "kda_out", il); + } else { + ggml_tensor * attn_input = cur; + ggml_tensor * q_all; + if (layer.wq_a) { + q_all = ggml_mul_mat(ctx0, layer.wq_a, cur); + cb(q_all, "q_a", il); + q_all = build_norm(q_all, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + cb(q_all, "q_a_norm", il); + q_all = ggml_mul_mat(ctx0, layer.wq_b, q_all); + cb(q_all, "q_b", il); + } else { + q_all = ggml_mul_mat(ctx0, layer.wq, cur); + } + ggml_tensor * q_nope = ggml_view_3d(ctx0, q_all, qk_nope_head_dim, n_head, n_tokens, + ggml_row_size(q_all->type, qk_head_dim), + ggml_row_size(q_all->type, qk_head_dim) * n_head, 0); + ggml_tensor * q_pe = ggml_view_3d(ctx0, q_all, qk_rope_head_dim, n_head, n_tokens, + ggml_row_size(q_all->type, qk_head_dim), + ggml_row_size(q_all->type, qk_head_dim) * n_head, + ggml_row_size(q_all->type, qk_nope_head_dim)); + + ggml_tensor * kv_all = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); + ggml_tensor * kv = ggml_view_2d(ctx0, kv_all, kv_lora_rank, n_tokens, + ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), 0); + ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_all, qk_rope_head_dim, 1, n_tokens, + ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), + ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), + ggml_row_size(kv_all->type, kv_lora_rank)); + + q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + kv = build_norm(kv, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + q_nope = ggml_mul_mat(ctx0, layer.wk_b, q_nope); + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + + ggml_tensor * q = ggml_concat(ctx0, q_nope, q_pe, 0); + kv = ggml_reshape_3d(ctx0, kv, kv_lora_rank, 1, n_tokens); + ggml_tensor * k = ggml_concat(ctx0, kv, k_pe, 0); + + cur = build_attn(inp_attn, nullptr, nullptr, nullptr, + q, k, kv, nullptr, nullptr, layer.wv_b, kq_scale, il); + + ggml_tensor * attn_gate = ggml_mul_mat(ctx0, layer.wqkv_gate, attn_input); + attn_gate = ggml_sigmoid(ctx0, ggml_reshape_3d(ctx0, attn_gate, 1, n_head, n_tokens)); + cur = ggml_reshape_3d(ctx0, cur, v_head_dim, n_head, n_tokens); + cur = ggml_mul(ctx0, cur, attn_gate); + cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, cur, v_head_dim * n_head, n_tokens)); + cb(cur, "mla_out", il); + } + + if (il == n_layer - 1 && inp_out_ids && cparams.embeddings_nextn_masked) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + + if ((uint32_t) il < hparams.n_layer_dense_lead) { + cur = build_ffn(cur, + layer.ffn_up, nullptr, nullptr, + layer.ffn_gate, nullptr, nullptr, + layer.ffn_down, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + } else { + ggml_tensor * moe = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_SILU, + hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il); + ggml_tensor * shared = build_ffn(cur, + layer.ffn_up_shexp, nullptr, nullptr, + layer.ffn_gate_shexp, nullptr, nullptr, + layer.ffn_down_shexp, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + cur = ggml_add(ctx0, moe, shared); + } + + cur = ggml_add(ctx0, cur, ffn_inp); + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + inpL = cur; + } + + ggml_tensor * cur = build_norm(inpL, model.output_norm, nullptr, LLM_NORM_RMS, -1); + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + if (!cparams.embeddings_nextn_masked && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + ggml_build_forward_expand(gf, cur); +} + +llama_model_bailingmoe3::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) : + llm_graph_context(params) { + GGML_ASSERT(hparams.n_layer_nextn == 1 && "BailingMoE3 MTP requires one NextN layer"); + + const int il = hparams.n_layer() + cparams.nextn_layer_offset; + GGML_ASSERT(cparams.nextn_layer_offset >= 0 && + cparams.nextn_layer_offset < (int) hparams.n_layer_nextn && + "nextn_layer_offset out of range"); + const auto & layer = model.layers[il]; + + GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj"); + GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm"); + GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm"); + GGML_ASSERT(layer.nextn.shared_head_norm && "MTP block missing final norm"); + + const int64_t n_head = hparams.n_head(); + const int64_t qk_head_dim = hparams.n_embd_head_k_mla(); + const int64_t v_head_dim = hparams.n_embd_head_v_mla(); + const int64_t qk_rope_head_dim = hparams.n_rot(); + const int64_t qk_nope_head_dim = qk_head_dim - qk_rope_head_dim; + const int64_t kv_lora_rank = hparams.n_lora_kv; + const float kq_scale = 1.0f / sqrtf((float) qk_head_dim); + + auto inp = std::make_unique(hparams.n_embd); + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->tokens); + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens); + ggml_set_input(inp->embd); + ggml_set_name(inp->embd, "mtp_h_input"); + + ggml_tensor * tok_embd = ggml_get_rows(ctx0, model.tok_embd, inp->tokens); + ggml_tensor * h_norm = build_norm(inp->embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il); + ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il); + ggml_tensor * cur = ggml_mul_mat(ctx0, layer.nextn.eh_proj, ggml_concat(ctx0, e_norm, h_norm, 0)); + cb(cur, "mtp_eh_proj", il); + + res->add_input(std::move(inp)); + + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + auto * inp_attn = build_attn_inp_k(); + + ggml_tensor * inpSA = cur; + cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + ggml_tensor * attn_input = cur; + + ggml_tensor * q_all; + if (layer.wq_a) { + q_all = ggml_mul_mat(ctx0, layer.wq_a, cur); + cb(q_all, "q_a", il); + q_all = build_norm(q_all, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + cb(q_all, "q_a_norm", il); + q_all = ggml_mul_mat(ctx0, layer.wq_b, q_all); + cb(q_all, "q_b", il); + } else { + q_all = ggml_mul_mat(ctx0, layer.wq, cur); + } + ggml_tensor * q_nope = ggml_view_3d(ctx0, q_all, qk_nope_head_dim, n_head, n_tokens, + ggml_row_size(q_all->type, qk_head_dim), + ggml_row_size(q_all->type, qk_head_dim) * n_head, 0); + ggml_tensor * q_pe = ggml_view_3d(ctx0, q_all, qk_rope_head_dim, n_head, n_tokens, + ggml_row_size(q_all->type, qk_head_dim), + ggml_row_size(q_all->type, qk_head_dim) * n_head, + ggml_row_size(q_all->type, qk_nope_head_dim)); + + ggml_tensor * kv_all = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); + ggml_tensor * kv = ggml_view_2d(ctx0, kv_all, kv_lora_rank, n_tokens, + ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), 0); + ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_all, qk_rope_head_dim, 1, n_tokens, + ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), + ggml_row_size(kv_all->type, kv_lora_rank + qk_rope_head_dim), + ggml_row_size(kv_all->type, kv_lora_rank)); + + q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + kv = build_norm(kv, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + q_nope = ggml_mul_mat(ctx0, layer.wk_b, q_nope); + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + + ggml_tensor * q = ggml_concat(ctx0, q_nope, q_pe, 0); + kv = ggml_reshape_3d(ctx0, kv, kv_lora_rank, 1, n_tokens); + ggml_tensor * k = ggml_concat(ctx0, kv, k_pe, 0); + + cur = build_attn(inp_attn, nullptr, nullptr, nullptr, + q, k, kv, nullptr, nullptr, layer.wv_b, kq_scale, il); + + ggml_tensor * attn_gate = ggml_mul_mat(ctx0, layer.wqkv_gate, attn_input); + attn_gate = ggml_sigmoid(ctx0, ggml_reshape_3d(ctx0, attn_gate, 1, n_head, n_tokens)); + cur = ggml_reshape_3d(ctx0, cur, v_head_dim, n_head, n_tokens); + cur = ggml_mul(ctx0, cur, attn_gate); + cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, cur, v_head_dim * n_head, n_tokens)); + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + + ggml_tensor * moe = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_SILU, + hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il); + ggml_tensor * shared = build_ffn(cur, + layer.ffn_up_shexp, nullptr, nullptr, + layer.ffn_gate_shexp, nullptr, nullptr, + layer.ffn_down_shexp, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + cur = ggml_add(ctx0, moe, shared); + cur = ggml_add(ctx0, cur, ffn_inp); + cur = build_norm(cur, layer.nextn.shared_head_norm, nullptr, LLM_NORM_RMS, -1); + + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/models.h b/src/models/models.h index beab9f6bc7d..a9e85859b0f 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1628,6 +1628,25 @@ struct llama_model_bailingmoe2 : public llama_model_base { }; +struct llama_model_bailingmoe3 : public llama_model_base { + llama_model_bailingmoe3(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_build_delta_net_base { + graph(const llama_model & model, const llm_graph_params & params); + + const llama_model & model; + }; + + struct graph_mtp : public llm_graph_context { + graph_mtp(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_seed_oss : public llama_model_base { llama_model_seed_oss(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 2cdf3573982..14f6495f774 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -103,6 +103,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { || arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_KIMI_LINEAR + || arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_MISTRAL4) { n_embd = 128; n_head = 1; @@ -141,7 +142,8 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_FULL_ATTENTION_INTERVAL, uint32_t(2)); if (arch == LLM_ARCH_PLAMO2 || arch == LLM_ARCH_JAMBA || arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE || - arch == LLM_ARCH_GRANITE_HYBRID || arch == LLM_ARCH_LFM2 || arch == LLM_ARCH_LFM2MOE || arch == LLM_ARCH_KIMI_LINEAR) { + arch == LLM_ARCH_GRANITE_HYBRID || arch == LLM_ARCH_LFM2 || arch == LLM_ARCH_LFM2MOE || arch == LLM_ARCH_KIMI_LINEAR || + arch == LLM_ARCH_BAILINGMOE3) { GGML_ASSERT(n_layer >= 2); std::vector n_head_per_layer; n_head_per_layer.reserve(n_layer); @@ -160,6 +162,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { || arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_KIMI_LINEAR + || arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_MISTRAL4) { ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, uint32_t(576)); ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, uint32_t(512)); @@ -229,6 +232,12 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_SSM_TIME_STEP_RANK, n_head); ms.add_kv(LLM_KV_SSM_GROUP_COUNT, arch == LLM_ARCH_PLAMO2 ? 0 : uint32_t(2)); ms.add_kv(LLM_KV_KDA_HEAD_DIM, uint32_t(128)); + ms.add_kv(LLM_KV_KDA_SAFE_GATE, true); + ms.add_kv(LLM_KV_KDA_GATE_LOWER_BOUND, -5.0f); + if (arch == LLM_ARCH_BAILINGMOE3) { + ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP, std::vector({0.0f, 4.0f})); + ms.add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, std::vector({0.0f, 5.0f})); + } ms.add_kv(LLM_KV_WKV_HEAD_SIZE, n_embd/n_head); ms.add_kv(LLM_KV_SHORTCONV_L_CACHE, uint32_t(3)); @@ -341,6 +350,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_EXAONE_MOE: case LLM_ARCH_BAILINGMOE: case LLM_ARCH_BAILINGMOE2: + case LLM_ARCH_BAILINGMOE3: case LLM_ARCH_DOTS1: case LLM_ARCH_AFMOE: case LLM_ARCH_ERNIE4_5: @@ -575,6 +585,9 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg } const std::string config_name = moe ? "MoE" : "Dense"; gguf_context_ptr gguf_ctx = get_gguf_ctx(arch, moe); + if (arch == LLM_ARCH_BAILINGMOE3) { + GGML_ASSERT(gguf_remove_key(gguf_ctx.get(), "bailingmoe3.kda.safe_gate") >= 0); + } std::pair model_and_ctx_cpu; std::vector logits_cpu; for (device_config & dc : dev_configs) {