Support NemotronH_Omni_Reasoning_V3 in SGLang

python/sglang/srt/arg_groups/model_overrides/nemotron_h.py

@@ -1,6 +1,7 @@
 """Config-time override declarations for nemotron_h.
 
-Architectures: NemotronHForCausalLM, NemotronHPuzzleForCausalLM.
+Architectures: NemotronHForCausalLM, NemotronHPuzzleForCausalLM,
+NemotronH_Omni_Reasoning_V3.
 """
 
 import logging
@@ -17,7 +18,11 @@
 logger = logging.getLogger(__name__)
 
 
-@_register_for("NemotronHForCausalLM", "NemotronHPuzzleForCausalLM")
+@_register_for(
+    "NemotronHForCausalLM",
+    "NemotronHPuzzleForCausalLM",
+    "NemotronH_Omni_Reasoning_V3",
+)
 def _nemotron_h_overrides(server_args: Any, hf_config: Any) -> dict:
     """NemotronH quantization / MoE runner / attention backend defaults
     (absorbed from the retired arg_groups/nemotron_h_hook.py; the mamba radix
@@ -35,7 +40,12 @@ def _nemotron_h_overrides(server_args: Any, hf_config: Any) -> dict:
     ]
     quantization = cfg.quantization
     if is_modelopt:
-        assert model_config.hf_config.mlp_hidden_act == "relu2"
+        language_config = (
+            model_config.hf_text_config
+            if model_arch == "NemotronH_Omni_Reasoning_V3"
+            else hf_config
+        )
+        assert language_config.mlp_hidden_act == "relu2"
         if model_config.quantization == "modelopt":
             quant_algo = model_config.hf_config.quantization_config["quant_algo"]
             if quant_algo == "MIXED_PRECISION":

python/sglang/srt/arg_groups/overrides.py

@@ -456,6 +456,7 @@ def _step3p_overrides(server_args: Any, hf_config: Any) -> dict:
         "MiniCPMV4_6ForConditionalGeneration",
         "NemotronHForCausalLM",
         "NemotronHPuzzleForCausalLM",
+        "NemotronH_Omni_Reasoning_V3",
         "FalconH1ForCausalLM",
         "JetNemotronForCausalLM",
         "JetVLMForConditionalGeneration",
@@ -488,6 +489,7 @@ def _step3p_overrides(server_args: Any, hf_config: Any) -> dict:
         "Glm5NextForConditionalGeneration",
         "NemotronHForCausalLM",
         "NemotronHPuzzleForCausalLM",
+        "NemotronH_Omni_Reasoning_V3",
         # KDA-based: same MambaPool ping-pong machinery as GDN; requires the
         # KDA backend's track-snapshot writes (decode + extend) so donated
         # slots hold real states for prefix-cache restores.
@@ -974,6 +976,7 @@ def _sparse_head_overlap_disable(view: Any) -> dict:
         "Qwen3_5ForConditionalGeneration",
         "NemotronHForCausalLM",
         "NemotronHPuzzleForCausalLM",
+        "NemotronH_Omni_Reasoning_V3",
     }
 )
 

python/sglang/srt/configs/__init__.py

@@ -57,6 +57,7 @@
 from sglang.srt.configs.nano_nemotron_vl import (
     NemotronH_Nano_Omni_Reasoning_V3_Config,
     NemotronH_Nano_VL_V2_Config,
+    NemotronH_Omni_Reasoning_V3_Config,
 )
 from sglang.srt.configs.nemotron_h import NemotronHConfig, NemotronHPuzzleConfig
 from sglang.srt.configs.olmo3 import Olmo3Config
@@ -137,6 +138,7 @@
     "NemotronHPuzzleConfig",
     "NemotronH_Nano_VL_V2_Config",
     "NemotronH_Nano_Omni_Reasoning_V3_Config",
+    "NemotronH_Omni_Reasoning_V3_Config",
     "NanbeigeConfig",
     "JetNemotronConfig",
     "JetVLMConfig",

python/sglang/srt/configs/model_config.py

@@ -892,6 +892,14 @@ def _config_draft_model(self):
             self.hf_config.architectures[0] = "ExaoneMoEForCausalLMMTP"
             self.hf_config.num_nextn_predict_layers = 1
 
+        if (
+            is_draft_model
+            and self.hf_config.architectures[0] == "NemotronH_Omni_Reasoning_V3"
+        ):
+            self.hf_config = self.hf_text_config
+            self.hf_config.architectures = ["NemotronHForCausalLMMTP"]
+            self.hf_config.num_nextn_predict_layers = 1
+
         if is_draft_model and self.hf_config.architectures[0] in [
             "NemotronHForCausalLM",
             "NemotronHPuzzleForCausalLM",
@@ -2071,6 +2079,7 @@ def is_generation_model(model_architectures: List[str], is_embedding: bool = Fal
     "MossVLForConditionalGeneration",
     "NemotronH_Nano_VL_V2",
     "NemotronH_Nano_Omni_Reasoning_V3",
+    "NemotronH_Omni_Reasoning_V3",
     "MuseGlimmerForConditionalGeneration",
     "PixtralForConditionalGeneration",
     "Qwen2AudioForConditionalGeneration",

python/sglang/srt/configs/nano_nemotron_vl.py

@@ -30,6 +30,15 @@ def float_triplet(seq: Any):
     return a, b, c
 
 
+def _nemotron_h_compatible_config(config: dict) -> dict:
+    config = dict(config)
+    aliases = {"linear_attention": "mamba", "full_attention": "attention"}
+    for field in ("layers_block_type", "mtp_layers_block_type"):
+        if config.get(field) is not None:
+            config[field] = [aliases.get(value, value) for value in config[field]]
+    return config
+
+
 class NemotronH_Nano_VL_V2_Config(PretrainedConfig):
     model_type = "NemotronH_Nano_VL_V2"
     is_composition = True
@@ -166,3 +175,15 @@ def __init__(self, *args, **kwargs):
         # Explicit __init__ prevents PretrainedConfig.__init_subclass__ from
         # replacing the parent's custom __init__ with a dataclass-generated one.
         super().__init__(*args, **kwargs)
+
+
+class NemotronH_Omni_Reasoning_V3_Config(NemotronH_Nano_Omni_Reasoning_V3_Config):
+    model_type = "nemotron_h_omni"
+
+    def __init__(self, *args, **kwargs):
+        args = list(args)
+        if len(args) > 1 and args[1] is not None:
+            args[1] = _nemotron_h_compatible_config(args[1])
+        elif kwargs.get("llm_config") is not None:
+            kwargs["llm_config"] = _nemotron_h_compatible_config(kwargs["llm_config"])
+        super().__init__(*args, **kwargs)

python/sglang/srt/managers/scheduler.py

@@ -973,10 +973,11 @@ def init_mamba_backend(self) -> None:
             initialize_mamba_selective_state_update_backend(self.server_args)
 
     def init_moe_gemm_config(self):
-        # For the MM models, check the text_config for MoE settings
-        config_to_check = getattr(
-            self.model_config.hf_config, "text_config", self.model_config.hf_config
-        )
+        config_to_check = self.model_config.hf_config
+        if hasattr(self.model_config.hf_config, "text_config"):
+            config_to_check = self.model_config.hf_config.text_config
+        elif hasattr(self.model_config, "hf_text_config"):
+            config_to_check = self.model_config.hf_text_config
 
         # Different MoE architectures expose the per-token expert count under
         # different attribute names (e.g. Gemma4 uses ``top_k_experts``,

python/sglang/srt/models/nano_nemotron_vl.py

@@ -16,6 +16,7 @@
 # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/nano_nemotron_vl.py
 
 import logging
+from collections import deque
 from typing import Iterable
 
 import torch
@@ -178,6 +179,18 @@ def pixel_shuffle(self, x: torch.Tensor, scale_factor: float = 0.5) -> torch.Ten
             x = x.permute(0, 2, 1, 3).contiguous()
         return x
 
+    def _normalize_vision_features(self, features: torch.Tensor) -> torch.Tensor:
+        return features
+
+    def _load_extra_weight(self, name: str, weight: torch.Tensor) -> None:
+        return
+
+    def get_embed_and_head(self):
+        return self.language_model.get_embed_and_head()
+
+    def set_embed_and_head(self, embed, head):
+        self.language_model.set_embed_and_head(embed, head)
+
     def extract_feature_dynamic(self, pixel_values_list: list[torch.Tensor]):
         """Extract features from variable-size images (dynamic resolution).
 
