From: Ruixiang Wang Date: Mon, 10 Aug 2026 08:25:24 +0000 (+0200) Subject: model: add MTP support for Nemotron model (#26725) X-Git-Tag: upstream/0.0.10438~94 X-Git-Url: https://git.djapps.eu/?a=commitdiff_plain;h=7a20b417f4526cae073bd997af5020cea3e7ccbe;p=pkg%2Fggml%2Fsources%2Fllama.cpp model: add MTP support for Nemotron model (#26725) * model: add MTP support for Nemotron Nano model * model: add mtp_flags for nemotron model * address review comments --- diff --git a/conversion/nemotron.py b/conversion/nemotron.py index 0572b42ca..e5075020c 100644 --- a/conversion/nemotron.py +++ b/conversion/nemotron.py @@ -197,6 +197,7 @@ class NemotronHModel(GraniteHybridModel): """Hybrid mamba2/attention model from NVIDIA""" model_arch = gguf.MODEL_ARCH.NEMOTRON_H is_moe: bool = False + supports_mtp_export = True def __init__(self, *args, **kwargs): # We have to determine the correct model architecture (MoE vs non-MoE) before @@ -236,6 +237,25 @@ class NemotronHModel(GraniteHybridModel): self._ssm_layers = [i for i, val in enumerate(pattern) if val == "mamba"] self._mlp_layers = [i for i, val in enumerate(pattern) if val == "moe"] + # `--no-mtp` drops it entirely; `--mtp` exports only the MTP head + self._mtp_bid: int | None = None + if self.is_moe and not self.no_mtp: + n_nextn = self.hparams.get("num_nextn_predict_layers", 0) or 0 + if n_nextn > 0: + assert n_nextn == 1, ( + "NemotronH MTP conversion currently supports num_nextn_predict_layers == 1" + ) + self._mtp_bid = self.block_count + self.block_count += 1 + # The folded MTP block carries both an attention sub-layer and a + # MoE sub-layer, so register it as both so the per-layer metadata arrays cover it + self._attn_layers.append(self._mtp_bid) + self._mlp_layers.append(self._mtp_bid) + self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) + + if self.mtp_only and self._mtp_bid is None: + raise ValueError("--mtp was requested, but this model does not contain a supported MTP head") + def get_attn_layers(self): pattern = self.hparams.get("hybrid_override_pattern") or self.hparams.get("layers_block_type") if pattern is None: @@ -246,6 +266,36 @@ class NemotronHModel(GraniteHybridModel): return [i for i, val in enumerate(pattern) if val == "attention"] + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + if name.startswith("mtp."): + # --no-mtp: drop the MTP head entirely + if cls.no_mtp: + return None + elif cls.mtp_only: + # --mtp: export the MTP head plus the tensors it shares with the target model + keep = name in ( + "backbone.embeddings.weight", + "backbone.norm_f.weight", + "lm_head.weight", + ) + if not keep: + return None + return super().filter_tensors((name, gen)) + + def prepare_metadata(self, vocab_only: bool): + from_dir = self.fname_out.is_dir() + super().prepare_metadata(vocab_only=vocab_only) + + if not self.mtp_only or not from_dir: + return + output_type: str = self.ftype.name.partition("_")[2] + fname_default: str = gguf.naming_convention( + self.metadata.name, self.metadata.basename, self.metadata.finetune, + self.metadata.version, size_label=None, output_type=output_type, model_type=None) + self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf" + def set_gguf_parameters(self): super().set_gguf_parameters() @@ -284,6 +334,10 @@ class NemotronHModel(GraniteHybridModel): if (latent_size := self.hparams.get("moe_latent_size")) is not None: self.gguf_writer.add_moe_latent_size(latent_size) + # MTP head: number of trailing NextN blocks + if self._mtp_bid is not None: + self.gguf_writer.add_nextn_predict_layers(self.hparams["num_nextn_predict_layers"]) + def set_vocab(self): # The NemotronH config uses pattern characters (e.g. '-') that may not # be supported by the installed transformers version. AutoTokenizer @@ -350,15 +404,24 @@ class NemotronHModel(GraniteHybridModel): if not self.is_moe: self.gguf_writer.add_add_bos_token(True) + _MTP_SPECIAL_RENAMES = { + "mtp.layers.0.enorm.weight": "model.layers.{bid}.enorm.weight", + "mtp.layers.0.hnorm.weight": "model.layers.{bid}.hnorm.weight", + "mtp.layers.0.eh_proj.weight": "model.layers.{bid}.eh_proj.weight", + "mtp.layers.1.norm.weight": "model.layers.{bid}.post_attention_layernorm.weight", + "mtp.layers.1.final_layernorm.weight": "model.layers.{bid}.shared_head.norm.weight", + } + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: - if self.is_moe and bid is not None: - # Skip Multi-Token Prediction (MTP) tensors. These are used for - # for speculative decoding but we don't include them in this model - # conversion. See https://github.com/ggml-org/llama.cpp/pull/18886 - if name.startswith("mtp."): - logger.info(f"gguf: Skipping MTP (Speculative) layer: {name}") - return + # mtp.layers.0: NextN input fusion + attention + # mtp.layers.1: MoE + final head norm + if self._mtp_bid is not None and name.startswith(("mtp.layers.0.", "mtp.layers.1.")): + suffix = name.split(".", 3)[3] + bid = self._mtp_bid + renamed = self._MTP_SPECIAL_RENAMES.get(name) + name = renamed.format(bid=bid) if renamed else f"backbone.layers.{bid}.{suffix}" + if self.is_moe and bid is not None: if name.endswith("mixer.gate.e_score_correction.bias"): yield from ModelBase.modify_tensors(self, data_torch, name, bid) return diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 304100f7f..e6740287f 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -3846,6 +3846,12 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_DOWN_SHEXP, MODEL_TENSOR.FFN_UP_SHEXP, MODEL_TENSOR.FFN_EXP_PROBS_B, + # NextN/MTP (draft head) + MODEL_TENSOR.ATTN_POST_NORM, + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], MODEL_ARCH.EXAONE: [ MODEL_TENSOR.TOKEN_EMBD, diff --git a/src/llama-model.cpp b/src/llama-model.cpp index b4575b82a..9316636d6 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2231,6 +2231,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE); + const bool mtp_on_hybrid_nemotron = + params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && arch == LLM_ARCH_NEMOTRON_H_MOE; + if (llm_arch_is_recurrent(arch)) { res = new llama_memory_recurrent( *this, @@ -2241,7 +2244,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, cparams.n_seq_max, cparams.n_rs_seq, nullptr); - } else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen) { + } else if (llm_arch_is_hybrid(arch) && !mtp_on_hybrid_qwen && !mtp_on_hybrid_nemotron) { // The main difference between hybrid architectures is the // layer filters, so pick the right one here llama_memory_hybrid::layer_filter_cb filter_attn = nullptr; @@ -2322,7 +2325,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, }; } - if (mtp_on_hybrid_qwen) { + if (mtp_on_hybrid_qwen || mtp_on_hybrid_nemotron) { filter = [&](uint32_t il) { return il >= hparams.n_layer(); }; } diff --git a/src/models/models.h b/src/models/models.h index a8908da42..12412ef53 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1461,6 +1461,10 @@ struct llama_model_nemotron_h_moe : public llama_model_nemotron_h { using graph = llama_model_nemotron_h::graph; + 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; }; diff --git a/src/models/nemotron-h-moe.cpp b/src/models/nemotron-h-moe.cpp index a59cc6c9f..4d03f49e0 100644 --- a/src/models/nemotron-h-moe.cpp +++ b/src/models/nemotron-h-moe.cpp @@ -1,6 +1,156 @@ #include "models.h" std::unique_ptr llama_model_nemotron_h_moe::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); } +// MTP draft head for Nemotron-H MoE +llama_model_nemotron_h_moe::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) + : llm_graph_context(params) { + GGML_ASSERT(hparams.n_layer_nextn == 1 && "NEMOTRON_H_MOE MTP currently supports a single MTP block"); + + const int64_t n_embd_head = hparams.n_embd_head_v(); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + + const int il = hparams.n_layer(); + const auto & layer = model.layers[il]; + + GGML_ASSERT(layer.nextn.eh_proj && layer.nextn.enorm && layer.nextn.hnorm); + GGML_ASSERT(layer.ffn_gate_inp); + + // token embedding weights + ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd; + GGML_ASSERT(tok_embd_w != nullptr && "NEMOTRON_H_MOE MTP requires token embeddings"); + + 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_inp(), n_tokens); + ggml_set_input(inp->embd); + + ggml_tensor * tok_embd; + if (ubatch.token) { + tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens); + } else { + tok_embd = inp->embd; + } + cb(tok_embd, "mtp_tok_embd", il); + + inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens); + ggml_set_input(inp->h); + ggml_set_name(inp->h, "mtp_h_input"); + + ggml_tensor * h_embd = inp->h; + + res->add_input(std::move(inp)); + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + // attention fills KV over all tokens, but the MoE is position-wise: gather output rows before + // it to save FFN compute (unless unmasked embeddings_nextn needs the full-length hidden state) + const bool emit_h_nextn = cparams.embeddings_nextn; + const bool crop_before_ffn = inp_out_ids && (!emit_h_nextn || cparams.embeddings_nextn_masked); + + auto * inp_attn = build_attn_inp_kv(); + + ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il); + cb(h_norm, "mtp_hnorm", il); + ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il); + cb(e_norm, "mtp_enorm", il); + + ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0); + cb(concat, "mtp_concat", il); + + ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s); + cb(cur, "mtp_eh_proj", il); + + // dense NoPE attention sub-layer (mtp.layers.0) + ggml_tensor * inpSA = cur; + cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_attn_norm", il); + + { + auto [Qcur, Kcur, Vcur] = build_qkv(layer, cur, n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il); + const float kq_scale = hparams.f_attention_scale == 0.0f + ? 1.0f / sqrtf(float(n_embd_head)) : hparams.f_attention_scale; + cur = build_attn(inp_attn, layer.wo, layer.wo_b, layer.wo_s, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + cb(cur, "mtp_attn_out", il); + } + + cur = ggml_add(ctx0, cur, inpSA); + cb(cur, "mtp_attn_residual", il); + + // gather the output rows here so the MoE FFN below only runs on the positions we keep + if (crop_before_ffn) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + + // MoE FFN sub-layer (mtp.layers.1) + ggml_tensor * ffn_residual = cur; + cur = build_norm(cur, layer.attn_post_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_attn_post_norm", il); + + { + ggml_tensor * router_logits = build_lora_mm(layer.ffn_gate_inp, cur); + cb(router_logits, "mtp_ffn_moe_logits", il); + + ggml_tensor * moe_out = + build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + nullptr, // no gate + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_RELU_SQR, hparams.expert_weights_norm, + hparams.expert_weights_scale, + LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID, + il, + router_logits, nullptr, + layer.ffn_up_exps_s, + nullptr, // no gate + 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); + + cur = ggml_add(ctx0, moe_out, ffn_shexp); + cb(cur, "mtp_ffn_out", il); + } + + cur = ggml_add(ctx0, cur, ffn_residual); + cb(cur, "mtp_post_ffn", il); + + // final head norm: the MTP head has its own LayerNorm + GGML_ASSERT(layer.nextn.shared_head_norm && "NEMOTRON_H_MOE MTP: missing final head norm"); + cur = build_norm(cur, layer.nextn.shared_head_norm, nullptr, LLM_NORM, -1); + + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + if (!crop_before_ffn && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + + // LM head + ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output; + ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s; + GGML_ASSERT(head_w != nullptr && "NEMOTRON_H_MOE MTP requires an output projection"); + cur = build_lora_mm(head_w, cur, head_s); + cb(cur, "result_output", -1); + + res->t_logits = cur; + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/nemotron-h.cpp b/src/models/nemotron-h.cpp index a45626934..cd2af3179 100644 --- a/src/models/nemotron-h.cpp +++ b/src/models/nemotron-h.cpp @@ -7,13 +7,18 @@ void llama_model_nemotron_h::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank); ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group); + // NextN/MTP: optional draft head appended as extra trailing block(s) + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); + GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < n_layer_all"); + // A layer is recurrent IFF the n_head_kv value is set to 0 and - // the n_ff value is set to 0 - for (uint32_t i = 0; i < hparams.n_layer(); ++i) { - hparams.is_recr_impl[i] = (hparams.n_head_kv(i) == 0 && hparams.n_ff(i) == 0); + // the n_ff value is set to 0. Appended MTP blocks are dense (non-recurrent) + for (uint32_t i = 0; i < hparams.n_layer_all; ++i) { + hparams.is_recr_impl[i] = i < hparams.n_layer() && hparams.n_head_kv(i) == 0 && hparams.n_ff(i) == 0; } ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps); // MTP head final_layernorm ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false); ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false); @@ -30,9 +35,13 @@ void llama_model_nemotron_h::load_arch_hparams(llama_model_loader & ml) { } } -void llama_model_nemotron_h::load_arch_tensors(llama_model_loader &) { +void llama_model_nemotron_h::load_arch_tensors(llama_model_loader & ml) { LLAMA_LOAD_LOCALS; + const bool mtp_only = hparams.n_layer_nextn > 0 && ml.get_weight("blk.0.attn_norm.weight") == nullptr; + const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; + const int mtp_flags = !ml.load_mtp ? TENSOR_SKIP : 0; + // mamba2 Mixer SSM params // NOTE: int64_t for tensor dimensions const int64_t d_conv = hparams.ssm_d_conv; @@ -60,61 +69,94 @@ void llama_model_nemotron_h::load_arch_tensors(llama_model_loader &) { auto & layer = layers[i]; // all blocks use the attn norm - layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, trunk_flags); if (hparams.is_recr(i)) { // ssm layers - layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), {n_embd, d_in_proj}, 0); + layer.ssm_in = create_tensor(tn(LLM_TENSOR_SSM_IN, "weight", i), {n_embd, d_in_proj}, trunk_flags); - layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {d_conv, d_inner + 2*n_group*d_state}, 0); + layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {d_conv, d_inner + 2*n_group*d_state}, trunk_flags); layer.ssm_conv1d_b = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "bias", i), {d_inner + 2*n_group*d_state}, TENSOR_NOT_REQUIRED); - layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {n_ssm_head}, 0); + layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {n_ssm_head}, trunk_flags); // no "weight" suffix for these - layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_ssm_head}, 0); - layer.ssm_d = create_tensor(tn(LLM_TENSOR_SSM_D, i), {1, n_ssm_head}, 0); + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_ssm_head}, trunk_flags); + layer.ssm_d = create_tensor(tn(LLM_TENSOR_SSM_D, i), {1, n_ssm_head}, trunk_flags); - layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {d_inner / n_group, n_group}, 0); + layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {d_inner / n_group, n_group}, trunk_flags); // out_proj - layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), {d_inner, n_embd}, 0); + layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", i), {d_inner, n_embd}, trunk_flags); } else if (hparams.n_ff(i) == 0) { // attention layers (with optional bias) const int64_t n_head_i = hparams.n_head(i); const int64_t n_embd_k_gqa_i = hparams.n_embd_k_gqa(i); const int64_t n_embd_v_gqa_i = hparams.n_embd_v_gqa(i); - create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, 0); - layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, 0); + create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, trunk_flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, trunk_flags); layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED); } else { if (n_expert != 0) { const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used; const int64_t n_ff_shexp = hparams.n_ff_shexp; - layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert}, 0); - layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert }, 0); + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), { n_embd, n_expert}, trunk_flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert }, trunk_flags); // MoE branch layer.ffn_latent_down = create_tensor(tn(LLM_TENSOR_FFN_LATENT_DOWN, "weight", i), {n_embd, moe_n_embd}, TENSOR_NOT_REQUIRED); layer.ffn_latent_up = create_tensor(tn(LLM_TENSOR_FFN_LATENT_UP, "weight", i), {moe_n_embd, n_embd}, 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}, 0); - layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, moe_n_embd, n_expert}, trunk_flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {moe_n_embd, n_ff_exp, n_expert}, trunk_flags); // Shared expert branch - layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, 0); - layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, trunk_flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, trunk_flags); } else { // mlp layers - layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { hparams.n_ff(i), n_embd}, 0); - layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, hparams.n_ff(i)}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { hparams.n_ff(i), n_embd}, trunk_flags); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, hparams.n_ff(i)}, trunk_flags); layer.ffn_down_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED); layer.ffn_up_b = create_tensor(tn(LLM_TENSOR_FFN_UP, "bias", i), {hparams.n_ff(i)}, TENSOR_NOT_REQUIRED); } } } + + // NextN/MTP draft head: each predict layer folds an attention sub-layer and a MoE + // sub-layer into a single trailing block + for (int i = n_layer; i < n_layer_all; ++i) { + auto & layer = layers[i]; + + const int64_t n_head_i = hparams.n_head(i); + const int64_t n_embd_k_gqa_i = hparams.n_embd_k_gqa(i); + const int64_t n_embd_v_gqa_i = hparams.n_embd_v_gqa(i); + const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used; + const int64_t n_ff_shexp = hparams.n_ff_shexp; + + // NextN input-fusion tensors + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, mtp_flags); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, mtp_flags); + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2*n_embd, n_embd}, mtp_flags); + layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, mtp_flags); + + // attention sub-layer + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, mtp_flags); + create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head_i, n_embd_k_gqa_i, n_embd_v_gqa_i, mtp_flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head_i, n_embd}, mtp_flags); + layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, mtp_flags | TENSOR_NOT_REQUIRED); + + // MoE sub-layer + 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_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); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, mtp_flags); + } } std::unique_ptr llama_model_nemotron_h::build_arch_graph(const llm_graph_params & params) const { @@ -153,7 +195,7 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_ cur = build_ffn_layer(cur, model, il); } - if (il == n_layer - 1 && inp_out_ids) { + 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); } @@ -170,6 +212,14 @@ llama_model_nemotron_h::graph::graph(const llama_model & model, const llm_graph_ cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); + // seed for the MTP/NextN draft head + 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;