diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 5e61f61..9861e32 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -441,6 +441,12 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_FFN_GATE_EXP, "blk.%d.ffn_gate.%d" }, { LLM_TENSOR_FFN_DOWN_EXP, "blk.%d.ffn_down.%d" }, { LLM_TENSOR_FFN_UP_EXP, "blk.%d.ffn_up.%d" }, + { LLM_TENSOR_FFN_GATE_EXPS_TIER1, "blk.%d.ffn_gate_exps_tier1" }, + { LLM_TENSOR_FFN_UP_EXPS_TIER1, "blk.%d.ffn_up_exps_tier1" }, + { LLM_TENSOR_FFN_DOWN_EXPS_TIER1, "blk.%d.ffn_down_exps_tier1" }, + { LLM_TENSOR_FFN_GATE_EXPS_TIER2, "blk.%d.ffn_gate_exps_tier2" }, + { LLM_TENSOR_FFN_UP_EXPS_TIER2, "blk.%d.ffn_up_exps_tier2" }, + { LLM_TENSOR_FFN_DOWN_EXPS_TIER2, "blk.%d.ffn_down_exps_tier2" }, { LLM_TENSOR_FFN_GATE_EXPS, "blk.%d.ffn_gate_exps" }, { LLM_TENSOR_FFN_GATE_UP_EXPS, "blk.%d.ffn_gate_up_exps" }, { LLM_TENSOR_FFN_DOWN_EXPS, "blk.%d.ffn_down_exps" }, @@ -889,6 +895,12 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_DEC_ATTN_REL_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, {LLM_TENSOR_ENC_ATTN_REL_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, {LLM_TENSOR_FFN_DOWN_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_GATE_EXPS_TIER1, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_UP_EXPS_TIER1, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_DOWN_EXPS_TIER1, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_GATE_EXPS_TIER2, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_UP_EXPS_TIER2, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, + {LLM_TENSOR_FFN_DOWN_EXPS_TIER2, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, {LLM_TENSOR_FFN_GATE_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, {LLM_TENSOR_FFN_UP_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, {LLM_TENSOR_FFN_GATE_UP_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT_ID}}, diff --git a/src/llama-arch.h b/src/llama-arch.h index ca7d55a..9026666 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -469,6 +469,12 @@ enum llm_tensor { LLM_TENSOR_FFN_DOWN_EXPS, // merged experts LLM_TENSOR_FFN_GATE_EXPS, LLM_TENSOR_FFN_UP_EXPS, + LLM_TENSOR_FFN_GATE_EXPS_TIER1, + LLM_TENSOR_FFN_UP_EXPS_TIER1, + LLM_TENSOR_FFN_DOWN_EXPS_TIER1, + LLM_TENSOR_FFN_GATE_EXPS_TIER2, + LLM_TENSOR_FFN_UP_EXPS_TIER2, + LLM_TENSOR_FFN_DOWN_EXPS_TIER2, LLM_TENSOR_FFN_GATE_UP_EXPS, LLM_TENSOR_FFN_DOWN_SHEXP, LLM_TENSOR_FFN_GATE_SHEXP, diff --git a/src/llama-model.h b/src/llama-model.h index 0d7352a..024952b 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -329,6 +329,12 @@ struct llama_layer { struct ggml_tensor * ffn_gate_exps = nullptr; struct ggml_tensor * ffn_down_exps = nullptr; struct ggml_tensor * ffn_up_exps = nullptr; + struct ggml_tensor * ffn_gate_exps_tier1 = nullptr; + struct ggml_tensor * ffn_up_exps_tier1 = nullptr; + struct ggml_tensor * ffn_down_exps_tier1 = nullptr; + struct ggml_tensor * ffn_gate_exps_tier2 = nullptr; + struct ggml_tensor * ffn_up_exps_tier2 = nullptr; + struct ggml_tensor * ffn_down_exps_tier2 = nullptr; struct ggml_tensor * ffn_gate_up_exps = nullptr; struct ggml_tensor * ffn_gate_inp_b = nullptr; struct ggml_tensor * ffn_gate_exps_b = nullptr; diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index acfdd5b..1fcb5c7 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -198,9 +198,32 @@ void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { layer.ple_conv1d = create_tensor(tn(LLM_TENSOR_PLE_CONV1D, "weight", il), { hparams.ple_conv_kernel, hc_dim }, 0); } - layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, 0); - layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, 0); - create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, 0); + const auto * t1 = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_GATE_EXPS_TIER1, "weight", il).str().c_str()); + const bool mixed = t1 != nullptr; + const auto * router_meta = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il).str().c_str()); + if (!router_meta) throw std::runtime_error("missing expert router tensor"); + const int64_t layer_experts = mixed ? router_meta->ne[1] : n_expert; + if (layer_experts < n_expert_used || layer_experts > n_expert) throw std::runtime_error("invalid layer expert count"); + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, layer_experts }, 0); + if (mixed) { + const auto * t0 = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", il).str().c_str()); + const auto * t2 = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_GATE_EXPS_TIER2, "weight", il).str().c_str()); + if (!t0 || !t2 || n_ff_exp != 640) throw std::runtime_error("missing mixed-tier expert tensors"); + const int64_t n0 = t0->ne[2], n1 = t1->ne[2], n2 = t2->ne[2]; + if (n0 < 1 || n0 > 64 || n1 < 1 || n1 > 64 || n2 < 1 || n2 > 64 || n0+n1+n2 != layer_experts) + throw std::runtime_error("invalid mixed-tier bank sizes"); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n0 }, 0); + create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n0, 0); + layer.ffn_gate_exps_tier1 = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS_TIER1, "weight", il), { n_embd, n_ff_exp, n1 }, 0); + layer.ffn_up_exps_tier1 = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS_TIER1, "weight", il), { n_embd, n_ff_exp, n1 }, 0); + layer.ffn_down_exps_tier1 = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS_TIER1, "weight", il), { 768, n_embd, n1 }, 0); + layer.ffn_gate_exps_tier2 = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS_TIER2, "weight", il), { n_embd, n_ff_exp, n2 }, 0); + layer.ffn_up_exps_tier2 = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS_TIER2, "weight", il), { n_embd, n_ff_exp, n2 }, 0); + layer.ffn_down_exps_tier2 = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS_TIER2, "weight", il), { n_ff_exp, n_embd, n2 }, 0); + } else { + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, 0); + create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, 0); + } layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, 0); layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0); @@ -904,8 +927,52 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear( ggml_tensor * llama_model_qwen4exp::graph::build_layer_ffn(ggml_tensor * cur, const int il) { GGML_ASSERT(model.layers[il].ffn_gate_inp != nullptr); - ggml_tensor * moe_out = - build_moe_ffn(cur, + ggml_tensor * moe_out = nullptr; + if (model.layers[il].ffn_gate_exps_tier1 != nullptr) { + const auto & layer = model.layers[il]; + const int64_t n_tokens = cur->ne[1]; + ggml_tensor * logits = build_lora_mm(layer.ffn_gate_inp, cur); + ggml_tensor * probs = ggml_soft_max(ctx0, logits); + cb(probs, "ffn_moe_probs", il); + ggml_tensor * ids = ggml_argsort_top_k(ctx0, probs, n_expert_used); + cb(ids, "ffn_moe_topk", il); + ggml_tensor * selected = ggml_get_rows(ctx0, ggml_reshape_3d(ctx0, probs, 1, layer.ffn_gate_inp->ne[1], n_tokens), ids); + selected = ggml_reshape_2d(ctx0, selected, n_expert_used, n_tokens); + ggml_tensor * denom = ggml_clamp(ctx0, ggml_sum_rows(ctx0, selected), 6.103515625e-5f, INFINITY); + selected = ggml_div(ctx0, selected, denom); + if (hparams.expert_weights_scale != 0.0f && hparams.expert_weights_scale != 1.0f) { + selected = ggml_scale(ctx0, selected, hparams.expert_weights_scale); + } + ggml_tensor * ids_f = ggml_cast(ctx0, ids, GGML_TYPE_F32); + ggml_tensor * inp = ggml_reshape_3d(ctx0, cur, cur->ne[0], 1, n_tokens); + ggml_tensor * weighted = nullptr; + ggml_tensor * ups[] = {layer.ffn_up_exps, layer.ffn_up_exps_tier1, layer.ffn_up_exps_tier2}; + ggml_tensor * gates[] = {layer.ffn_gate_exps, layer.ffn_gate_exps_tier1, layer.ffn_gate_exps_tier2}; + ggml_tensor * downs[] = {layer.ffn_down_exps, layer.ffn_down_exps_tier1, layer.ffn_down_exps_tier2}; + int64_t tier_start = 0; + for (int tier = 0; tier < 3; ++tier) { + const int64_t tier_size = ups[tier]->ne[2]; + const float base = float(tier_start); + ggml_tensor * local_f = ggml_scale_bias(ctx0, ids_f, 1.0f, -base); + ggml_tensor * local = ggml_cast(ctx0, ggml_clamp(ctx0, local_f, 0.0f, float(tier_size - 1)), GGML_TYPE_I32); + ggml_tensor * low = ggml_step(ctx0, ggml_scale_bias(ctx0, local_f, 1.0f, 0.5f)); + ggml_tensor * high = ggml_step(ctx0, ggml_scale_bias(ctx0, local_f, -1.0f, float(tier_size) - 0.5f)); + tier_start += tier_size; + ggml_tensor * mask = ggml_mul(ctx0, low, high); + ggml_tensor * weights = ggml_reshape_3d(ctx0, ggml_mul(ctx0, selected, mask), 1, n_expert_used, n_tokens); + ggml_tensor * up = build_lora_mm_id(ups[tier], inp, local); + ggml_tensor * gate = build_lora_mm_id(gates[tier], inp, local); + ggml_tensor * act = ggml_swiglu_split(ctx0, gate, up); + if (tier == 1) act = ggml_pad(ctx0, act, 128, 0, 0, 0); + ggml_tensor * result = ggml_mul(ctx0, build_lora_mm_id(downs[tier], act, local), weights); + weighted = weighted ? ggml_add(ctx0, weighted, result) : result; + } + for (int i = 0; i < n_expert_used; ++i) { + ggml_tensor * piece = ggml_view_2d(ctx0, weighted, cur->ne[0], n_tokens, weighted->nb[2], i*weighted->nb[1]); + moe_out = moe_out ? ggml_add(ctx0, moe_out, piece) : piece; + } + } else { + moe_out = build_moe_ffn(cur, model.layers[il].ffn_gate_inp, model.layers[il].ffn_up_exps, model.layers[il].ffn_gate_exps, @@ -919,6 +986,7 @@ ggml_tensor * llama_model_qwen4exp::graph::build_layer_ffn(ggml_tensor * cur, co model.layers[il].ffn_up_exps_s, model.layers[il].ffn_gate_exps_s, model.layers[il].ffn_down_exps_s); + } cb(moe_out, "ffn_moe_out", il); // shared experts, as in the Qwen3Next reference