Fix ModelOpt mixed-precision NVFP4 routing

python/sglang/srt/configs/model_config.py

@@ -1147,10 +1147,16 @@ def _parse_modelopt_quant_config(self, quant_config_dict: dict) -> Optional[dict
 
         if quant_algo == "MIXED_PRECISION":
             architectures = getattr(self.hf_config, "architectures", []) or []
-            if getattr(self.hf_config, "model_type", None) == "nemotron_h" or any(
-                arch.startswith("NemotronH") for arch in architectures
+            has_per_module_rules = bool(json_quant_configs.get("quantized_layers"))
+            if (
+                has_per_module_rules
+                or getattr(self.hf_config, "model_type", None) == "nemotron_h"
+                or any(arch.startswith("NemotronH") for arch in architectures)
             ):
-                return {"quant_method": "modelopt_mixed", "quant_algo": quant_algo}
+                return {
+                    "quant_method": "modelopt_mixed",
+                    "quant_algo": quant_algo,
+                }
             return {"quant_method": "w4afp8", "quant_algo": quant_algo}
         elif quant_algo and ("FP4" in quant_algo or "NVFP4" in quant_algo):
             return {"quant_method": "modelopt_fp4", "quant_algo": quant_algo}

python/sglang/srt/layers/logits_processor.py

@@ -33,6 +33,7 @@
     get_dp_dtype,
     get_dp_hidden_size,
 )
+from sglang.srt.layers.quantization.unquant import UnquantizedEmbeddingMethod
 from sglang.srt.layers.triton_ops.softcap import softcap_inplace_logits as fused_softcap
 from sglang.srt.layers.utils.logprob import (
     InputLogprobsResult,
@@ -886,6 +887,8 @@ def _compute_lm_head(
         if hasattr(lm_head, "set_lora") and hasattr(lm_head, "apply_lora"):
             # This is a LoRA-wrapped module, use its forward method
             logits = lm_head(hidden_states)
+        elif not isinstance(lm_head.quant_method, UnquantizedEmbeddingMethod):
+            logits = lm_head.quant_method.apply(lm_head, hidden_states, embedding_bias)
         elif hasattr(lm_head, "weight"):
             # Normal linear layer
             if self.use_fp32_lm_head:

python/sglang/srt/layers/quantization/modelopt_quant.py

@@ -741,19 +741,21 @@ def get_quant_method(
     ) -> Optional[QuantizeMethodBase]:
         from sglang.srt.layers.linear import LinearBase
         from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
+        from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
 
         quant_algo = self._resolve_quant_algo(prefix)
 
-        if isinstance(layer, LinearBase):
+        is_lm_head = isinstance(layer, ParallelLMHead)
+        if isinstance(layer, LinearBase) or is_lm_head:
             if is_layer_skipped(
                 prefix, self.exclude_modules, self.packed_modules_mapping
             ) or self.is_layer_excluded(prefix):
-                return UnquantizedLinearMethod()
+                return None if is_lm_head else UnquantizedLinearMethod()
             if quant_algo == "FP8":
                 return ModelOptFp8LinearMethod(self.fp8_config)
-            if quant_algo == "NVFP4":
+            if quant_algo and quant_algo.endswith("NVFP4"):
                 return ModelOptFp4LinearMethod(self.nvfp4_config)
-            return UnquantizedLinearMethod()
+            return None if is_lm_head else UnquantizedLinearMethod()
 
         if self.kv_cache_quant_algo and isinstance(layer, RadixAttention):
             return ModelOptFp8KVCacheMethod(self.fp8_config)
@@ -763,7 +765,7 @@ def get_quant_method(
                 return None
             if quant_algo == "FP8":
                 return ModelOptFp8MoEMethod(self.fp8_config)
-            if quant_algo == "NVFP4":
+            if quant_algo and quant_algo.endswith("NVFP4"):
                 return ModelOptNvFp4FusedMoEMethod(self.nvfp4_config)
             return None
 

python/sglang/srt/layers/vocab_parallel_embedding.py

@@ -459,6 +459,9 @@ def weight_loader(self, param: Parameter, loaded_weight: torch.Tensor):
         # If parameter does not have output dim, then it should
         # be copied onto all gpus (e.g. g_idx for act_order gptq).
         if output_dim is None:
+            if param.data.numel() == 1 and loaded_weight.numel() == 1:
+                param.data.reshape(-1).copy_(loaded_weight.reshape(-1))
+                return
             assert param.data.shape == loaded_weight.shape
             param.data.copy_(loaded_weight)
             return

python/sglang/srt/models/qwen3_5.py

@@ -1377,18 +1377,19 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
                 if "mlp.experts" in name:
                     continue
 
-                name = name.replace(weight_name, param_name)
+                mapped_name = name.replace(weight_name, param_name)
                 # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
+                if mapped_name.endswith(".bias") and mapped_name not in params_dict:
                     continue
                 # Skip layers on other devices.
                 # if is_pp_missing_parameter(name, self):
                 #     continue
-                if name not in params_dict:
+                if mapped_name not in params_dict:
                     continue
-                param = params_dict[name]
+                param = params_dict[mapped_name]
                 weight_loader = getattr(param, "weight_loader")
                 weight_loader(param, loaded_weight, shard_id)
+                name = mapped_name
                 break
             else:
                 # Skip loading extra bias for GPTQ models.
@@ -1735,18 +1736,19 @@ def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
                 if "visual" in name or "mlp.experts" in name:
                     continue
 
-                name = name.replace(weight_name, param_name)
+                mapped_name = name.replace(weight_name, param_name)
                 # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
+                if mapped_name.endswith(".bias") and mapped_name not in params_dict:
                     continue
                 # Skip layers on other devices.
                 # if is_pp_missing_parameter(name, self):
                 #     continue
-                if name not in params_dict:
+                if mapped_name not in params_dict:
                     continue
-                param = params_dict[name]
+                param = params_dict[mapped_name]
                 weight_loader = getattr(param, "weight_loader")
                 weight_loader(param, loaded_weight, shard_id)
+                name = mapped_name
                 break
             else:
                 if "visual" in name:

test/registered/unit/layers/test_logits_processor_quantized_lm_head.py

@@ -0,0 +1,31 @@
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import torch
+
+from sglang.srt.layers.logits_processor import LogitsProcessor
+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 TestQuantizedLmHead(CustomTestCase):
+    def test_quantized_lm_head_uses_quant_method(self):
+        processor = LogitsProcessor.__new__(LogitsProcessor)
+        hidden_states = torch.randn(2, 8)
+        expected = torch.randn(2, 16)
+        quant_method = MagicMock()
+        quant_method.apply.return_value = expected
+        lm_head = SimpleNamespace(quant_method=quant_method)
+
+        actual = processor._compute_lm_head(hidden_states, lm_head)
+
+        self.assertIs(actual, expected)
+        quant_method.apply.assert_called_once_with(lm_head, hidden_states, None)
+
+
+if __name__ == "__main__":
+    import unittest
+
+    unittest.main()

test/registered/unit/model_loader/test_modelopt_loader.py

@@ -8,6 +8,7 @@
 import unittest
 from unittest.mock import MagicMock, patch
 
+import torch
 import torch.nn as nn
 
 from sglang.srt.configs.device_config import DeviceConfig
@@ -16,8 +17,13 @@
 from sglang.srt.layers.modelopt_utils import QUANT_CFG_CHOICES
 from sglang.srt.layers.quantization.modelopt_quant import (
     ModelOptFp4Config,
+    ModelOptFp4LinearMethod,
     ModelOptMixedPrecisionConfig,
 )
+from sglang.srt.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
 from sglang.srt.model_loader.loader import ModelOptModelLoader
 from sglang.srt.models.utils import WeightsMapper
 from sglang.srt.utils import get_device
@@ -626,14 +632,19 @@ def test_non_modelopt_quant_method_unchanged(self):
 
 
 class TestModelOptMixedPrecisionConfig(CustomTestCase):
-    def test_nemotron_mixed_precision_uses_modelopt_mixed(self):
+    def test_per_module_mixed_precision_uses_modelopt_mixed(self):
         model_config = ModelConfig.__new__(ModelConfig)
         model_config.hf_config = MagicMock()
-        model_config.hf_config.model_type = "nemotron_h"
-        model_config.hf_config.architectures = ["NemotronHForCausalLM"]
+        model_config.hf_config.model_type = "qwen3_5"
+        model_config.hf_config.architectures = ["Qwen3_5ForConditionalGeneration"]
 
         result = model_config._parse_modelopt_quant_config(
-            {"quantization": {"quant_algo": "MIXED_PRECISION"}}
+            {
+                "quantization": {
+                    "quant_algo": "MIXED_PRECISION",
+                    "quantized_layers": {"lm_head": {"quant_algo": "W4A16_NVFP4"}},
+                }
+            }
         )
 
         self.assertEqual(result["quant_method"], "modelopt_mixed")
@@ -652,6 +663,42 @@ def test_mixed_precision_uses_nvfp4_min_capability(self):
             ModelOptFp4Config.get_min_capability(),
         )
 
+    def test_nvidia_nvfp4_algorithm_label_quantizes_lm_head(self):
+        quant_config = ModelOptMixedPrecisionConfig.from_config(
+            {
+                "quant_algo": "MIXED_PRECISION",
+                "quantized_layers": {
+                    "lm_head": {
+                        "quant_algo": "W4A16_NVFP4",
+                        "group_size": 16,
+                    },
+                },
+                "packed_modules_mapping": {},
+            }
+        )
+        lm_head = ParallelLMHead.__new__(ParallelLMHead)
+        method = quant_config.get_quant_method(lm_head, "lm_head")
+        self.assertIsInstance(method, ModelOptFp4LinearMethod)
+
+    def test_unquantized_lm_head_uses_default_embedding_method(self):
+        quant_config = ModelOptMixedPrecisionConfig.from_config(
+            {
+                "quant_algo": "MIXED_PRECISION",
+                "quantized_layers": {
+                    "model.layers.0.mlp.down_proj": {"quant_algo": "FP8"},
+                },
+                "packed_modules_mapping": {},
+            }
+        )
+        lm_head = ParallelLMHead.__new__(ParallelLMHead)
+        self.assertIsNone(quant_config.get_quant_method(lm_head, "lm_head"))
+
+    def test_vocab_loader_accepts_equivalent_scalar_scale_shape(self):
+        layer = VocabParallelEmbedding.__new__(VocabParallelEmbedding)
+        param = nn.Parameter(torch.zeros(1), requires_grad=False)
+        layer.weight_loader(param, torch.tensor(0.25))
+        self.assertEqual(param.item(), 0.25)
+
     def test_mixed_precision_quant_layer_resolution_after_mapping(self):
         quant_config = ModelOptMixedPrecisionConfig.from_config(
             {