Support NemotronHPuzzleForCausalLM (#24429)

Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com>

python/sglang/srt/configs/__init__.py

@@ -26,7 +26,7 @@ from sglang.srt.configs.nano_nemotron_vl import (
     NemotronH_Nano_Omni_Reasoning_V3_Config,
     NemotronH_Nano_VL_V2_Config,
 )
-from sglang.srt.configs.nemotron_h import NemotronHConfig
+from sglang.srt.configs.nemotron_h import NemotronHConfig, NemotronHPuzzleConfig
 from sglang.srt.configs.olmo3 import Olmo3Config
 from sglang.srt.configs.qwen3_5 import Qwen3_5Config, Qwen3_5MoeConfig
 from sglang.srt.configs.qwen3_asr import Qwen3ASRConfig
@@ -70,6 +70,7 @@ __all__ = [
     "MiniCPMV4_6Config",
     "MiniCPMV4_6VisionConfig",
     "NemotronHConfig",
+    "NemotronHPuzzleConfig",
     "NemotronH_Nano_VL_V2_Config",
     "NemotronH_Nano_Omni_Reasoning_V3_Config",
     "JetNemotronConfig",

python/sglang/srt/configs/model_config.py

@@ -480,7 +480,10 @@ class ModelConfig:
             self.hf_config.architectures[0] = "ExaoneMoEForCausalLMMTP"
             self.hf_config.num_nextn_predict_layers = 1
 
-        if is_draft_model and self.hf_config.architectures[0] == "NemotronHForCausalLM":
+        if is_draft_model and self.hf_config.architectures[0] in [
+            "NemotronHForCausalLM",
+            "NemotronHPuzzleForCausalLM",
+        ]:
             self.hf_config.architectures[0] = "NemotronHForCausalLMMTP"
             self.hf_config.num_nextn_predict_layers = 1
 

python/sglang/srt/configs/nemotron_h.py

@@ -17,6 +17,9 @@
 
 """NemotronH model configuration"""
 
+import copy
+from typing import Any
+
 from transformers.configuration_utils import PretrainedConfig
 from transformers.utils import logging
 
@@ -240,6 +243,7 @@ class NemotronHConfig(PretrainedConfig):
 
     def __init__(
         self,
+        *,
         vocab_size=131072,
         tie_word_embeddings=False,
         hidden_size=4096,
@@ -504,3 +508,52 @@ class NemotronHConfig(PretrainedConfig):
             MLP: "mlp",
         }
         return [pattern_mapping[char] for char in pattern]
+
+    def get_nemotron_h_config_for_layer(self, layer_idx: int) -> "NemotronHConfig":
+        return self
+
+    def get_mtp_config(self) -> "NemotronHConfig":
+        return self
+
+    @property
+    def max_n_routed_experts(self) -> int:
+        return self.n_routed_experts
+
+
+class NemotronHPuzzleConfig(NemotronHConfig):
+    model_type = "nemotron_h_puzzle"
+    has_no_defaults_at_init = True
+
+    def __init__(
+        self,
+        *,
+        block_configs: list[dict[str, Any]],
+        mtp_block_configs: list[dict[str, Any]] | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.block_configs = block_configs
+        self.mtp_block_configs = mtp_block_configs
+
+    def get_nemotron_h_config_for_layer(self, layer_idx: int) -> NemotronHConfig:
+        layer_config = copy.copy(self)
+        for key, value in self.block_configs[layer_idx].items():
+            setattr(layer_config, key, value)
+        return layer_config
+
+    def get_mtp_config(self) -> NemotronHConfig:
+        assert self.mtp_block_configs
+        mtp_config = copy.copy(self)
+        mtp_config.block_configs = self.mtp_block_configs
+        return mtp_config
+
+    @property
+    def max_n_routed_experts(self) -> int:
+        block_n_routed_experts = [
+            block["n_routed_experts"]
+            for block in self.block_configs
+            if block["block_type"] == "moe"
+        ]
+        max_experts = max(block_n_routed_experts)
+        assert max_experts > 0
+        return max_experts

python/sglang/srt/models/nemotron_h.py

@@ -358,9 +358,10 @@ class NemotronHMoEDecoderLayer(nn.Module):
         prefix: str = "",
     ) -> None:
         super().__init__()
+        layer_config = config.get_nemotron_h_config_for_layer(layer_idx)
 
         self.mixer = NemotronHMoE(
-            config,
+            layer_config,
             layer_idx=layer_idx,
             quant_config=quant_config,
             prefix=f"{prefix}.mixer",
@@ -510,6 +511,7 @@ class NemotronHAttention(nn.Module):
             self.scaling,
             num_kv_heads=self.num_kv_heads,
             layer_id=layer_idx,
+            sliding_window_size=config.sliding_window,
             quant_config=quant_config,
             prefix=add_prefix("attn", prefix),
         )
@@ -533,9 +535,10 @@ class NemotronHAttentionDecoderLayer(nn.Module):
         prefix: str = "",
     ) -> None:
         super().__init__()
+        layer_config = config.get_nemotron_h_config_for_layer(layer_idx)
 
         self.mixer = NemotronHAttention(
-            config,
+            layer_config,
             layer_idx,
             quant_config,
             prefix=f"{prefix}.mixer",
@@ -904,7 +907,7 @@ class NemotronHForCausalLM(nn.Module):
             ckpt_gate_proj_name="up_proj",
             ckpt_down_proj_name="down_proj",
             ckpt_up_proj_name="",
-            num_experts=self.config.n_routed_experts,
+            num_experts=self.config.max_n_routed_experts,
         )
 
         params_dict = dict(self.named_parameters())
@@ -1004,7 +1007,11 @@ class NemotronHForCausalLM(nn.Module):
                         logger.warning(f"Parameter {name} not found in params_dict")
 
 
-EntryClass = [NemotronHForCausalLM]
+class NemotronHPuzzleForCausalLM(NemotronHForCausalLM):
+    pass
+
+
+EntryClass = [NemotronHForCausalLM, NemotronHPuzzleForCausalLM]
 
 
 @register_custom_op(mutates_args=["output"])

python/sglang/srt/models/nemotron_h_mtp.py

@@ -283,6 +283,7 @@ class NemotronHForCausalLMMTP(NemotronHForCausalLM):
         prefix: str = "",
     ):
         nn.Module.__init__(self)
