From e74c17e153311460fb195ad58fe98fd31300d803 Mon Sep 17 00:00:00 2001 From: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Date: Sun, 22 Mar 2026 17:13:58 +0200 Subject: [PATCH] Enable `NemotronHPuzzle` + `NemotronHMTP` (#37803) --- vllm/config/speculative.py | 2 +- vllm/model_executor/models/nemotron_h_mtp.py | 7 +++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index a4a48888a..8ff6d9753 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -280,7 +280,7 @@ class SpeculativeConfig: ) if ( - hf_config.model_type == "nemotron_h" + hf_config.model_type in {"nemotron_h", "nemotron_h_puzzle"} and hasattr(hf_config, "num_nextn_predict_layers") and hf_config.num_nextn_predict_layers > 0 ): diff --git a/vllm/model_executor/models/nemotron_h_mtp.py b/vllm/model_executor/models/nemotron_h_mtp.py index dcadf1f33..12551d425 100644 --- a/vllm/model_executor/models/nemotron_h_mtp.py +++ b/vllm/model_executor/models/nemotron_h_mtp.py @@ -395,13 +395,16 @@ class NemotronHMTP(nn.Module, SupportsPP): ] expert_params_mapping = [] - if hasattr(self.config, "n_routed_experts") and self.config.n_routed_experts: + num_experts = getattr(self.config, "n_routed_experts", None) + if getattr(self.config, "model_type", None) == "nemotron_h_puzzle": + num_experts = self.config.mtp_n_routed_experts + if num_experts is not None: expert_params_mapping = FusedMoE.make_expert_params_mapping( self, ckpt_gate_proj_name="up_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="", # Empty - non-gated MoE - num_experts=self.config.n_routed_experts, + num_experts=num_experts, num_redundant_experts=self.num_redundant_experts, )