@@ -190,6 +203,7 @@ def extract_feature_dynamic(self, pixel_values_list: list[torch.Tensor]):
         offset = 0
         for i, num_patches in enumerate(num_patches_list):
             img_feats = features[0, offset : offset + num_patches]
+            img_feats = self._normalize_vision_features(img_feats)
             h_patches = pixel_values_list[i].shape[-2] // patch_size
             w_patches = pixel_values_list[i].shape[-1] // patch_size
             img_feats = img_feats.reshape(1, h_patches, w_patches, -1)
@@ -203,6 +217,7 @@ def extract_feature_dynamic(self, pixel_values_list: list[torch.Tensor]):
     def extract_video_feature_temporal(self, pixel_values, num_frames):
         """Extract video features with temporal compression (tubelet grouping)."""
         vit_embeds = self.vision_model(pixel_values, num_frames=num_frames)
+        vit_embeds = self._normalize_vision_features(vit_embeds)
         num_tubelets = vit_embeds.shape[0]
         patch_size = self.config.patch_size
         h_patches = pixel_values.shape[-2] // patch_size
@@ -217,6 +232,9 @@ def extract_video_feature_temporal(self, pixel_values, num_frames):
     def get_input_embeddings(self):
         return self.language_model.get_input_embeddings()
 
