From efe38480cd2e6b77d8fbd25c2c7fea5fdba64f19 Mon Sep 17 00:00:00 2001 From: Ryan Stewart Date: Wed, 19 Aug 2026 21:53:53 +0000 Subject: [PATCH 01/13] add native Nemotron 3.5 Super VL support Signed-off-by: Ryan Stewart --- python/sglang/srt/arg_groups/overrides.py | 16 +++- python/sglang/srt/configs/__init__.py | 2 + python/sglang/srt/configs/model_config.py | 9 ++ python/sglang/srt/configs/nano_nemotron_vl.py | 21 +++++ python/sglang/srt/models/nano_nemotron_vl.py | 82 ++++++++++++++++++- python/sglang/srt/models/nemotron_h_mtp.py | 3 + .../multimodal/processors/nano_nemotron_vl.py | 9 +- python/sglang/srt/server_args.py | 2 +- .../srt/utils/hf_transformers/common.py | 2 + .../unit/configs/test_model_config.py | 15 ++++ .../models/test_nemotron_h_weight_loading.py | 68 ++++++++++++++- 11 files changed, 222 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 645c7854e762..fbd606ee144e 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1462,7 +1462,11 @@ def _inkling_overrides(server_args: Any, hf_config: Any) -> dict: return overrides -@_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 @@ -1480,7 +1484,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": @@ -1771,6 +1780,7 @@ def _step3p_overrides(server_args: Any, hf_config: Any) -> dict: "MiniCPMV4_6ForConditionalGeneration", "NemotronHForCausalLM", "NemotronHPuzzleForCausalLM", + "NemotronH_Omni_Reasoning_V3", "FalconH1ForCausalLM", "JetNemotronForCausalLM", "JetVLMForConditionalGeneration", @@ -1800,6 +1810,7 @@ def _step3p_overrides(server_args: Any, hf_config: Any) -> dict: "GraniteMoeHybridForCausalLM", "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. @@ -2233,6 +2244,7 @@ def _sparse_head_overlap_disable(view: Any) -> dict: "Qwen3_5ForConditionalGeneration", "NemotronHForCausalLM", "NemotronHPuzzleForCausalLM", + "NemotronH_Omni_Reasoning_V3", } ) diff --git a/python/sglang/srt/configs/__init__.py b/python/sglang/srt/configs/__init__.py index 492760909ba5..941608087589 100644 --- a/python/sglang/srt/configs/__init__.py +++ b/python/sglang/srt/configs/__init__.py @@ -46,6 +46,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 @@ -113,6 +114,7 @@ "NemotronHPuzzleConfig", "NemotronH_Nano_VL_V2_Config", "NemotronH_Nano_Omni_Reasoning_V3_Config", + "NemotronH_Omni_Reasoning_V3_Config", "JetNemotronConfig", "JetVLMConfig", "MiniCPMHybridConfig", diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 5ec73e5410c6..b15f79f09675 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -759,6 +759,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", @@ -1892,6 +1900,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", diff --git a/python/sglang/srt/configs/nano_nemotron_vl.py b/python/sglang/srt/configs/nano_nemotron_vl.py index 888f94a1b433..1733eceb6ef9 100644 --- a/python/sglang/srt/configs/nano_nemotron_vl.py +++ b/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) diff --git a/python/sglang/srt/models/nano_nemotron_vl.py b/python/sglang/srt/models/nano_nemotron_vl.py index 021d33ad8bfe..7e1d5cd2f116 100644 --- a/python/sglang/srt/models/nano_nemotron_vl.py +++ b/python/sglang/srt/models/nano_nemotron_vl.py @@ -178,6 +178,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 +202,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 +216,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 @@ -229,6 +243,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 @@ -379,6 +394,8 @@ def is_sound_weights(name: str) -> bool: vision_weights.append((hf_key, w)) elif is_sound_weights(name): sound_weights.append((name, w)) + else: + self._load_extra_weight(name, w) self.language_model.load_weights(llm_weights) self.vision_model.load_weights(vision_weights) @@ -390,4 +407,67 @@ 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): + return + if self.vision_final_layernorm is None: + raise ValueError(f"Unexpected vision projector weight: {name}") + parameter = dict(self.vision_final_layernorm.named_parameters())[ + name.removeprefix(prefix) + ] + 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, +] diff --git a/python/sglang/srt/models/nemotron_h_mtp.py b/python/sglang/srt/models/nemotron_h_mtp.py index f569baf1c9a1..458d1a39d5dd 100644 --- a/python/sglang/srt/models/nemotron_h_mtp.py +++ b/python/sglang/srt/models/nemotron_h_mtp.py @@ -366,6 +366,9 @@ def forward( def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]], is_mtp: bool = False ): + weights = ( + (name.removeprefix("language_model."), weight) for name, weight in weights + ) super().load_weights(weights, is_mtp=True) def set_lm_head_from_target(self, target_lm_head: nn.Module) -> None: diff --git a/python/sglang/srt/multimodal/processors/nano_nemotron_vl.py b/python/sglang/srt/multimodal/processors/nano_nemotron_vl.py index 32bab7ae4470..9060d13a253d 100644 --- a/python/sglang/srt/multimodal/processors/nano_nemotron_vl.py +++ b/python/sglang/srt/multimodal/processors/nano_nemotron_vl.py @@ -22,6 +22,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.managers.schedule_batch import ( Modality, @@ -31,6 +32,7 @@ from sglang.srt.models.nano_nemotron_vl import ( NemotronH_Nano_Omni_Reasoning_V3, NemotronH_Nano_VL_V2, + NemotronH_Omni_Reasoning_V3, ) from sglang.srt.models.parakeet import ParakeetExtractor from sglang.srt.multimodal.audio_from_video import extract_audio_from_video_bytes @@ -57,7 +59,11 @@ class NanoNemotronVLImageProcessor(BaseMultimodalProcessor): - models = [NemotronH_Nano_VL_V2, NemotronH_Nano_Omni_Reasoning_V3] + models = [ + NemotronH_Nano_VL_V2, + NemotronH_Nano_Omni_Reasoning_V3, + NemotronH_Omni_Reasoning_V3, + ] gpu_image_decode = ( False # NanoNemotronVL processes loaded image as PIL image explicitly ) @@ -69,6 +75,7 @@ def __init__(self, hf_config, server_args, _image_processor, *args, **kwargs): { NemotronH_Nano_VL_V2_Config: NemotronH_Nano_VL_V2, NemotronH_Nano_Omni_Reasoning_V3_Config: NemotronH_Nano_Omni_Reasoning_V3, + NemotronH_Omni_Reasoning_V3_Config: NemotronH_Omni_Reasoning_V3, }, ) Image.MAX_IMAGE_PIXELS = None diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 4a7820f28030..e9d87cb2e3ab 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -9905,7 +9905,7 @@ def mamba_cache_chunk_size(self) -> int: # Must match sglang.kernels.ops.attention.fla.chunk_delta_h.CHUNK_SIZE FLA_CHUNK_SIZE = 64 - hf_config = self.get_model_config().hf_config + hf_config = self.get_model_config().hf_text_config chunk_size = getattr(hf_config, "mamba_chunk_size", FLA_CHUNK_SIZE) page_size = resolved_view(self).page_size assert ( diff --git a/python/sglang/srt/utils/hf_transformers/common.py b/python/sglang/srt/utils/hf_transformers/common.py index 5ffa8530c4d7..8a5c863cc696 100644 --- a/python/sglang/srt/utils/hf_transformers/common.py +++ b/python/sglang/srt/utils/hf_transformers/common.py @@ -58,6 +58,7 @@ MuseGlimmerConfig, NemotronH_Nano_Omni_Reasoning_V3_Config, NemotronH_Nano_VL_V2_Config, + NemotronH_Omni_Reasoning_V3_Config, NemotronHConfig, NemotronHPuzzleConfig, Olmo3Config, @@ -119,6 +120,7 @@ Dots3Config, NemotronH_Nano_VL_V2_Config, NemotronH_Nano_Omni_Reasoning_V3_Config, + NemotronH_Omni_Reasoning_V3_Config, NemotronHConfig, NemotronHPuzzleConfig, DeepseekVLV2Config, diff --git a/test/registered/unit/configs/test_model_config.py b/test/registered/unit/configs/test_model_config.py index 778928154eff..baca934fc569 100644 --- a/test/registered/unit/configs/test_model_config.py +++ b/test/registered/unit/configs/test_model_config.py @@ -69,6 +69,21 @@ def test_qwen35_mtp_depth_is_synced_to_text_config(self): self.assertEqual(config.hf_config.num_nextn_predict_layers, 1) self.assertEqual(config.hf_text_config.num_nextn_predict_layers, 1) + def test_nemotron_h_omni_mtp_uses_language_model_config(self): + config = object.__new__(ModelConfig) + config.is_draft_model = True + config.speculative_algorithm = "EAGLE" + config.hf_config = SimpleNamespace( + architectures=["NemotronH_Omni_Reasoning_V3"] + ) + config.hf_text_config = SimpleNamespace(architectures=["NemotronHForCausalLM"]) + + config._config_draft_model() + + self.assertIs(config.hf_config, config.hf_text_config) + self.assertEqual(config.hf_config.architectures, ["NemotronHForCausalLMMTP"]) + self.assertEqual(config.hf_config.num_nextn_predict_layers, 1) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/models/test_nemotron_h_weight_loading.py b/test/registered/unit/models/test_nemotron_h_weight_loading.py index ed4c21d234ff..80287837db83 100644 --- a/test/registered/unit/models/test_nemotron_h_weight_loading.py +++ b/test/registered/unit/models/test_nemotron_h_weight_loading.py @@ -15,6 +15,7 @@ import torch from sglang.srt.models.nemotron_h import NemotronHForCausalLM +from sglang.srt.models.nemotron_h_mtp import NemotronHForCausalLMMTP class _FakePPGroup: @@ -43,9 +44,19 @@ def weight_loader(self, param, loaded_weight): self.loaded_weight = loaded_weight +class _RecordingStackedParam: + def __init__(self): + self.loads = [] + + def weight_loader(self, param, loaded_weight, shard_id): + self.loads.append((param, loaded_weight, shard_id)) + + class TestNemotronHWeightLoading(unittest.TestCase): - def _make_minimal_model(self, named_parameters=()): - model = object.__new__(NemotronHForCausalLM) + def _make_minimal_model( + self, named_parameters=(), model_class=NemotronHForCausalLM + ): + model = object.__new__(model_class) model.config = SimpleNamespace(n_routed_experts=2, max_n_routed_experts=2) model.model = SimpleNamespace() model.pp_group = _FakePPGroup() @@ -134,6 +145,59 @@ def test_mtp_keeps_shared_embed_tokens_and_lm_head(self): skipped.loaded_weight, "non-MTP target weight should be skipped" ) + def test_mtp_strips_multimodal_language_model_prefix(self): + embed = _RecordingParam() + head = _RecordingParam() + mtp_layer = _RecordingParam() + model = self._make_minimal_model( + [ + ("model.embed_tokens.weight", embed), + ("lm_head.weight", head), + ("model.layers.0.norm.weight", mtp_layer), + ], + model_class=NemotronHForCausalLMMTP, + ) + model.remap_prefix = {"backbone": "model"} + model.remap_substr = {"embeddings": "embed_tokens"} + + w_embed, w_head, w_mtp = (torch.ones(1) for _ in range(3)) + model.load_weights( + [ + ("language_model.backbone.embeddings.weight", w_embed), + ("language_model.lm_head.weight", w_head), + ("language_model.mtp.layers.0.norm.weight", w_mtp), + ] + ) + + self.assertIs(embed.loaded_weight, w_embed) + self.assertIs(head.loaded_weight, w_head) + self.assertIs(mtp_layer.loaded_weight, w_mtp) + + def test_split_qkv_fp8_scales_load_into_fused_parameter(self): + input_scale = _RecordingStackedParam() + model = self._make_minimal_model( + [("model.layers.7.mixer.qkv_proj.input_scale", input_scale)] + ) + model.stacked_params_mapping = NemotronHForCausalLM.stacked_params_mapping + + q_scale, k_scale, v_scale = (torch.tensor(value) for value in (1, 2, 3)) + model.load_weights( + [ + ("model.layers.7.mixer.q_proj.input_scale", q_scale), + ("model.layers.7.mixer.k_proj.input_scale", k_scale), + ("model.layers.7.mixer.v_proj.input_scale", v_scale), + ] + ) + + self.assertEqual( + input_scale.loads, + [ + (input_scale, q_scale, "q"), + (input_scale, k_scale, "k"), + (input_scale, v_scale, "v"), + ], + ) + if __name__ == "__main__": unittest.main() From b1af8deae2bed9a7990c1978eee549513e29047c Mon Sep 17 00:00:00 2001 From: Ryan Stewart Date: Wed, 19 Aug 2026 21:54:08 +0000 Subject: [PATCH 02/13] load quantized RADIO exports for Nemotron 3.5 Super VL Signed-off-by: Ryan Stewart --- python/sglang/srt/models/radio.py | 57 +++- .../test_nemotron_h_omni_registration.py | 289 ++++++++++++++++++ 2 files changed, 340 insertions(+), 6 deletions(-) create mode 100644 test/registered/unit/models/test_nemotron_h_omni_registration.py diff --git a/python/sglang/srt/models/radio.py b/python/sglang/srt/models/radio.py index df295f20200d..7a606d4b0e29 100644 --- a/python/sglang/srt/models/radio.py +++ b/python/sglang/srt/models/radio.py @@ -42,6 +42,41 @@ norm_t: TypeAlias = tuple[float, float, float] | torch.Tensor +def _map_hf_radio_weight_name(name: str) -> tuple[str, str | None] | None: + prefix = "radio_model.hf_model." + if not name.startswith(prefix): + return None + + name = name.removeprefix(prefix) + if name == "summary_idxs": + return None + + embedding_names = { + "embeddings.cls_register_token": "model.patch_generator.cls_token.token", + "embeddings.patch_projection": "model.patch_generator.embedder", + "embeddings.position_embedding": "model.patch_generator.pos_embed", + "embeddings.video_patch_projection": "model.patch_generator.video_embedder", + } + for source, target in embedding_names.items(): + if name == source or name.startswith(f"{source}."): + return name.replace(source, target, 1), None + + name = name.replace("encoder.layer.", "model.encoder.layers.", 1) + attention_names = { + ".attention.attention.query.": (".attn.attn.qkv_proj.", "q"), + ".attention.attention.key.": (".attn.attn.qkv_proj.", "k"), + ".attention.attention.value.": (".attn.attn.qkv_proj.", "v"), + ".attention.output.dense.": (".attn.attn.proj.", None), + } + for source, (target, shard_id) in attention_names.items(): + if source in name: + return name.replace(source, target, 1), shard_id + + name = name.replace(".layer_scale1.lambda1", ".ls1") + name = name.replace(".layer_scale2.lambda1", ".ls2") + return name, None + + def _ntuple(n): def parse(x): if isinstance(x, Iterable) and not isinstance(x, str): @@ -588,15 +623,25 @@ def load_weights(self, weights) -> set[str]: weights_list = list(weights) for name, weight in weights_list: - if not name.startswith("radio_model."): - # Skip non-radio weights - continue - name = replace_substrings(name, remap_substrings) - name = replace_prefix(name, remap_prefixes) + loaded_shard_id = None + if name.startswith("radio_model.hf_model."): + mapped_weight = _map_hf_radio_weight_name(name) + if mapped_weight is None: + continue + name, loaded_shard_id = mapped_weight + else: + if not name.startswith("radio_model."): + # Skip non-radio weights + continue + name = replace_substrings(name, remap_substrings) + name = replace_prefix(name, remap_prefixes) if name and name in params_dict: param = params_dict[name] weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, weight) + if loaded_shard_id is None: + weight_loader(param, weight) + else: + weight_loader(param, weight, loaded_shard_id) loaded_params.add(name) if "video_embedder" in name: self.model.patch_generator._video_embedder_loaded = True diff --git a/test/registered/unit/models/test_nemotron_h_omni_registration.py b/test/registered/unit/models/test_nemotron_h_omni_registration.py new file mode 100644 index 000000000000..2a23c28dd5ef --- /dev/null +++ b/test/registered/unit/models/test_nemotron_h_omni_registration.py @@ -0,0 +1,289 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch +import torch.nn as nn + +from sglang.srt.arg_groups.overrides import ( + _FLASHINFER_ALLREDUCE_FUSION_ARCHS, + _MAMBA_EXTRA_BUFFER_ARCHS, + _MAMBA_RADIX_CACHE_ARCHS, + _MODEL_OVERRIDE_FNS, +) +from sglang.srt.configs.model_config import is_multimodal_model +from sglang.srt.configs.nano_nemotron_vl import ( + NemotronH_Omni_Reasoning_V3_Config, +) +from sglang.srt.layers.quantization.modelopt_quant import ( + ModelOptMixedPrecisionConfig, +) +from sglang.srt.models.nano_nemotron_vl import NemotronH_Omni_Reasoning_V3 +from sglang.srt.models.nemotron_h import NemotronHForCausalLM +from sglang.srt.models.radio import RadioModel, _map_hf_radio_weight_name +from sglang.srt.multimodal.processors.nano_nemotron_vl import ( + NanoNemotronVLImageProcessor, +) +from sglang.srt.server_args import ServerArgs +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + + +class _RecordingWeight: + def __init__(self): + self.loads = [] + + def weight_loader(self, param, weight, shard_id=None): + self.loads.append((param, weight, shard_id)) + + +class TestNemotronHOmniRegistration(CustomTestCase): + def test_config_uses_checkpoint_model_type(self): + config = NemotronH_Omni_Reasoning_V3_Config( + vision_config={"args": {"model": "radio"}}, + llm_config={}, + architectures=["NemotronH_Omni_Reasoning_V3"], + ) + + self.assertEqual(config.model_type, "nemotron_h_omni") + self.assertEqual(config.architectures, ["NemotronH_Omni_Reasoning_V3"]) + + def test_config_normalizes_current_nemotron_h_layer_names(self): + llm_config = { + "layers_block_type": ["linear_attention", "moe", "full_attention"], + "num_nextn_predict_layers": 1, + "mtp_layers_block_type": ["full_attention", "moe"], + } + + config = NemotronH_Omni_Reasoning_V3_Config( + vision_config={"args": {"model": "radio"}}, + llm_config=llm_config, + ) + + self.assertEqual( + config.llm_config.layers_block_type, + ["mamba", "moe", "attention"], + ) + self.assertEqual( + config.llm_config.mtp_layers_block_type, + ["attention", "moe"], + ) + self.assertEqual( + llm_config["layers_block_type"], + ["linear_attention", "moe", "full_attention"], + ) + + def test_model_and_processor_register_new_architecture(self): + from sglang.srt.models.registry import ModelRegistry + + model_class, architecture = ModelRegistry.resolve_model_cls( + "NemotronH_Omni_Reasoning_V3" + ) + + self.assertIs(model_class, NemotronH_Omni_Reasoning_V3) + self.assertEqual(architecture, "NemotronH_Omni_Reasoning_V3") + self.assertIn(NemotronH_Omni_Reasoning_V3, NanoNemotronVLImageProcessor.models) + + def test_new_architecture_is_multimodal(self): + self.assertTrue(is_multimodal_model(["NemotronH_Omni_Reasoning_V3"])) + + def test_new_architecture_uses_nemotron_h_runtime_policy(self): + architecture = "NemotronH_Omni_Reasoning_V3" + + self.assertIn(architecture, _MODEL_OVERRIDE_FNS) + self.assertIn(architecture, _MAMBA_RADIX_CACHE_ARCHS) + self.assertIn(architecture, _MAMBA_EXTRA_BUFFER_ARCHS) + self.assertIn(architecture, _FLASHINFER_ALLREDUCE_FUSION_ARCHS) + + def test_mixed_precision_resolves_fused_qkv_from_split_layers(self): + quant_config = ModelOptMixedPrecisionConfig.from_config( + { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + f"language_model.model.layers.7.mixer.{projection}": { + "quant_algo": "FP8" + } + for projection in ("q_proj", "k_proj", "v_proj") + }, + "packed_modules_mapping": ( + NemotronH_Omni_Reasoning_V3.packed_modules_mapping + ), + } + ) + + self.assertEqual( + NemotronH_Omni_Reasoning_V3.packed_modules_mapping, + NemotronHForCausalLM.packed_modules_mapping, + ) + self.assertEqual( + quant_config._resolve_quant_algo( + "language_model.model.layers.7.mixer.qkv_proj" + ), + "FP8", + ) + + def test_mamba_cache_chunk_size_uses_language_config(self): + server_args = object.__new__(ServerArgs) + server_args.get_model_config = lambda: SimpleNamespace( + hf_config=SimpleNamespace(), + hf_text_config=SimpleNamespace(mamba_chunk_size=128), + ) + + with patch( + "sglang.srt.server_args.resolved_view", + return_value=SimpleNamespace(page_size=64), + ): + self.assertEqual(server_args.mamba_cache_chunk_size, 128) + + def test_multimodal_wrapper_exposes_language_embed_and_head(self): + model = object.__new__(NemotronH_Omni_Reasoning_V3) + nn.Module.__init__(model) + embed = object() + head = object() + model.language_model = SimpleNamespace( + get_embed_and_head=lambda: (embed, head), + lm_head=head, + ) + + self.assertEqual(model.get_embed_and_head(), (embed, head)) + self.assertIs(model.lm_head, head) + + def test_super_vision_final_layernorm_is_loaded_and_applied(self): + model = object.__new__(NemotronH_Omni_Reasoning_V3) + nn.Module.__init__(model) + model.mlp1 = nn.Sequential() + model.vision_final_layernorm = nn.LayerNorm(2) + model.language_model = SimpleNamespace(load_weights=lambda weights: None) + model.vision_model = SimpleNamespace(load_weights=lambda weights: None) + model.sound_encoder = None + + weight = torch.tensor([2.0, 3.0]) + bias = torch.tensor([0.5, -0.5]) + model.load_weights( + [ + ("vision_projector.vision_final_layernorm.weight", weight), + ("vision_projector.vision_final_layernorm.bias", bias), + ] + ) + + features = torch.tensor([[1.0, 3.0]]) + expected = nn.functional.layer_norm(features, (2,), weight, bias) + torch.testing.assert_close(model._normalize_vision_features(features), expected) + + def test_super_hf_vision_and_projector_names_are_remapped(self): + remap = NemotronH_Omni_Reasoning_V3._remap_checkpoint_weight_name + + self.assertEqual( + remap("vision_model.embeddings.position_embedding"), + "vision_model.radio_model.hf_model.embeddings.position_embedding", + ) + self.assertEqual( + remap("vision_model.embeddings.video_patch_projection.weight"), + "vision_model.radio_model.hf_model.embeddings.video_patch_projection.weight", + ) + self.assertEqual( + remap("vision_projector.mlp1.linear1.weight"), + "mlp1.1.weight", + ) + self.assertEqual( + remap("vision_model.radio_model.model.patch_generator.pos_embed"), + "vision_model.radio_model.model.patch_generator.pos_embed", + ) + + def test_hf_radio_loader_maps_embeddings_and_split_qkv(self): + model = object.__new__(RadioModel) + nn.Module.__init__(model) + position_embedding = _RecordingWeight() + qkv_weight = _RecordingWeight() + model.named_parameters = lambda: iter( + [ + ("model.patch_generator.pos_embed", position_embedding), + ( + "model.encoder.layers.0.attn.attn.qkv_proj.weight", + qkv_weight, + ), + ] + ) + + position = torch.ones(1) + query, key, value = (torch.full((1,), value) for value in (2, 3, 4)) + loaded = model.load_weights( + [ + ( + "radio_model.hf_model.embeddings.position_embedding", + position, + ), + ( + "radio_model.hf_model.encoder.layer.0.attention.attention.query.weight", + query, + ), + ( + "radio_model.hf_model.encoder.layer.0.attention.attention.key.weight", + key, + ), + ( + "radio_model.hf_model.encoder.layer.0.attention.attention.value.weight", + value, + ), + ("radio_model.hf_model.summary_idxs", torch.tensor([0, 1])), + ] + ) + + self.assertEqual( + loaded, + { + "model.patch_generator.pos_embed", + "model.encoder.layers.0.attn.attn.qkv_proj.weight", + }, + ) + self.assertEqual( + position_embedding.loads, [(position_embedding, position, None)] + ) + self.assertEqual( + qkv_weight.loads, + [ + (qkv_weight, query, "q"), + (qkv_weight, key, "k"), + (qkv_weight, value, "v"), + ], + ) + + def test_hf_radio_name_mapping_covers_encoder_parameters(self): + cases = { + "embeddings.video_patch_projection.weight": ( + "model.patch_generator.video_embedder.weight", + None, + ), + "encoder.layer.1.attention.output.dense.weight": ( + "model.encoder.layers.1.attn.attn.proj.weight", + None, + ), + "encoder.layer.2.layer_scale1.lambda1": ( + "model.encoder.layers.2.ls1", + None, + ), + "encoder.layer.3.layer_scale2.lambda1": ( + "model.encoder.layers.3.ls2", + None, + ), + "encoder.layer.4.mlp.fc1.bias": ( + "model.encoder.layers.4.mlp.fc1.bias", + None, + ), + "encoder.layer.5.norm2.weight": ( + "model.encoder.layers.5.norm2.weight", + None, + ), + } + for source, expected in cases.items(): + with self.subTest(source=source): + self.assertEqual( + _map_hf_radio_weight_name(f"radio_model.hf_model.{source}"), + expected, + ) + + +if __name__ == "__main__": + unittest.main() From 74834a7bdfa3eca52b25827ade0fdf5d71283c40 Mon Sep 17 00:00:00 2001 From: elvischenv <219235043+elvischenv@users.noreply.github.com> Date: Tue, 25 Aug 2026 03:14:52 +0800 Subject: [PATCH 03/13] initialize MoE config for VLMs that use llm_config init_moe_gemm_config only looked for hf_config.text_config, so VLMs whose language sub-config is named llm_config (NemotronH_Omni_Reasoning_V3) fell back to the top-level config, found no MoE top-k attribute, and skipped initialize_moe_config entirely. get_moe_runner_backend() then returned auto instead of the configured backend, leaving NVFP4 MoE weights prepared for CUTLASS while the runner resolved to FlashInfer TRTLLM: NotImplementedError: Unsupported moe_runner_backend for NVFP4 MoE: MoeRunnerBackend.FLASHINFER_TRTLLM. Fall back to hf_text_config, which resolves llm_config / language_config / thinker_config and returns the top-level config when there is no sub-config. Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com> Co-Authored-By: Claude Opus 5 (1M context) (cherry picked from commit 64ff8e99d689a83d77be6fce80186d9ab492ba17) --- python/sglang/srt/managers/scheduler.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index bc50b2fcaf01..aec3a89e50cb 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -886,10 +886,12 @@ 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 - ) + # For the MM models, check the text_config or hf_text_config for MoE settings + 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``, From 61a761d2da5d6d60cf1d44fc7760850f3ef66b8e Mon Sep 17 00:00:00 2001 From: Ryan Stewart Date: Tue, 25 Aug 2026 17:13:12 +0000 Subject: [PATCH 04/13] use canonical text config for MoE initialization Signed-off-by: Ryan Stewart --- python/sglang/srt/managers/scheduler.py | 7 +- .../managers/test_scheduler_moe_config.py | 78 +++++++++++++++++++ 2 files changed, 79 insertions(+), 6 deletions(-) create mode 100644 test/registered/unit/managers/test_scheduler_moe_config.py diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index aec3a89e50cb..73751afb97ce 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -886,12 +886,7 @@ 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 or hf_text_config for MoE settings - 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 + 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``, diff --git a/test/registered/unit/managers/test_scheduler_moe_config.py b/test/registered/unit/managers/test_scheduler_moe_config.py new file mode 100644 index 000000000000..1021480d0f88 --- /dev/null +++ b/test/registered/unit/managers/test_scheduler_moe_config.py @@ -0,0 +1,78 @@ +"""Tests for scheduler MoE configuration discovery.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +import sglang.srt.managers.scheduler as scheduler_module +from sglang.srt.managers.scheduler import Scheduler + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestSchedulerMoeConfig(CustomTestCase): + def _run_init(self, *, hf_config, hf_text_config): + scheduler = Scheduler.__new__(Scheduler) + scheduler.server_args = object() + scheduler.model_config = SimpleNamespace( + hf_config=hf_config, + hf_text_config=hf_text_config, + ) + + with ( + patch.object(scheduler_module, "initialize_moe_config") as init_moe, + patch.object(scheduler_module, "initialize_fp8_gemm_config"), + patch.object(scheduler_module, "initialize_fp4_gemm_config"), + patch.object(scheduler_module, "initialize_bf16_gemm_config"), + patch.object( + scheduler_module, + "require_mlp_sync", + return_value=False, + ), + ): + scheduler.init_moe_gemm_config() + + return scheduler, init_moe + + def test_uses_canonical_text_config_for_llm_config_vlm(self): + outer_config = SimpleNamespace() + llm_config = SimpleNamespace(num_experts_per_tok=8) + + scheduler, init_moe = self._run_init( + hf_config=outer_config, + hf_text_config=llm_config, + ) + + init_moe.assert_called_once_with(scheduler.server_args) + self.assertFalse(scheduler.require_mlp_sync) + + def test_canonical_text_config_takes_precedence(self): + outer_config = SimpleNamespace(text_config=SimpleNamespace()) + llm_config = SimpleNamespace(num_experts_per_tok=8) + + scheduler, init_moe = self._run_init( + hf_config=outer_config, + hf_text_config=llm_config, + ) + + init_moe.assert_called_once_with(scheduler.server_args) + + def test_skips_moe_initialization_for_non_moe_text_config(self): + outer_config = SimpleNamespace() + text_config = SimpleNamespace() + + _, init_moe = self._run_init( + hf_config=outer_config, + hf_text_config=text_config, + ) + + init_moe.assert_not_called() + + +if __name__ == "__main__": + unittest.main() From 5d762fa0bb2a95ba044139aad4ce0e8ab8590ad7 Mon Sep 17 00:00:00 2001 From: Ryan Stewart Date: Tue, 25 Aug 2026 23:35:28 +0000 Subject: [PATCH 05/13] validate Super VL checkpoint weights Signed-off-by: Ryan Stewart --- python/sglang/srt/models/nano_nemotron_vl.py | 10 +- python/sglang/srt/models/radio.py | 8 +- .../unit/models/test_nano_nemotron_vl.py | 104 ++++++++++++++ test/registered/unit/models/test_radio.py | 132 ++++++++++++++++++ 4 files changed, 249 insertions(+), 5 deletions(-) create mode 100644 test/registered/unit/models/test_nano_nemotron_vl.py create mode 100644 test/registered/unit/models/test_radio.py diff --git a/python/sglang/srt/models/nano_nemotron_vl.py b/python/sglang/srt/models/nano_nemotron_vl.py index 7e1d5cd2f116..81d419ccfcbd 100644 --- a/python/sglang/srt/models/nano_nemotron_vl.py +++ b/python/sglang/srt/models/nano_nemotron_vl.py @@ -438,12 +438,14 @@ def _normalize_vision_features(self, features: torch.Tensor) -> torch.Tensor: def _load_extra_weight(self, name: str, weight: torch.Tensor) -> None: prefix = "vision_projector.vision_final_layernorm." if not name.startswith(prefix): - return + 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 = dict(self.vision_final_layernorm.named_parameters())[ - name.removeprefix(prefix) - ] + 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 diff --git a/python/sglang/srt/models/radio.py b/python/sglang/srt/models/radio.py index 7a606d4b0e29..1e2fde3b62c1 100644 --- a/python/sglang/srt/models/radio.py +++ b/python/sglang/srt/models/radio.py @@ -623,8 +623,10 @@ def load_weights(self, weights) -> set[str]: weights_list = list(weights) for name, weight in weights_list: + source_name = name + is_hf_export = name.startswith("radio_model.hf_model.") loaded_shard_id = None - if name.startswith("radio_model.hf_model."): + if is_hf_export: mapped_weight = _map_hf_radio_weight_name(name) if mapped_weight is None: continue @@ -645,6 +647,10 @@ def load_weights(self, weights) -> set[str]: loaded_params.add(name) if "video_embedder" in name: self.model.patch_generator._video_embedder_loaded = True + elif is_hf_export: + raise ValueError( + f"Unexpected HF RADIO weight: {source_name} (mapped to {name})" + ) return loaded_params diff --git a/test/registered/unit/models/test_nano_nemotron_vl.py b/test/registered/unit/models/test_nano_nemotron_vl.py new file mode 100644 index 000000000000..768a3e987977 --- /dev/null +++ b/test/registered/unit/models/test_nano_nemotron_vl.py @@ -0,0 +1,104 @@ +"""Unit tests for native Nemotron-H Omni model integration.""" + +import unittest +from types import SimpleNamespace + +import torch +import torch.nn as nn + +from sglang.srt.models.nano_nemotron_vl import NemotronH_Omni_Reasoning_V3 +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class TestNemotronHOmniModel(CustomTestCase): + def test_model_registry_resolves_new_architecture(self): + from sglang.srt.models.registry import ModelRegistry + + model_class, architecture = ModelRegistry.resolve_model_cls( + "NemotronH_Omni_Reasoning_V3" + ) + + self.assertIs(model_class, NemotronH_Omni_Reasoning_V3) + self.assertEqual(architecture, "NemotronH_Omni_Reasoning_V3") + + def test_exposes_language_embed_and_head(self): + model = object.__new__(NemotronH_Omni_Reasoning_V3) + nn.Module.__init__(model) + embed = object() + head = object() + model.language_model = SimpleNamespace( + get_embed_and_head=lambda: (embed, head), + lm_head=head, + ) + + self.assertEqual(model.get_embed_and_head(), (embed, head)) + self.assertIs(model.lm_head, head) + + def test_vision_final_layernorm_is_loaded_and_applied(self): + model = object.__new__(NemotronH_Omni_Reasoning_V3) + nn.Module.__init__(model) + model.mlp1 = nn.Sequential() + model.vision_final_layernorm = nn.LayerNorm(2) + model.language_model = SimpleNamespace(load_weights=lambda weights: None) + model.vision_model = SimpleNamespace(load_weights=lambda weights: None) + model.sound_encoder = None + + weight = torch.tensor([2.0, 3.0]) + bias = torch.tensor([0.5, -0.5]) + model.load_weights( + [ + ("vision_projector.vision_final_layernorm.weight", weight), + ("vision_projector.vision_final_layernorm.bias", bias), + ] + ) + + features = torch.tensor([[1.0, 3.0]]) + expected = nn.functional.layer_norm(features, (2,), weight, bias) + torch.testing.assert_close(model._normalize_vision_features(features), expected) + + def test_hf_vision_and_projector_names_are_remapped(self): + remap = NemotronH_Omni_Reasoning_V3._remap_checkpoint_weight_name + + self.assertEqual( + remap("vision_model.embeddings.position_embedding"), + "vision_model.radio_model.hf_model.embeddings.position_embedding", + ) + self.assertEqual( + remap("vision_model.embeddings.video_patch_projection.weight"), + ( + "vision_model.radio_model.hf_model.embeddings." + "video_patch_projection.weight" + ), + ) + self.assertEqual( + remap("vision_projector.mlp1.linear1.weight"), + "mlp1.1.weight", + ) + self.assertEqual( + remap("vision_model.radio_model.model.patch_generator.pos_embed"), + "vision_model.radio_model.model.patch_generator.pos_embed", + ) + + def test_unexpected_checkpoint_weight_raises(self): + model = object.__new__(NemotronH_Omni_Reasoning_V3) + nn.Module.__init__(model) + model.mlp1 = nn.Sequential() + model.vision_final_layernorm = nn.LayerNorm(2) + + cases = ( + ("vision_projector.unknown.weight", "Unexpected Nemotron-H Omni"), + ( + "vision_projector.vision_final_layernorm.running_mean", + "Unexpected vision projector weight", + ), + ) + for name, message in cases: + with self.subTest(name=name), self.assertRaisesRegex(ValueError, message): + model.load_weights([(name, torch.ones(1))]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/models/test_radio.py b/test/registered/unit/models/test_radio.py new file mode 100644 index 000000000000..ff21ee80498d --- /dev/null +++ b/test/registered/unit/models/test_radio.py @@ -0,0 +1,132 @@ +"""Unit tests for RADIO checkpoint weight loading.""" + +import unittest +from types import SimpleNamespace + +import torch +import torch.nn as nn + +from sglang.srt.models.radio import RadioModel +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class _RecordingWeight: + def __init__(self): + self.loads = [] + + def weight_loader(self, param, weight, shard_id=None): + self.loads.append((param, weight, shard_id)) + + +class TestRadioWeightLoading(CustomTestCase): + def _make_model(self, named_parameters=()): + model = object.__new__(RadioModel) + nn.Module.__init__(model) + model.named_parameters = lambda: iter(named_parameters) + model.model = SimpleNamespace( + patch_generator=SimpleNamespace(_video_embedder_loaded=False) + ) + return model + + def test_hf_export_maps_embeddings_and_split_qkv(self): + position_embedding = _RecordingWeight() + qkv_weight = _RecordingWeight() + model = self._make_model( + [ + ("model.patch_generator.pos_embed", position_embedding), + ("model.encoder.layers.0.attn.attn.qkv_proj.weight", qkv_weight), + ] + ) + + position = torch.ones(1) + query, key, value = (torch.full((1,), value) for value in (2, 3, 4)) + loaded = model.load_weights( + [ + ("radio_model.hf_model.embeddings.position_embedding", position), + ( + "radio_model.hf_model.encoder.layer.0.attention.attention." + "query.weight", + query, + ), + ( + "radio_model.hf_model.encoder.layer.0.attention.attention." + "key.weight", + key, + ), + ( + "radio_model.hf_model.encoder.layer.0.attention.attention." + "value.weight", + value, + ), + ("radio_model.hf_model.summary_idxs", torch.tensor([0, 1])), + ] + ) + + self.assertEqual( + loaded, + { + "model.patch_generator.pos_embed", + "model.encoder.layers.0.attn.attn.qkv_proj.weight", + }, + ) + self.assertEqual( + position_embedding.loads, [(position_embedding, position, None)] + ) + self.assertEqual( + qkv_weight.loads, + [ + (qkv_weight, query, "q"), + (qkv_weight, key, "k"), + (qkv_weight, value, "v"), + ], + ) + + def test_hf_export_loads_encoder_parameters(self): + cases = { + "embeddings.video_patch_projection.weight": ( + "model.patch_generator.video_embedder.weight" + ), + "encoder.layer.1.attention.output.dense.weight": ( + "model.encoder.layers.1.attn.attn.proj.weight" + ), + "encoder.layer.2.layer_scale1.lambda1": "model.encoder.layers.2.ls1", + "encoder.layer.3.layer_scale2.lambda1": "model.encoder.layers.3.ls2", + "encoder.layer.4.mlp.fc1.bias": "model.encoder.layers.4.mlp.fc1.bias", + "encoder.layer.5.norm2.weight": "model.encoder.layers.5.norm2.weight", + } + for source, target in cases.items(): + with self.subTest(source=source): + parameter = _RecordingWeight() + model = self._make_model([(target, parameter)]) + weight = torch.ones(1) + + self.assertEqual( + model.load_weights( + [(f"radio_model.hf_model.{source}", weight)] + ), + {target}, + ) + self.assertEqual(parameter.loads, [(parameter, weight, None)]) + + def test_unmapped_hf_export_weight_raises(self): + model = self._make_model() + + with self.assertRaisesRegex(ValueError, "Unexpected HF RADIO weight"): + model.load_weights( + [("radio_model.hf_model.encoder.layer.0.unknown.weight", torch.ones(1))] + ) + + def test_legacy_unknown_weight_remains_ignored(self): + model = self._make_model() + + self.assertEqual( + model.load_weights([("radio_model.unknown.weight", torch.ones(1))]), + set(), + ) + + +if __name__ == "__main__": + unittest.main() From 3982dab2899fb14a255fb7f07d8f56738bd9a010 Mon Sep 17 00:00:00 2001 From: Ryan Stewart Date: Tue, 25 Aug 2026 23:35:58 +0000 Subject: [PATCH 06/13] organize Super VL tests by source component Signed-off-by: Ryan Stewart --- .../unit/configs/test_model_config.py | 4 + .../configs/test_nano_nemotron_vl_config.py | 52 ++++ .../managers/test_scheduler_moe_config.py | 4 +- .../unit/model_loader/test_modelopt_loader.py | 29 ++ .../test_nemotron_h_omni_registration.py | 289 ------------------ .../models/test_nemotron_h_weight_loading.py | 17 +- test/registered/unit/models/test_radio.py | 4 +- .../test_nano_nemotron_vl_processor.py | 24 ++ .../unit/server_args/test_server_args.py | 29 ++ test/registered/unit/test_model_overrides.py | 64 +++- 10 files changed, 204 insertions(+), 312 deletions(-) create mode 100644 test/registered/unit/configs/test_nano_nemotron_vl_config.py delete mode 100644 test/registered/unit/models/test_nemotron_h_omni_registration.py create mode 100644 test/registered/unit/multimodal/test_nano_nemotron_vl_processor.py diff --git a/test/registered/unit/configs/test_model_config.py b/test/registered/unit/configs/test_model_config.py index baca934fc569..370ac5009e93 100644 --- a/test/registered/unit/configs/test_model_config.py +++ b/test/registered/unit/configs/test_model_config.py @@ -7,6 +7,7 @@ ModelConfig, get_hybrid_layer_ids, is_embedding_gemma, + is_multimodal_model, ) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -54,6 +55,9 @@ def test_does_not_misclassify_causal_gemma3(self): class TestDraftModelConfig(CustomTestCase): + def test_nemotron_h_omni_is_multimodal(self): + self.assertTrue(is_multimodal_model(["NemotronH_Omni_Reasoning_V3"])) + def test_qwen35_mtp_depth_is_synced_to_text_config(self): config = object.__new__(ModelConfig) config.is_draft_model = True diff --git a/test/registered/unit/configs/test_nano_nemotron_vl_config.py b/test/registered/unit/configs/test_nano_nemotron_vl_config.py new file mode 100644 index 000000000000..50b3e1321237 --- /dev/null +++ b/test/registered/unit/configs/test_nano_nemotron_vl_config.py @@ -0,0 +1,52 @@ +"""Unit tests for Nano Nemotron VL configuration compatibility.""" + +import unittest + +from sglang.srt.configs.nano_nemotron_vl import ( + NemotronH_Omni_Reasoning_V3_Config, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +class TestNemotronHOmniConfig(CustomTestCase): + def test_uses_checkpoint_model_type(self): + config = NemotronH_Omni_Reasoning_V3_Config( + vision_config={"args": {"model": "radio"}}, + llm_config={}, + architectures=["NemotronH_Omni_Reasoning_V3"], + ) + + self.assertEqual(config.model_type, "nemotron_h_omni") + self.assertEqual(config.architectures, ["NemotronH_Omni_Reasoning_V3"]) + + def test_normalizes_current_nemotron_h_layer_names(self): + llm_config = { + "layers_block_type": ["linear_attention", "moe", "full_attention"], + "num_nextn_predict_layers": 1, + "mtp_layers_block_type": ["full_attention", "moe"], + } + + config = NemotronH_Omni_Reasoning_V3_Config( + vision_config={"args": {"model": "radio"}}, + llm_config=llm_config, + ) + + self.assertEqual( + config.llm_config.layers_block_type, + ["mamba", "moe", "attention"], + ) + self.assertEqual( + config.llm_config.mtp_layers_block_type, + ["attention", "moe"], + ) + self.assertEqual( + llm_config["layers_block_type"], + ["linear_attention", "moe", "full_attention"], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/managers/test_scheduler_moe_config.py b/test/registered/unit/managers/test_scheduler_moe_config.py index 1021480d0f88..9433638b9acb 100644 --- a/test/registered/unit/managers/test_scheduler_moe_config.py +++ b/test/registered/unit/managers/test_scheduler_moe_config.py @@ -48,7 +48,7 @@ def test_uses_canonical_text_config_for_llm_config_vlm(self): hf_text_config=llm_config, ) - init_moe.assert_called_once_with(scheduler.server_args) + init_moe.assert_called_once_with() self.assertFalse(scheduler.require_mlp_sync) def test_canonical_text_config_takes_precedence(self): @@ -60,7 +60,7 @@ def test_canonical_text_config_takes_precedence(self): hf_text_config=llm_config, ) - init_moe.assert_called_once_with(scheduler.server_args) + init_moe.assert_called_once_with() def test_skips_moe_initialization_for_non_moe_text_config(self): outer_config = SimpleNamespace() diff --git a/test/registered/unit/model_loader/test_modelopt_loader.py b/test/registered/unit/model_loader/test_modelopt_loader.py index ea642a0eb637..20ee58cef465 100644 --- a/test/registered/unit/model_loader/test_modelopt_loader.py +++ b/test/registered/unit/model_loader/test_modelopt_loader.py @@ -36,6 +36,8 @@ ) from sglang.srt.model_loader.weight_utils import get_quant_config from sglang.srt.models.minimax_m3 import MiniMaxM3SparseForCausalLM +from sglang.srt.models.nano_nemotron_vl import NemotronH_Omni_Reasoning_V3 +from sglang.srt.models.nemotron_h import NemotronHForCausalLM from sglang.srt.models.utils import WeightsMapper from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_cuda_ci @@ -696,6 +698,33 @@ def test_unquantized_modelopt_fp4_preserves_modelopt_workflows(self): class TestModelOptMixedPrecisionConfig(CustomTestCase): + def test_nemotron_h_omni_resolves_fused_qkv_from_split_layers(self): + quant_config = ModelOptMixedPrecisionConfig.from_config( + { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + f"language_model.model.layers.7.mixer.{projection}": { + "quant_algo": "FP8" + } + for projection in ("q_proj", "k_proj", "v_proj") + }, + "packed_modules_mapping": ( + NemotronH_Omni_Reasoning_V3.packed_modules_mapping + ), + } + ) + + self.assertEqual( + NemotronH_Omni_Reasoning_V3.packed_modules_mapping, + NemotronHForCausalLM.packed_modules_mapping, + ) + self.assertEqual( + quant_config._resolve_quant_algo( + "language_model.model.layers.7.mixer.qkv_proj" + ), + "FP8", + ) + def test_fp8_pb_wo_dispatches_to_native_block_fp8(self): quant_config = ModelOptMixedPrecisionConfig.from_config( { diff --git a/test/registered/unit/models/test_nemotron_h_omni_registration.py b/test/registered/unit/models/test_nemotron_h_omni_registration.py deleted file mode 100644 index 2a23c28dd5ef..000000000000 --- a/test/registered/unit/models/test_nemotron_h_omni_registration.py +++ /dev/null @@ -1,289 +0,0 @@ -import unittest -from types import SimpleNamespace -from unittest.mock import patch - -import torch -import torch.nn as nn - -from sglang.srt.arg_groups.overrides import ( - _FLASHINFER_ALLREDUCE_FUSION_ARCHS, - _MAMBA_EXTRA_BUFFER_ARCHS, - _MAMBA_RADIX_CACHE_ARCHS, - _MODEL_OVERRIDE_FNS, -) -from sglang.srt.configs.model_config import is_multimodal_model -from sglang.srt.configs.nano_nemotron_vl import ( - NemotronH_Omni_Reasoning_V3_Config, -) -from sglang.srt.layers.quantization.modelopt_quant import ( - ModelOptMixedPrecisionConfig, -) -from sglang.srt.models.nano_nemotron_vl import NemotronH_Omni_Reasoning_V3 -from sglang.srt.models.nemotron_h import NemotronHForCausalLM -from sglang.srt.models.radio import RadioModel, _map_hf_radio_weight_name -from sglang.srt.multimodal.processors.nano_nemotron_vl import ( - NanoNemotronVLImageProcessor, -) -from sglang.srt.server_args import ServerArgs -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase - -register_cpu_ci(est_time=10, suite="base-a-test-cpu") - - -class _RecordingWeight: - def __init__(self): - self.loads = [] - - def weight_loader(self, param, weight, shard_id=None): - self.loads.append((param, weight, shard_id)) - - -class TestNemotronHOmniRegistration(CustomTestCase): - def test_config_uses_checkpoint_model_type(self): - config = NemotronH_Omni_Reasoning_V3_Config( - vision_config={"args": {"model": "radio"}}, - llm_config={}, - architectures=["NemotronH_Omni_Reasoning_V3"], - ) - - self.assertEqual(config.model_type, "nemotron_h_omni") - self.assertEqual(config.architectures, ["NemotronH_Omni_Reasoning_V3"]) - - def test_config_normalizes_current_nemotron_h_layer_names(self): - llm_config = { - "layers_block_type": ["linear_attention", "moe", "full_attention"], - "num_nextn_predict_layers": 1, - "mtp_layers_block_type": ["full_attention", "moe"], - } - - config = NemotronH_Omni_Reasoning_V3_Config( - vision_config={"args": {"model": "radio"}}, - llm_config=llm_config, - ) - - self.assertEqual( - config.llm_config.layers_block_type, - ["mamba", "moe", "attention"], - ) - self.assertEqual( - config.llm_config.mtp_layers_block_type, - ["attention", "moe"], - ) - self.assertEqual( - llm_config["layers_block_type"], - ["linear_attention", "moe", "full_attention"], - ) - - def test_model_and_processor_register_new_architecture(self): - from sglang.srt.models.registry import ModelRegistry - - model_class, architecture = ModelRegistry.resolve_model_cls( - "NemotronH_Omni_Reasoning_V3" - ) - - self.assertIs(model_class, NemotronH_Omni_Reasoning_V3) - self.assertEqual(architecture, "NemotronH_Omni_Reasoning_V3") - self.assertIn(NemotronH_Omni_Reasoning_V3, NanoNemotronVLImageProcessor.models) - - def test_new_architecture_is_multimodal(self): - self.assertTrue(is_multimodal_model(["NemotronH_Omni_Reasoning_V3"])) - - def test_new_architecture_uses_nemotron_h_runtime_policy(self): - architecture = "NemotronH_Omni_Reasoning_V3" - - self.assertIn(architecture, _MODEL_OVERRIDE_FNS) - self.assertIn(architecture, _MAMBA_RADIX_CACHE_ARCHS) - self.assertIn(architecture, _MAMBA_EXTRA_BUFFER_ARCHS) - self.assertIn(architecture, _FLASHINFER_ALLREDUCE_FUSION_ARCHS) - - def test_mixed_precision_resolves_fused_qkv_from_split_layers(self): - quant_config = ModelOptMixedPrecisionConfig.from_config( - { - "quant_algo": "MIXED_PRECISION", - "quantized_layers": { - f"language_model.model.layers.7.mixer.{projection}": { - "quant_algo": "FP8" - } - for projection in ("q_proj", "k_proj", "v_proj") - }, - "packed_modules_mapping": ( - NemotronH_Omni_Reasoning_V3.packed_modules_mapping - ), - } - ) - - self.assertEqual( - NemotronH_Omni_Reasoning_V3.packed_modules_mapping, - NemotronHForCausalLM.packed_modules_mapping, - ) - self.assertEqual( - quant_config._resolve_quant_algo( - "language_model.model.layers.7.mixer.qkv_proj" - ), - "FP8", - ) - - def test_mamba_cache_chunk_size_uses_language_config(self): - server_args = object.__new__(ServerArgs) - server_args.get_model_config = lambda: SimpleNamespace( - hf_config=SimpleNamespace(), - hf_text_config=SimpleNamespace(mamba_chunk_size=128), - ) - - with patch( - "sglang.srt.server_args.resolved_view", - return_value=SimpleNamespace(page_size=64), - ): - self.assertEqual(server_args.mamba_cache_chunk_size, 128) - - def test_multimodal_wrapper_exposes_language_embed_and_head(self): - model = object.__new__(NemotronH_Omni_Reasoning_V3) - nn.Module.__init__(model) - embed = object() - head = object() - model.language_model = SimpleNamespace( - get_embed_and_head=lambda: (embed, head), - lm_head=head, - ) - - self.assertEqual(model.get_embed_and_head(), (embed, head)) - self.assertIs(model.lm_head, head) - - def test_super_vision_final_layernorm_is_loaded_and_applied(self): - model = object.__new__(NemotronH_Omni_Reasoning_V3) - nn.Module.__init__(model) - model.mlp1 = nn.Sequential() - model.vision_final_layernorm = nn.LayerNorm(2) - model.language_model = SimpleNamespace(load_weights=lambda weights: None) - model.vision_model = SimpleNamespace(load_weights=lambda weights: None) - model.sound_encoder = None - - weight = torch.tensor([2.0, 3.0]) - bias = torch.tensor([0.5, -0.5]) - model.load_weights( - [ - ("vision_projector.vision_final_layernorm.weight", weight), - ("vision_projector.vision_final_layernorm.bias", bias), - ] - ) - - features = torch.tensor([[1.0, 3.0]]) - expected = nn.functional.layer_norm(features, (2,), weight, bias) - torch.testing.assert_close(model._normalize_vision_features(features), expected) - - def test_super_hf_vision_and_projector_names_are_remapped(self): - remap = NemotronH_Omni_Reasoning_V3._remap_checkpoint_weight_name - - self.assertEqual( - remap("vision_model.embeddings.position_embedding"), - "vision_model.radio_model.hf_model.embeddings.position_embedding", - ) - self.assertEqual( - remap("vision_model.embeddings.video_patch_projection.weight"), - "vision_model.radio_model.hf_model.embeddings.video_patch_projection.weight", - ) - self.assertEqual( - remap("vision_projector.mlp1.linear1.weight"), - "mlp1.1.weight", - ) - self.assertEqual( - remap("vision_model.radio_model.model.patch_generator.pos_embed"), - "vision_model.radio_model.model.patch_generator.pos_embed", - ) - - def test_hf_radio_loader_maps_embeddings_and_split_qkv(self): - model = object.__new__(RadioModel) - nn.Module.__init__(model) - position_embedding = _RecordingWeight() - qkv_weight = _RecordingWeight() - model.named_parameters = lambda: iter( - [ - ("model.patch_generator.pos_embed", position_embedding), - ( - "model.encoder.layers.0.attn.attn.qkv_proj.weight", - qkv_weight, - ), - ] - ) - - position = torch.ones(1) - query, key, value = (torch.full((1,), value) for value in (2, 3, 4)) - loaded = model.load_weights( - [ - ( - "radio_model.hf_model.embeddings.position_embedding", - position, - ), - ( - "radio_model.hf_model.encoder.layer.0.attention.attention.query.weight", - query, - ), - ( - "radio_model.hf_model.encoder.layer.0.attention.attention.key.weight", - key, - ), - ( - "radio_model.hf_model.encoder.layer.0.attention.attention.value.weight", - value, - ), - ("radio_model.hf_model.summary_idxs", torch.tensor([0, 1])), - ] - ) - - self.assertEqual( - loaded, - { - "model.patch_generator.pos_embed", - "model.encoder.layers.0.attn.attn.qkv_proj.weight", - }, - ) - self.assertEqual( - position_embedding.loads, [(position_embedding, position, None)] - ) - self.assertEqual( - qkv_weight.loads, - [ - (qkv_weight, query, "q"), - (qkv_weight, key, "k"), - (qkv_weight, value, "v"), - ], - ) - - def test_hf_radio_name_mapping_covers_encoder_parameters(self): - cases = { - "embeddings.video_patch_projection.weight": ( - "model.patch_generator.video_embedder.weight", - None, - ), - "encoder.layer.1.attention.output.dense.weight": ( - "model.encoder.layers.1.attn.attn.proj.weight", - None, - ), - "encoder.layer.2.layer_scale1.lambda1": ( - "model.encoder.layers.2.ls1", - None, - ), - "encoder.layer.3.layer_scale2.lambda1": ( - "model.encoder.layers.3.ls2", - None, - ), - "encoder.layer.4.mlp.fc1.bias": ( - "model.encoder.layers.4.mlp.fc1.bias", - None, - ), - "encoder.layer.5.norm2.weight": ( - "model.encoder.layers.5.norm2.weight", - None, - ), - } - for source, expected in cases.items(): - with self.subTest(source=source): - self.assertEqual( - _map_hf_radio_weight_name(f"radio_model.hf_model.{source}"), - expected, - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/models/test_nemotron_h_weight_loading.py b/test/registered/unit/models/test_nemotron_h_weight_loading.py index 80287837db83..6c2bb15e3738 100644 --- a/test/registered/unit/models/test_nemotron_h_weight_loading.py +++ b/test/registered/unit/models/test_nemotron_h_weight_loading.py @@ -1,13 +1,4 @@ -""" -Unit tests for NemotronHForCausalLM.load_weights. - -Regression test for Nemotron-H expert scale checkpoint tensors that map to -parameters absent from the current runtime model. -""" - -from sglang.test.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=4, suite="base-a-test-cpu") +"""Unit tests for Nemotron-H target and MTP checkpoint weight loading.""" import unittest from types import SimpleNamespace @@ -16,6 +7,10 @@ from sglang.srt.models.nemotron_h import NemotronHForCausalLM from sglang.srt.models.nemotron_h_mtp import NemotronHForCausalLMMTP +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=4, suite="base-a-test-cpu") class _FakePPGroup: @@ -52,7 +47,7 @@ def weight_loader(self, param, loaded_weight, shard_id): self.loads.append((param, loaded_weight, shard_id)) -class TestNemotronHWeightLoading(unittest.TestCase): +class TestNemotronHWeightLoading(CustomTestCase): def _make_minimal_model( self, named_parameters=(), model_class=NemotronHForCausalLM ): diff --git a/test/registered/unit/models/test_radio.py b/test/registered/unit/models/test_radio.py index ff21ee80498d..3ef21eccba76 100644 --- a/test/registered/unit/models/test_radio.py +++ b/test/registered/unit/models/test_radio.py @@ -104,9 +104,7 @@ def test_hf_export_loads_encoder_parameters(self): weight = torch.ones(1) self.assertEqual( - model.load_weights( - [(f"radio_model.hf_model.{source}", weight)] - ), + model.load_weights([(f"radio_model.hf_model.{source}", weight)]), {target}, ) self.assertEqual(parameter.loads, [(parameter, weight, None)]) diff --git a/test/registered/unit/multimodal/test_nano_nemotron_vl_processor.py b/test/registered/unit/multimodal/test_nano_nemotron_vl_processor.py new file mode 100644 index 000000000000..cee894743245 --- /dev/null +++ b/test/registered/unit/multimodal/test_nano_nemotron_vl_processor.py @@ -0,0 +1,24 @@ +"""Unit tests for the Nano Nemotron VL processor registry.""" + +import unittest + +from sglang.srt.models.nano_nemotron_vl import NemotronH_Omni_Reasoning_V3 +from sglang.srt.multimodal.processors.nano_nemotron_vl import ( + NanoNemotronVLImageProcessor, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestNanoNemotronVLProcessor(CustomTestCase): + def test_supports_nemotron_h_omni(self): + self.assertIn( + NemotronH_Omni_Reasoning_V3, + NanoNemotronVLImageProcessor.models, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 9407f3b57eaa..8b4710c5fc59 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -206,6 +206,35 @@ def test_config_nested_dict_args_are_json(self): os.unlink(config_file) +class TestMambaCacheChunkSize(CustomTestCase): + def test_uses_canonical_text_config_for_composed_model(self): + server_args = object.__new__(ServerArgs) + server_args.get_model_config = lambda: SimpleNamespace( + hf_config=SimpleNamespace(mamba_chunk_size=32), + hf_text_config=SimpleNamespace(mamba_chunk_size=128), + ) + + with patch( + "sglang.srt.server_args.resolved_view", + return_value=SimpleNamespace(page_size=64), + ): + self.assertEqual(server_args.mamba_cache_chunk_size, 128) + + def test_preserves_text_only_config_behavior(self): + server_args = object.__new__(ServerArgs) + text_config = SimpleNamespace(mamba_chunk_size=128) + server_args.get_model_config = lambda: SimpleNamespace( + hf_config=text_config, + hf_text_config=text_config, + ) + + with patch( + "sglang.srt.server_args.resolved_view", + return_value=SimpleNamespace(page_size=64), + ): + self.assertEqual(server_args.mamba_cache_chunk_size, 128) + + class TestMmEncoderDataParallelLogging(CustomTestCase): def test_logs_when_encoder_dp_has_no_parallelism(self): server_args = ServerArgs( diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index b02515befcf8..479be95ba89f 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -981,6 +981,45 @@ def test_nemotron_h_explicit_split_attention_backend_wins(self): self.assertNotIn("attention_backend", overrides) self.assertNotIn("speculative_draft_attention_backend", overrides) + def test_nemotron_h_omni_uses_inner_text_config(self): + outer_config = SimpleNamespace( + architectures=["NemotronH_Omni_Reasoning_V3"], + quantization_config={"quant_algo": "NVFP4"}, + ) + model_config = SimpleNamespace( + quantization="modelopt", + hf_config=outer_config, + hf_text_config=SimpleNamespace(mlp_hidden_act="relu2"), + ) + server_args = SimpleNamespace( + quantization=None, + moe_runner_backend="auto", + moe_a2a_backend="none", + attention_backend=None, + get_model_config=lambda: model_config, + ) + + with ( + patch.object(overrides_module, "is_sm100_supported", return_value=False), + patch.object(overrides_module, "is_cuda", return_value=False), + ): + self.assertEqual( + collect_model_override_declarations( + "NemotronH_Omni_Reasoning_V3", + server_args, + outer_config, + ), + [ + ( + "_nemotron_h_overrides", + { + "quantization": "modelopt_fp4", + "moe_runner_backend": "flashinfer_cutlass", + }, + ) + ], + ) + def test_nemotron_h_w4a16_moe_rejects_a2a_backend(self): from sglang.srt.arg_groups.overrides import _nemotron_h_overrides @@ -1807,6 +1846,12 @@ def _view(arch="Qwen3MoeForCausalLM", **kw): _flashinfer_allreduce_fusion_auto_enable(_view()), {"flashinfer_allreduce_fusion_backend": "auto"}, ) + self.assertEqual( + _flashinfer_allreduce_fusion_auto_enable( + _view(arch="NemotronH_Omni_Reasoning_V3") + ), + {"flashinfer_allreduce_fusion_backend": "auto"}, + ) # guards: unsupported arch / tp==1 / dp attention / a2a backend self.assertEqual( _flashinfer_allreduce_fusion_auto_enable( @@ -2178,13 +2223,18 @@ def _view(arch, layer_types=None, **kw): ) # NemotronH routes through the pass (covered by the guard union, # not the branch chain — its hook invokes the handler) - self.assertEqual( - _mamba_radix_cache_resolution(_view("NemotronHForCausalLM")), - { - "uses_mamba_radix_cache": True, - "mamba_radix_cache_strategy": "extra_buffer", - }, - ) + for architecture in ( + "NemotronHForCausalLM", + "NemotronH_Omni_Reasoning_V3", + ): + with self.subTest(architecture=architecture): + self.assertEqual( + _mamba_radix_cache_resolution(_view(architecture)), + { + "uses_mamba_radix_cache": True, + "mamba_radix_cache_strategy": "extra_buffer", + }, + ) # GraniteMoeHybrid is guarded on mamba layer types self.assertEqual( _mamba_radix_cache_resolution( From d794cadd96a324e64cf19aed5885e443198cb13d Mon Sep 17 00:00:00 2001 From: Ryan Stewart Date: Wed, 26 Aug 2026 17:23:27 +0000 Subject: [PATCH 07/13] stream Nemotron-H VL weights during loading Signed-off-by: Ryan Stewart --- python/sglang/srt/models/nano_nemotron_vl.py | 48 ++++++++--------- .../unit/models/test_nano_nemotron_vl.py | 51 +++++++++++++++++++ 2 files changed, 72 insertions(+), 27 deletions(-) diff --git a/python/sglang/srt/models/nano_nemotron_vl.py b/python/sglang/srt/models/nano_nemotron_vl.py index 81d419ccfcbd..0c028522f188 100644 --- a/python/sglang/srt/models/nano_nemotron_vl.py +++ b/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 @@ -373,34 +374,27 @@ 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)) - else: - self._load_extra_weight(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): diff --git a/test/registered/unit/models/test_nano_nemotron_vl.py b/test/registered/unit/models/test_nano_nemotron_vl.py index 768a3e987977..f268b09a2638 100644 --- a/test/registered/unit/models/test_nano_nemotron_vl.py +++ b/test/registered/unit/models/test_nano_nemotron_vl.py @@ -87,6 +87,9 @@ def test_unexpected_checkpoint_weight_raises(self): nn.Module.__init__(model) model.mlp1 = nn.Sequential() model.vision_final_layernorm = nn.LayerNorm(2) + model.language_model = SimpleNamespace(load_weights=lambda weights: None) + model.vision_model = SimpleNamespace(load_weights=lambda weights: None) + model.sound_encoder = None cases = ( ("vision_projector.unknown.weight", "Unexpected Nemotron-H Omni"), @@ -99,6 +102,54 @@ def test_unexpected_checkpoint_weight_raises(self): with self.subTest(name=name), self.assertRaisesRegex(ValueError, message): model.load_weights([(name, torch.ones(1))]) + def test_language_weights_are_streamed_and_remaining_components_are_routed(self): + model = object.__new__(NemotronH_Omni_Reasoning_V3) + nn.Module.__init__(model) + model.mlp1 = nn.Sequential() + model.vision_final_layernorm = None + source_exhausted = False + loaded_language_weights = [] + loaded_vision_weights = [] + loaded_sound_weights = [] + + def source_weights(): + nonlocal source_exhausted + yield "language_model.model.layer.weight", torch.ones(1) + yield "vision_model.radio_model.encoder.weight", torch.ones(1) + yield "sound_encoder.projection.weight", torch.ones(1) + source_exhausted = True + + def load_language_weights(weights): + self.assertFalse(source_exhausted) + loaded_language_weights.append(next(weights)) + + def load_vision_weights(weights): + self.assertFalse(source_exhausted) + loaded_vision_weights.extend(weights) + + def load_sound_weights(weights): + self.assertFalse(source_exhausted) + loaded_sound_weights.extend(weights) + + model.language_model = SimpleNamespace(load_weights=load_language_weights) + model.vision_model = SimpleNamespace(load_weights=load_vision_weights) + model.sound_encoder = SimpleNamespace(load_weights=load_sound_weights) + + model.load_weights(source_weights()) + + self.assertTrue(source_exhausted) + self.assertEqual( + [name for name, _ in loaded_language_weights], ["model.layer.weight"] + ) + self.assertEqual( + [name for name, _ in loaded_vision_weights], + ["radio_model.encoder.weight"], + ) + self.assertEqual( + [name for name, _ in loaded_sound_weights], + ["sound_encoder.projection.weight"], + ) + if __name__ == "__main__": unittest.main() From 53ebe8a452c5e60a7f44f9a56e7af9f42ab196a3 Mon Sep 17 00:00:00 2001 From: Ryan Stewart Date: Thu, 27 Aug 2026 01:57:50 +0000 Subject: [PATCH 08/13] delegate DFlash layer capture for Nemotron-H VL Signed-off-by: Ryan Stewart --- python/sglang/srt/models/nano_nemotron_vl.py | 3 +++ test/registered/unit/models/test_nano_nemotron_vl.py | 12 ++++++++++++ 2 files changed, 15 insertions(+) diff --git a/python/sglang/srt/models/nano_nemotron_vl.py b/python/sglang/srt/models/nano_nemotron_vl.py index 0c028522f188..f35109ae1574 100644 --- a/python/sglang/srt/models/nano_nemotron_vl.py +++ b/python/sglang/srt/models/nano_nemotron_vl.py @@ -232,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] diff --git a/test/registered/unit/models/test_nano_nemotron_vl.py b/test/registered/unit/models/test_nano_nemotron_vl.py index f268b09a2638..329035225a85 100644 --- a/test/registered/unit/models/test_nano_nemotron_vl.py +++ b/test/registered/unit/models/test_nano_nemotron_vl.py @@ -37,6 +37,18 @@ def test_exposes_language_embed_and_head(self): self.assertEqual(model.get_embed_and_head(), (embed, head)) self.assertIs(model.lm_head, head) + def test_delegates_dflash_capture_to_language_model(self): + model = object.__new__(NemotronH_Omni_Reasoning_V3) + nn.Module.__init__(model) + captured_layer_ids = [] + model.language_model = SimpleNamespace( + set_dflash_layers_to_capture=captured_layer_ids.extend + ) + + model.set_dflash_layers_to_capture([1, 22, 43, 64, 85]) + + self.assertEqual(captured_layer_ids, [1, 22, 43, 64, 85]) + def test_vision_final_layernorm_is_loaded_and_applied(self): model = object.__new__(NemotronH_Omni_Reasoning_V3) nn.Module.__init__(model) From 2e8e5683e90462e1ff615d3deb047faec940ed60 Mon Sep 17 00:00:00 2001 From: Ryan Stewart Date: Thu, 27 Aug 2026 01:58:31 +0000 Subject: [PATCH 09/13] reuse multimodal embeddings in Nemotron-H MTP Signed-off-by: Ryan Stewart --- python/sglang/srt/models/nemotron_h_mtp.py | 16 ++++- .../unit/models/test_nemotron_h_mtp.py | 63 +++++++++++++++++++ 2 files changed, 78 insertions(+), 1 deletion(-) create mode 100644 test/registered/unit/models/test_nemotron_h_mtp.py diff --git a/python/sglang/srt/models/nemotron_h_mtp.py b/python/sglang/srt/models/nemotron_h_mtp.py index 458d1a39d5dd..9b2df060a474 100644 --- a/python/sglang/srt/models/nemotron_h_mtp.py +++ b/python/sglang/srt/models/nemotron_h_mtp.py @@ -293,7 +293,21 @@ def forward( inputs_embeds: torch.Tensor | None = None, ) -> torch.Tensor: if inputs_embeds is None: - inputs_embeds = self.get_input_embeddings(input_ids) + inputs_embeds = forward_batch.mm_input_embeds + if ( + forward_batch.forward_mode.is_extend() + and forward_batch.contains_mm_inputs() + and not forward_batch.forward_mode.is_draft_extend_v2() + ): + assert inputs_embeds is not None + last_indices = ( + forward_batch.extend_start_loc + forward_batch.extend_seq_lens - 1 + ).long() + inputs_embeds[last_indices] = self.get_input_embeddings( + input_ids[last_indices] + ) + if inputs_embeds is None: + inputs_embeds = self.get_input_embeddings(input_ids) hidden_states = forward_batch.spec_info.hidden_states residual = None diff --git a/test/registered/unit/models/test_nemotron_h_mtp.py b/test/registered/unit/models/test_nemotron_h_mtp.py new file mode 100644 index 000000000000..be96e300e118 --- /dev/null +++ b/test/registered/unit/models/test_nemotron_h_mtp.py @@ -0,0 +1,63 @@ +"""Unit tests for Nemotron-H MTP model behavior.""" + +import unittest +from types import SimpleNamespace + +import torch +import torch.nn as nn + +from sglang.srt.models.nemotron_h_mtp import NemotronHMultiTokenPredictor +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class _RecordingLayer(nn.Module): + def __init__(self): + super().__init__() + self.inputs_embeds = None + + def forward(self, *, inputs_embeds, hidden_states, residual, forward_batch): + self.inputs_embeds = inputs_embeds + return hidden_states, residual + + +class TestNemotronHMultiTokenPredictor(CustomTestCase): + def test_multimodal_prefill_reuses_target_embeddings(self): + model = object.__new__(NemotronHMultiTokenPredictor) + nn.Module.__init__(model) + model.embed_tokens = nn.Embedding(8, 2) + model.embed_tokens.weight.data.copy_(torch.arange(16).reshape(8, 2)) + model.pattern_len = 1 + layer = _RecordingLayer() + model.layers = nn.ModuleDict({"0": layer}) + + target_embeddings = torch.tensor( + [[101.0, 102.0], [103.0, 104.0], [105.0, 106.0]] + ) + forward_batch = SimpleNamespace( + mm_input_embeds=target_embeddings.clone(), + forward_mode=SimpleNamespace( + is_extend=lambda: True, + is_draft_extend_v2=lambda: False, + ), + contains_mm_inputs=lambda: True, + extend_start_loc=torch.tensor([0]), + extend_seq_lens=torch.tensor([3]), + spec_info=SimpleNamespace(hidden_states=torch.zeros(3, 2)), + ) + + model( + input_ids=torch.tensor([100, 101, 2]), + positions=torch.arange(3), + forward_batch=forward_batch, + ) + + expected = target_embeddings.clone() + expected[-1] = model.embed_tokens(torch.tensor(2)) + torch.testing.assert_close(layer.inputs_embeds, expected) + + +if __name__ == "__main__": + unittest.main() From 9c744d496b6d76cfc88c231844fba93d7f7ad4b0 Mon Sep 17 00:00:00 2001 From: rystewart-nvidia Date: Thu, 3 Sep 2026 15:11:20 +0000 Subject: [PATCH 10/13] map quantized MTP metadata for Nemotron-H Signed-off-by: rystewart-nvidia --- python/sglang/srt/models/nemotron_h_mtp.py | 5 ++++ .../unit/models/test_nemotron_h_mtp.py | 28 ++++++++++++++++++- 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/models/nemotron_h_mtp.py b/python/sglang/srt/models/nemotron_h_mtp.py index 6f3a08302ac3..58047f95b4de 100644 --- a/python/sglang/srt/models/nemotron_h_mtp.py +++ b/python/sglang/srt/models/nemotron_h_mtp.py @@ -38,6 +38,7 @@ NemotronHMoEDecoderLayer, ) from sglang.srt.models.nemotron_h_utils import is_attn_layer +from sglang.srt.models.utils import WeightsMapper from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import add_prefix @@ -323,6 +324,10 @@ def forward( class NemotronHForCausalLMMTP(NemotronHForCausalLM): + hf_to_sglang_mapper = NemotronHForCausalLM.hf_to_sglang_mapper | WeightsMapper( + orig_to_new_prefix={"language_model.mtp.": "mtp."} + ) + def __init__( self, config: NemotronHConfig, diff --git a/test/registered/unit/models/test_nemotron_h_mtp.py b/test/registered/unit/models/test_nemotron_h_mtp.py index d3973952de15..20e5bec04f48 100644 --- a/test/registered/unit/models/test_nemotron_h_mtp.py +++ b/test/registered/unit/models/test_nemotron_h_mtp.py @@ -6,7 +6,13 @@ import torch import torch.nn as nn -from sglang.srt.models.nemotron_h_mtp import NemotronHMultiTokenPredictor +from sglang.srt.layers.quantization.modelopt_quant import ( + ModelOptMixedPrecisionConfig, +) +from sglang.srt.models.nemotron_h_mtp import ( + NemotronHForCausalLMMTP, + NemotronHMultiTokenPredictor, +) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -86,5 +92,25 @@ def test_multimodal_prefill_reuses_target_embeddings(self): torch.testing.assert_close(layer.inputs_embeds, expected) +class TestNemotronHForCausalLMMTP(CustomTestCase): + def test_maps_quantized_mtp_metadata(self): + quant_config = ModelOptMixedPrecisionConfig.from_config( + { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "language_model.mtp.layers.0.mixer.q_proj": {"quant_algo": "FP8"} + }, + } + ) + quant_config.apply_weight_name_mapper( + NemotronHForCausalLMMTP.hf_to_sglang_mapper + ) + + self.assertEqual( + quant_config._resolve_quant_algo("mtp.layers.0.mixer.q_proj"), + "FP8", + ) + + if __name__ == "__main__": unittest.main() From 07f7f821b9792755b3657975650a03da4b6146c0 Mon Sep 17 00:00:00 2001 From: rystewart-nvidia Date: Fri, 4 Sep 2026 02:15:14 +0000 Subject: [PATCH 11/13] cleanup tests Signed-off-by: rystewart-nvidia --- .../configs/test_nano_nemotron_vl_config.py | 1 - .../managers/test_scheduler_moe_config.py | 78 ------------------- .../unit/model_loader/test_modelopt_loader.py | 5 -- .../unit/server_args/test_server_args.py | 14 ---- 4 files changed, 98 deletions(-) delete mode 100644 test/registered/unit/managers/test_scheduler_moe_config.py diff --git a/test/registered/unit/configs/test_nano_nemotron_vl_config.py b/test/registered/unit/configs/test_nano_nemotron_vl_config.py index 50b3e1321237..a8612ccb28c3 100644 --- a/test/registered/unit/configs/test_nano_nemotron_vl_config.py +++ b/test/registered/unit/configs/test_nano_nemotron_vl_config.py @@ -20,7 +20,6 @@ def test_uses_checkpoint_model_type(self): ) self.assertEqual(config.model_type, "nemotron_h_omni") - self.assertEqual(config.architectures, ["NemotronH_Omni_Reasoning_V3"]) def test_normalizes_current_nemotron_h_layer_names(self): llm_config = { diff --git a/test/registered/unit/managers/test_scheduler_moe_config.py b/test/registered/unit/managers/test_scheduler_moe_config.py deleted file mode 100644 index 9433638b9acb..000000000000 --- a/test/registered/unit/managers/test_scheduler_moe_config.py +++ /dev/null @@ -1,78 +0,0 @@ -"""Tests for scheduler MoE configuration discovery.""" - -import unittest -from types import SimpleNamespace -from unittest.mock import patch - -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel - -maybe_stub_sgl_kernel() - -import sglang.srt.managers.scheduler as scheduler_module -from sglang.srt.managers.scheduler import Scheduler - -register_cpu_ci(est_time=1, suite="base-a-test-cpu") - - -class TestSchedulerMoeConfig(CustomTestCase): - def _run_init(self, *, hf_config, hf_text_config): - scheduler = Scheduler.__new__(Scheduler) - scheduler.server_args = object() - scheduler.model_config = SimpleNamespace( - hf_config=hf_config, - hf_text_config=hf_text_config, - ) - - with ( - patch.object(scheduler_module, "initialize_moe_config") as init_moe, - patch.object(scheduler_module, "initialize_fp8_gemm_config"), - patch.object(scheduler_module, "initialize_fp4_gemm_config"), - patch.object(scheduler_module, "initialize_bf16_gemm_config"), - patch.object( - scheduler_module, - "require_mlp_sync", - return_value=False, - ), - ): - scheduler.init_moe_gemm_config() - - return scheduler, init_moe - - def test_uses_canonical_text_config_for_llm_config_vlm(self): - outer_config = SimpleNamespace() - llm_config = SimpleNamespace(num_experts_per_tok=8) - - scheduler, init_moe = self._run_init( - hf_config=outer_config, - hf_text_config=llm_config, - ) - - init_moe.assert_called_once_with() - self.assertFalse(scheduler.require_mlp_sync) - - def test_canonical_text_config_takes_precedence(self): - outer_config = SimpleNamespace(text_config=SimpleNamespace()) - llm_config = SimpleNamespace(num_experts_per_tok=8) - - scheduler, init_moe = self._run_init( - hf_config=outer_config, - hf_text_config=llm_config, - ) - - init_moe.assert_called_once_with() - - def test_skips_moe_initialization_for_non_moe_text_config(self): - outer_config = SimpleNamespace() - text_config = SimpleNamespace() - - _, init_moe = self._run_init( - hf_config=outer_config, - hf_text_config=text_config, - ) - - init_moe.assert_not_called() - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/unit/model_loader/test_modelopt_loader.py b/test/registered/unit/model_loader/test_modelopt_loader.py index bc6beded4d5c..4307c431e430 100644 --- a/test/registered/unit/model_loader/test_modelopt_loader.py +++ b/test/registered/unit/model_loader/test_modelopt_loader.py @@ -37,7 +37,6 @@ from sglang.srt.model_loader.weight_utils import get_quant_config from sglang.srt.models.minimax_m3 import MiniMaxM3SparseForCausalLM from sglang.srt.models.nano_nemotron_vl import NemotronH_Omni_Reasoning_V3 -from sglang.srt.models.nemotron_h import NemotronHForCausalLM from sglang.srt.models.utils import WeightsMapper from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_cuda_ci @@ -714,10 +713,6 @@ def test_nemotron_h_omni_resolves_fused_qkv_from_split_layers(self): } ) - self.assertEqual( - NemotronH_Omni_Reasoning_V3.packed_modules_mapping, - NemotronHForCausalLM.packed_modules_mapping, - ) self.assertEqual( quant_config._resolve_quant_algo( "language_model.model.layers.7.mixer.qkv_proj" diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 75c6ba946f14..cb4c5739c773 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -291,20 +291,6 @@ def test_uses_canonical_text_config_for_composed_model(self): ): self.assertEqual(mamba_cache_chunk_size(server_args), 128) - def test_preserves_text_only_config_behavior(self): - server_args = object.__new__(ServerArgs) - text_config = SimpleNamespace(mamba_chunk_size=128) - server_args._model_config = SimpleNamespace( - hf_config=text_config, - hf_text_config=text_config, - ) - - with patch( - "sglang.srt.arg_groups.overrides.resolved_view", - return_value=SimpleNamespace(page_size=64), - ): - self.assertEqual(mamba_cache_chunk_size(server_args), 128) - class TestMmEncoderDataParallelLogging(CustomTestCase): def test_logs_when_encoder_dp_has_no_parallelism(self): From faec5773b72fcb1cec1fd998fa6d8607d80c4355 Mon Sep 17 00:00:00 2001 From: Ryan Stewart Date: Fri, 4 Sep 2026 02:49:26 +0000 Subject: [PATCH 12/13] preserve existing config selection behavior Signed-off-by: Ryan Stewart --- python/sglang/srt/arg_groups/overrides.py | 2 +- python/sglang/srt/managers/scheduler.py | 6 +++++- .../unit/server_args/test_server_args.py | 16 ---------------- 3 files changed, 6 insertions(+), 18 deletions(-) diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 7d9fba992126..6aae2a2144d8 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -1857,7 +1857,7 @@ def mamba_cache_chunk_size(server_args: Any) -> int: # Must match sglang.kernels.ops.attention.fla.chunk_delta_h.CHUNK_SIZE FLA_CHUNK_SIZE = 64 - hf_config = model_config_of(server_args).hf_text_config + hf_config = model_config_of(server_args).hf_config chunk_size = getattr(hf_config, "mamba_chunk_size", FLA_CHUNK_SIZE) page_size = resolved_view(server_args).page_size assert ( diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 87b7291ce23c..5c2d2f6ac21e 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -926,7 +926,11 @@ def init_mamba_backend(self) -> None: initialize_mamba_selective_state_update_backend(self.server_args) def init_moe_gemm_config(self): - config_to_check = self.model_config.hf_text_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``, diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index cb4c5739c773..b81bdd2dfd24 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -41,7 +41,6 @@ ) from sglang.srt.arg_groups.overrides import ( cutedsl_moe_max_num_tokens, - mamba_cache_chunk_size, max_speculative_num_draft_tokens, resolution_result, ) @@ -277,21 +276,6 @@ def test_config_nested_dict_args_are_json(self): os.unlink(config_file) -class TestMambaCacheChunkSize(CustomTestCase): - def test_uses_canonical_text_config_for_composed_model(self): - server_args = object.__new__(ServerArgs) - server_args._model_config = SimpleNamespace( - hf_config=SimpleNamespace(mamba_chunk_size=32), - hf_text_config=SimpleNamespace(mamba_chunk_size=128), - ) - - with patch( - "sglang.srt.arg_groups.overrides.resolved_view", - return_value=SimpleNamespace(page_size=64), - ): - self.assertEqual(mamba_cache_chunk_size(server_args), 128) - - class TestMmEncoderDataParallelLogging(CustomTestCase): def test_logs_when_encoder_dp_has_no_parallelism(self): server_args = ServerArgs( From ee6482b26e75ecab4a2df30ec893762781c0c839 Mon Sep 17 00:00:00 2001 From: Ryan Stewart Date: Thu, 10 Sep 2026 17:13:15 +0000 Subject: [PATCH 13/13] preserve standalone LM heads for Nemotron-H MTP --- python/sglang/srt/models/nemotron_h_mtp.py | 53 ++++++++- .../unit/models/test_nemotron_h_mtp.py | 110 ++++++++++++++++++ 2 files changed, 159 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/models/nemotron_h_mtp.py b/python/sglang/srt/models/nemotron_h_mtp.py index b5a067b3b4a7..d64caae1fc88 100644 --- a/python/sglang/srt/models/nemotron_h_mtp.py +++ b/python/sglang/srt/models/nemotron_h_mtp.py @@ -27,6 +27,7 @@ from sglang.srt.layers.linear import ColumnParallelLinear from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.quantization import QuantizationConfig +from sglang.srt.layers.quantization.modelopt_quant import ModelOptNvFp4A16LinearMethod from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, @@ -338,6 +339,7 @@ def __init__( config = config.get_mtp_config() self.config = config self.quant_config = quant_config + self._owns_lm_head = False # Required for parent's load_weights self.pp_group = get_pp_group() @@ -385,13 +387,56 @@ def forward( def load_weights( self, weights: Iterable[tuple[str, torch.Tensor]], is_mtp: bool = False ): - weights = ( - (name.removeprefix("language_model."), weight) for name, weight in weights + has_mtp_layers = False + has_target_layers = False + head_weights = set() + + def normalized_weights(): + nonlocal has_mtp_layers, has_target_layers + for name, weight in weights: + name = name.removeprefix("language_model.") + has_mtp_layers |= name.startswith("mtp.layers.") + has_target_layers |= name.startswith( + ("backbone.layers.", "model.layers.") + ) + if name.startswith("lm_head."): + head_weights.add(name) + yield name, weight + + # Inspect names while streaming: buffering a full target checkpoint here + # would double its host-memory footprint during embedded MTP loading. + super().load_weights(normalized_weights(), is_mtp=True) + self._owns_lm_head = bool( + has_mtp_layers and not has_target_layers and head_weights ) - super().load_weights(weights, is_mtp=True) + if self._owns_lm_head: + expected = { + name + for name, _ in self.named_parameters() + if name.startswith("lm_head.") + } + if "lm_head.input_scale" in expected and isinstance( + self.lm_head.quant_method, ModelOptNvFp4A16LinearMethod + ): + # NVFP4A16 accepts this loader placeholder but never uses it. + expected.remove("lm_head.input_scale") + missing = (expected | {"lm_head.weight"}) - head_weights + if missing: + raise ValueError( + f"Incomplete standalone MTP lm_head: missing {sorted(missing)}" + ) + + def set_embed_and_head(self, embed, head): + if not self._owns_lm_head: + return super().set_embed_and_head(embed, head) + # Standalone MTP checkpoints can supply a differently quantized head. + # Share only the input embeddings; retain the entire loaded head module. + self.model.embed_tokens.weight = embed + torch.cuda.empty_cache() + torch.cuda.synchronize() def set_lm_head_from_target(self, target_lm_head: nn.Module) -> None: - if self.config.tie_word_embeddings: + if self.config.tie_word_embeddings or self._owns_lm_head: return self.lm_head = target_lm_head diff --git a/test/registered/unit/models/test_nemotron_h_mtp.py b/test/registered/unit/models/test_nemotron_h_mtp.py index 20e5bec04f48..d6de48d818ed 100644 --- a/test/registered/unit/models/test_nemotron_h_mtp.py +++ b/test/registered/unit/models/test_nemotron_h_mtp.py @@ -2,12 +2,14 @@ import unittest from types import SimpleNamespace +from unittest.mock import patch import torch import torch.nn as nn from sglang.srt.layers.quantization.modelopt_quant import ( ModelOptMixedPrecisionConfig, + ModelOptNvFp4A16LinearMethod, ) from sglang.srt.models.nemotron_h_mtp import ( NemotronHForCausalLMMTP, @@ -93,6 +95,114 @@ def test_multimodal_prefill_reuses_target_embeddings(self): class TestNemotronHForCausalLMMTP(CustomTestCase): + def _make_head_model(self): + model = object.__new__(NemotronHForCausalLMMTP) + nn.Module.__init__(model) + model.config = SimpleNamespace( + max_n_routed_experts=0, tie_word_embeddings=False + ) + model.pp_group = SimpleNamespace(is_first_rank=True, is_last_rank=True) + model.model = nn.Module() + model.model.embed_tokens = nn.Embedding(4, 2) + model.model.layers = nn.ModuleList([nn.Linear(2, 2, bias=False)]) + model.lm_head = nn.Linear(2, 4, bias=False) + model.lm_head.quant_method = None + model.lm_head.register_parameter( + "weight_scale", nn.Parameter(torch.zeros(1), requires_grad=False) + ) + return model + + def test_standalone_mtp_head_survives_both_target_sharing_calls(self): + # Replacing either the head weight or its module silently discards the + # external checkpoint's output projection (including quantization scales). + for prefix in ("", "language_model."): + with self.subTest(prefix=prefix): + model = self._make_head_model() + model.load_weights( + iter( + [ + (prefix + "mtp.layers.0.weight", torch.ones(2, 2)), + ( + prefix + "lm_head.weight", + torch.arange(8.0).reshape(4, 2), + ), + (prefix + "lm_head.weight_scale", torch.tensor([0.5])), + ] + ) + ) + draft_head = model.lm_head + draft_weight = draft_head.weight + target_embed = nn.Parameter(torch.ones(4, 2)) + target_head = nn.Linear(2, 4, bias=False) + with patch("torch.cuda.synchronize"), patch("torch.cuda.empty_cache"): + model.set_embed_and_head(target_embed, target_head.weight) + self.assertIs(model.lm_head.weight, draft_weight) + model.set_lm_head_from_target(target_head) + self.assertIs(model.lm_head, draft_head) + self.assertIs(model.model.embed_tokens.weight, target_embed) + torch.testing.assert_close( + model.lm_head(torch.ones(1, 2)), + torch.tensor([[1.0, 5.0, 9.0, 13.0]]), + ) + torch.testing.assert_close( + model.lm_head.weight_scale, torch.tensor([0.5]) + ) + + def test_embedded_and_headless_mtp_share_complete_target_head(self): + for embedded in (False, True): + with self.subTest(embedded=embedded): + model = self._make_head_model() + weights = [("mtp.layers.0.weight", torch.ones(2, 2))] + if embedded: + # Full checkpoints also contain lm_head tensors; their + # presence alone must not opt out of embedded head sharing. + weights += [ + ("lm_head.weight", torch.ones(4, 2)), + ("lm_head.weight_scale", torch.ones(1)), + ("backbone.layers.0.weight", torch.ones(2, 2)), + ] + model.load_weights(iter(weights)) + target_head = nn.Linear(2, 4, bias=False) + target_embed = nn.Parameter(torch.ones(4, 2)) + with patch("torch.cuda.synchronize"), patch("torch.cuda.empty_cache"): + model.set_embed_and_head(target_embed, target_head.weight) + model.set_lm_head_from_target(target_head) + self.assertIs(model.lm_head, target_head) + self.assertIs(model.model.embed_tokens.weight, target_embed) + + def test_incomplete_standalone_head_is_rejected(self): + for missing in ("weight", "weight_scale"): + with self.subTest(missing=missing): + model = self._make_head_model() + weights = { + "mtp.layers.0.weight": torch.ones(2, 2), + "lm_head.weight": torch.ones(4, 2), + "lm_head.weight_scale": torch.ones(1), + } + del weights["lm_head." + missing] + with self.assertRaisesRegex( + ValueError, "Incomplete standalone MTP lm_head" + ): + model.load_weights(iter(weights.items())) + + def test_w4a16_head_does_not_require_unused_input_scale(self): + model = self._make_head_model() + model.lm_head.quant_method = ModelOptNvFp4A16LinearMethod(quant_config=None) + model.lm_head.register_parameter( + "input_scale", nn.Parameter(torch.zeros(1), requires_grad=False) + ) + # NVFP4A16 registers this loader placeholder but discards it before + # inference. Requiring it would reject valid standalone W4A16 heads. + model.load_weights( + iter( + [ + ("mtp.layers.0.weight", torch.ones(2, 2)), + ("lm_head.weight", torch.ones(4, 2)), + ("lm_head.weight_scale", torch.ones(1)), + ] + ) + ) + def test_maps_quantized_mtp_metadata(self): quant_config = ModelOptMixedPrecisionConfig.from_config( {