diff --git a/src/models/nemotron-h-moe.cpp b/src/models/nemotron-h-moe.cpp index 4d03f49e0f..b4fb254300 100644 --- a/src/models/nemotron-h-moe.cpp +++ b/src/models/nemotron-h-moe.cpp @@ -100,6 +100,18 @@ llama_model_nemotron_h_moe::graph_mtp::graph_mtp(const llama_model & model, cons ggml_tensor * router_logits = build_lora_mm(layer.ffn_gate_inp, cur); cb(router_logits, "mtp_ffn_moe_logits", il); + ggml_tensor * ffn_shexp = build_ffn(cur, + layer.ffn_up_shexp, NULL, layer.ffn_up_shexp_s, + NULL, NULL, NULL, + layer.ffn_down_shexp, NULL, layer.ffn_down_shexp_s, + NULL, + LLM_FFN_RELU_SQR, LLM_FFN_PAR, il); + cb(ffn_shexp, "mtp_ffn_shexp", il); + + if (layer.ffn_latent_down) { + cur = ggml_mul_mat(ctx0, layer.ffn_latent_down, cur); + } + ggml_tensor * moe_out = build_moe_ffn(cur, layer.ffn_gate_inp, @@ -118,13 +130,9 @@ llama_model_nemotron_h_moe::graph_mtp::graph_mtp(const llama_model & model, cons layer.ffn_down_exps_s); cb(moe_out, "mtp_ffn_moe_out", il); - ggml_tensor * ffn_shexp = build_ffn(cur, - layer.ffn_up_shexp, NULL, layer.ffn_up_shexp_s, - NULL, NULL, NULL, - layer.ffn_down_shexp, NULL, layer.ffn_down_shexp_s, - NULL, - LLM_FFN_RELU_SQR, LLM_FFN_PAR, il); - cb(ffn_shexp, "mtp_ffn_shexp", il); + if (layer.ffn_latent_up) { + moe_out = ggml_mul_mat(ctx0, layer.ffn_latent_up, moe_out); + } cur = ggml_add(ctx0, moe_out, ffn_shexp); cb(cur, "mtp_ffn_out", il); diff --git a/src/models/nemotron-h.cpp b/src/models/nemotron-h.cpp index 24ed9a673e..be27650b06 100644 --- a/src/models/nemotron-h.cpp +++ b/src/models/nemotron-h.cpp @@ -172,6 +172,8 @@ void llama_model_nemotron_h::load_arch_tensors(llama_model_loader & ml) { layer.attn_post_norm = create_tensor(tn(LLM_TENSOR_ATTN_POST_NORM, "weight", i), {n_embd}, mtp_flags); layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, mtp_flags); layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, mtp_flags); + layer.ffn_latent_down = create_tensor(tn(LLM_TENSOR_FFN_LATENT_DOWN, "weight", i), {n_embd, moe_n_embd}, mtp_flags | TENSOR_NOT_REQUIRED); + layer.ffn_latent_up = create_tensor(tn(LLM_TENSOR_FFN_LATENT_UP, "weight", i), {moe_n_embd, n_embd}, mtp_flags | TENSOR_NOT_REQUIRED); layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, moe_n_embd, n_expert}, mtp_flags); layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, mtp_flags); layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, mtp_flags);