+        config = config.get_mtp_config()
         self.config = config
         self.quant_config = quant_config
         # Required for parent's load_weights

python/sglang/srt/server_args.py

@@ -2360,7 +2360,7 @@ class ServerArgs:
                 support_mamba_cache=True,
                 support_mamba_cache_extra_buffer=True,
             )
-        elif model_arch in ["NemotronHForCausalLM"]:
+        elif model_arch in ["NemotronHForCausalLM", "NemotronHPuzzleForCausalLM"]:
             from sglang.srt.arg_groups.nemotron_h_hook import (
                 apply_nemotron_h_defaults,
             )

python/sglang/srt/utils/hf_transformers/common.py

@@ -46,6 +46,7 @@ from sglang.srt.configs import (
     NemotronH_Nano_Omni_Reasoning_V3_Config,
     NemotronH_Nano_VL_V2_Config,
     NemotronHConfig,
+    NemotronHPuzzleConfig,
     Olmo3Config,
     Qwen3_5Config,
     Qwen3_5MoeConfig,
@@ -96,6 +97,7 @@ _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
         NemotronH_Nano_VL_V2_Config,
         NemotronH_Nano_Omni_Reasoning_V3_Config,
         NemotronHConfig,
+        NemotronHPuzzleConfig,
         DeepseekVLV2Config,
         Qwen3_5Config,
         Qwen3_5MoeConfig,