Feat/support nemotron h mtp (#17013)

Signed-off-by: Shahar Mor <smor@nvidia.com>
2026-01-14 16:30:35 +08:00
parent b025cff441
commit ba625c2d90
4 changed files with 378 additions and 1 deletions
@@ -300,6 +300,10 @@ class ModelConfig:
            self.hf_config.architectures[0] = "Qwen3NextForCausalLMMTP"
            self.hf_config.num_nextn_predict_layers = 1

+        if is_draft_model and self.hf_config.architectures[0] == "NemotronHForCausalLM":
+            self.hf_config.architectures[0] = "NemotronHForCausalLMMTP"
+            self.hf_config.num_nextn_predict_layers = 1
+
    def _derive_hybrid_model(self):
        # Use self.context_len after it has been initialized to prevent using context_len which may be None.
        self.is_hybrid_swa = is_hybrid_swa_model(self.hf_config.architectures)