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
This commit is contained in:
Ruixiang Wang
2026-08-10 10:25:24 +02:00
committed by GitHub
parent e23e9440eb
commit 7a20b417f4
6 changed files with 309 additions and 33 deletions
+5 -2
View File
@@ -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(); };
}