From 0abe6a85a51f2b7f1c3ca0e8f78944b609b94344 Mon Sep 17 00:00:00 2001 From: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> Date: Thu, 28 May 2026 02:12:44 +0300 Subject: [PATCH] Support NemotronHPuzzleForCausalLM (#24429) Signed-off-by: Netanel Haber <58652339+netanel-haber@users.noreply.github.com> --- python/sglang/srt/configs/__init__.py | 3 +- python/sglang/srt/configs/model_config.py | 5 +- python/sglang/srt/configs/nemotron_h.py | 53 +++++++++++++++++++ python/sglang/srt/models/nemotron_h.py | 15 ++++-- python/sglang/srt/models/nemotron_h_mtp.py | 1 + python/sglang/srt/server_args.py | 2 +- .../srt/utils/hf_transformers/common.py | 2 + 7 files changed, 74 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py index 23e908186..c9043f7ea 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/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", diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index eed82eee5..713dab9e1 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/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 diff --git a/python/sglang/srt/configs/nemotron_h.py b/python/sglang/srt/configs/nemotron_h.py index 710765a5f..916a5f46c 100644 --- a/python/sglang/srt/configs/nemotron_h.py +++ b/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 diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index 1840b0727..3b903c3b3 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/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"]) diff --git a/python/sglang/srt/models/nemotron_h_mtp.py b/python/sglang/srt/models/nemotron_h_mtp.py index dabd4b4ae..a88d2b806 100644 --- a/python/sglang/srt/models/nemotron_h_mtp.py +++ b/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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 5c1710e89..3f758a8b3 100644 --- a/python/sglang/srt/server_args.py +++ b/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, ) diff --git a/python/sglang/srt/utils/hf_transformers/common.py b/python/sglang/srt/utils/hf_transformers/common.py index 777acf2a4..0fd841888 100644 --- a/python/sglang/srt/utils/hf_transformers/common.py +++ b/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,