Skip to content

feat(bailingmoe3): add bailingmoe3 arch support #32

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
oscar1229 wants to merge 1 commit into
base: mtmd-backend
Choose a base branch
from
Open

feat(bailingmoe3): add bailingmoe3 arch support #32

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion src/llama-arch.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,7 @@ static const std::map<llm_arch, const char *> 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" },
Expand Down Expand Up @@ -294,7 +295,9 @@ static const std::map<llm_kv, const char *> 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" },

Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down
3 changes: 3 additions & 0 deletions src/llama-arch.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,

Expand Down
1 change: 1 addition & 0 deletions src/llama-context.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
2 changes: 2 additions & 0 deletions src/llama-hparams.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
9 changes: 7 additions & 2 deletions src/llama-model-saver.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<std::vector<uint32_t>>(const enum llm_kv, const std::vector<uint32_t> &, const bool);
template void llama_model_saver::add_kv<std::vector<float>>(const enum llm_kv, const std::vector<float> &, const bool);

void llama_model_saver::add_kv(const enum llm_kv key, const std::vector<std::string> & value) {
std::vector<const char *> tmp(value.size());
Expand Down Expand Up @@ -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<float>(
hparams.swiglu_clamp_exp.begin(), hparams.swiglu_clamp_exp.begin() + hparams.n_layer_all));
add_kv(LLM_KV_SWIGLU_CLAMP_SHEXP, std::vector<float>(
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);
Expand Down Expand Up @@ -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);

Expand Down
11 changes: 9 additions & 2 deletions src/llama-model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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";
Expand All @@ -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";
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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:
Expand Down
2 changes: 2 additions & 0 deletions src/llama-model.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand Down
7 changes: 6 additions & 1 deletion src/llama-quant.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Loading
Loading