model : extend Nemotron MTP support (#29018)

* first fix

* removed unnecessary declarations
This commit is contained in:
Kartik Gulia
2026-09-17 13:52:36 +02:00
committed by GitHub
parent ebbb185227
commit 7076180486
2 changed files with 17 additions and 7 deletions
+15 -7
View File
@@ -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);
+2
View File
@@ -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);