+    def set_dflash_layers_to_capture(self, layer_ids: list[int]) -> None:
+        self.language_model.set_dflash_layers_to_capture(layer_ids)
+
     def extract_feature(self, pixel_values):
         micro_batch_size = 128
         n = pixel_values.shape[0]
@@ -229,6 +247,7 @@ def extract_feature(self, pixel_values):
             batch_size = chunk.shape[0]
             vit_embeds = self.vision_model(chunk)
             vit_embeds = vit_embeds.to(dtype=self.model_dtype)
+            vit_embeds = self._normalize_vision_features(vit_embeds)
             vit_embeds = vit_embeds.reshape(batch_size, h_patches, w_patches, -1)
             vit_embeds = self.pixel_shuffle(
                 vit_embeds, scale_factor=self.downsample_ratio
@@ -358,36 +377,96 @@ def is_vision_weights(name: str) -> bool:
         def is_sound_weights(name: str) -> bool:
             return name.startswith("sound")
 
-        # Separate weights by component
-        llm_weights = []
-        vision_weights = []
-        sound_weights = []
-
-        for name, w in weights:
-            if is_llm(name):
-                # Strip 'language_model.' prefix for LLM weights
-                llm_weights.append((".".join(name.split(".")[1:]), w))
-            elif is_adapter_weights((name, w)):
-                # Load vision-language adapter weights directly
-                trimmed_name = ".".join(name.split(".")[1:])
-                param = adapter_dict[trimmed_name]
-                with torch.no_grad():
-                    default_weight_loader(param, w)
-            elif is_vision_weights(name):
-                # Convert: vision_model.radio_model.* → radio_model.*
-                hf_key = name[len("vision_model.") :]
-                vision_weights.append((hf_key, w))
-            elif is_sound_weights(name):
-                sound_weights.append((name, w))
-
+        def iter_llm_weights():
+            for name, w in weights:
+                if is_llm(name):
+                    yield (".".join(name.split(".")[1:]), w)
+                elif is_adapter_weights((name, w)):
+                    trimmed_name = ".".join(name.split(".")[1:])
+                    param = adapter_dict[trimmed_name]
+                    with torch.no_grad():
+                        default_weight_loader(param, w)
+                elif is_vision_weights(name):
+                    hf_key = name[len("vision_model.") :]
+                    self.vision_model.load_weights([(hf_key, w)])
+                elif is_sound_weights(name):
+                    if self.sound_encoder is not None:
+                        self.sound_encoder.load_weights([(name, w)])
+                else:
+                    self._load_extra_weight(name, w)
+
+        llm_weights = iter_llm_weights()
         self.language_model.load_weights(llm_weights)
-        self.vision_model.load_weights(vision_weights)
-        if self.sound_encoder is not None and len(sound_weights) > 0:
-            self.sound_encoder.load_weights(sound_weights)
+        deque(llm_weights, maxlen=0)
 
 
 class NemotronH_Nano_Omni_Reasoning_V3(NemotronH_Nano_VL_V2):
     pass
 
 
-EntryClass = [NemotronH_Nano_VL_V2, NemotronH_Nano_Omni_Reasoning_V3]
+class NemotronH_Omni_Reasoning_V3(NemotronH_Nano_VL_V2):
+    packed_modules_mapping = NemotronHForCausalLM.packed_modules_mapping
+    _hf_projector_weight_names = {
+        "vision_projector.mlp1.norm.": "mlp1.0.",
+        "vision_projector.mlp1.linear1.": "mlp1.1.",
+        "vision_projector.mlp1.linear2.": "mlp1.3.",
+    }
+
+    def __init__(self, config, quant_config=None, prefix: str = ""):
+        super().__init__(config, quant_config, prefix)
+        self.vision_final_layernorm = (
+            nn.LayerNorm(
+                config.vit_hidden_size,
+                eps=config.raw_vision_config.get("layer_norm_eps", 1e-6),
+            ).to(self.model_dtype)
+            if (config.llm_config.num_nextn_predict_layers or 0) > 0
+            else None
+        )
+
+    @property
+    def lm_head(self):
+        return self.language_model.lm_head
+
+    def _normalize_vision_features(self, features: torch.Tensor) -> torch.Tensor:
+        if self.vision_final_layernorm is None:
+            return features
+        return self.vision_final_layernorm(features)
+
+    def _load_extra_weight(self, name: str, weight: torch.Tensor) -> None:
+        prefix = "vision_projector.vision_final_layernorm."
+        if not name.startswith(prefix):
+            raise ValueError(f"Unexpected Nemotron-H Omni weight: {name}")
+        if self.vision_final_layernorm is None:
+            raise ValueError(f"Unexpected vision projector weight: {name}")
+        parameter_name = name.removeprefix(prefix)
+        parameters = dict(self.vision_final_layernorm.named_parameters())
+        if parameter_name not in parameters:
+            raise ValueError(f"Unexpected vision projector weight: {name}")
+        parameter = parameters[parameter_name]
+        default_weight_loader(parameter, weight)
+
+    @classmethod
+    def _remap_checkpoint_weight_name(cls, name: str) -> str:
+        for source, target in cls._hf_projector_weight_names.items():
+            if name.startswith(source):
+                return name.replace(source, target, 1)
+        if name.startswith("vision_model.") and not name.startswith(
+            "vision_model.radio_model."
+        ):
+            return name.replace(
+                "vision_model.", "vision_model.radio_model.hf_model.", 1
+            )
+        return name
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
+        super().load_weights(
+            (self._remap_checkpoint_weight_name(name), weight)
+            for name, weight in weights
+        )
+
+
+EntryClass = [
+    NemotronH_Nano_VL_V2,
+    NemotronH_Nano_Omni_Reasoning_V3,
+    NemotronH_Omni_Reasoning_V3,
+]