[ModelOpt FP4] Support online MoE weight quantization (#33115)

Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>

test/registered/models_e2e/test_qwen35_fp4_mtp.py

@@ -38,6 +38,8 @@ MTP_BASE_ARGS = [
     "trtllm_mha",
     "--quantization",
     "modelopt_fp4",
+    "--speculative-draft-model-quantization",
+    "modelopt_fp4",
     "--speculative-algorithm",
     "NEXTN",
     "--speculative-num-steps",

test/registered/unit/layers/quantization/test_modelopt_nvfp4.py

@@ -1,9 +1,16 @@
 import unittest
+from unittest.mock import patch
 
 import torch
+import torch.nn as nn
 
 from sglang.srt.layers.linear import MergedColumnParallelLinear, QKVParallelLinear
 from sglang.srt.layers.parameter import PerTensorScaleParameter
+from sglang.srt.layers.quantization.modelopt_quant import (
+    ModelOptFp4Config,
+    ModelOptFp4LinearMethod,
+)
+from sglang.srt.model_loader.weight_utils import default_weight_loader
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -75,6 +82,56 @@ class TestModelOptNvfp4(CustomTestCase):
 
         torch.testing.assert_close(scale, torch.tensor([0.25, 0.5]))
 
+    def test_missing_input_scale_defaults_to_one_and_checkpoint_overwrites(self):
+        config = ModelOptFp4Config(
+            is_checkpoint_nvfp4_serialized=True,
+            group_size=16,
+            use_per_token_activation=False,
+        )
+        layer = nn.Module()
+        ModelOptFp4LinearMethod(config).create_weights(
+            layer,
+            input_size_per_partition=16,
+            output_partition_sizes=[16],
+            input_size=16,
+            output_size=16,
+            params_dtype=torch.bfloat16,
+            weight_loader=default_weight_loader,
+        )
+
+        torch.testing.assert_close(layer.input_scale, torch.ones(1))
+        default_weight_loader(layer.input_scale, torch.tensor(0.25))
+        torch.testing.assert_close(layer.input_scale, torch.tensor([0.25]))
+
+    @patch(
+        "sglang.srt.layers.quantization.modelopt_quant.envs."
+        "SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION.get",
+        return_value=True,
+    )
+    def test_modelopt_fp4_per_token_activation_contract(self, _):
+        # Serialized ModelOpt FP4 retains the existing environment-controlled
+        # per-token activation path.
+        serialized_config = ModelOptFp4Config(
+            is_checkpoint_nvfp4_serialized=True,
+            group_size=16,
+        )
+        # Online modelopt_fp4 always uses per-tensor activation scaling, even
+        # when the serialized-checkpoint environment switch is enabled.
+        online_config = ModelOptFp4Config(
+            is_checkpoint_nvfp4_serialized=False,
+            group_size=16,
+        )
+
+        self.assertTrue(serialized_config.use_per_token_activation)
+        self.assertFalse(online_config.use_per_token_activation)
+        # nvfp4_online is the public interface for online per-token scaling.
+        with self.assertRaisesRegex(ValueError, "Use nvfp4_online"):
+            ModelOptFp4Config(
+                is_checkpoint_nvfp4_serialized=False,
+                group_size=16,
+                use_per_token_activation=True,
+            )
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/model_loader/test_modelopt_loader.py

@@ -6,6 +6,7 @@ applies NVIDIA Model Optimizer quantization to models during loading.
 """
 
 import unittest
+from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import torch
@@ -25,7 +26,12 @@ from sglang.srt.layers.quantization.modelopt_quant import (
     ModelOptMixedPrecisionConfig,
     ModelOptNvFp4A16LinearMethod,
 )
-from sglang.srt.model_loader.loader import DefaultModelLoader, ModelOptModelLoader
+from sglang.srt.model_loader.loader import (
+    DefaultModelLoader,
+    ModelOptModelLoader,
+    get_model_loader,
+)
+from sglang.srt.model_loader.weight_utils import get_quant_config
 from sglang.srt.models.minimax_m3 import MiniMaxM3SparseForCausalLM
 from sglang.srt.models.utils import WeightsMapper
 from sglang.srt.utils import get_device
@@ -605,6 +611,86 @@ class TestParseQuantHfConfig(CustomTestCase):
         self.assertEqual(result["quant_method"], "gptq")
         self.assertNotIn("quant_algo", result)
 
+    def test_inherited_draft_modelopt_fp4_accepts_fp8_checkpoint(self):
+        # ServerArgs has already copied the target's modelopt_fp4 request to the
+        # draft. Compatible FP8 metadata must not replace it with plain fp8.
+        self.model_config.quantization = "modelopt_fp4"
+        self.model_config.is_draft_model = True
+        self.model_config.is_draft_quantization_explicit = False
+        with (
+            patch.object(
+                self.model_config,
+                "_parse_quant_hf_config",
+                return_value={"quant_method": "fp8"},
+            ),
+            patch.object(
+                self.model_config,
+                "_find_quant_modelslim_config",
+                return_value=None,
+            ),
+        ):
+            self.model_config._verify_quantization()
+
+        # Keeping modelopt_fp4 selects online FP8-to-NVFP4 conversion for
+        # eligible MoE experts; this test stops at quantization-method routing.
+        self.assertEqual(self.model_config.quantization, "modelopt_fp4")
+
+
+class TestModelOptFp4LoaderSelection(CustomTestCase):
+    def test_draft_modelopt_fp4_uses_checkpoint_exclusions(self):
+        cases = (
+            # Excluded MTP experts are unpacked, so an explicit draft request
+            # replaces the serialized config with online weight quantization.
+            ("explicit embedded draft", True, ["mtp.layers.0*"], False),
+            # MTP experts present in the serialized checkpoint stay serialized.
+            ("explicit serialized draft", True, [], True),
+            # Inherited target quantization does not override draft exclusions.
+            ("inherited embedded draft", False, ["mtp.layers.0*"], True),
+        )
+        for name, is_explicit, ignored_layers, is_serialized in cases:
+            with self.subTest(name=name):
+                model_config = SimpleNamespace(
+                    model_path="target-model",
+                    quantization="modelopt_fp4",
+                    is_draft_model=True,
+                    is_draft_quantization_explicit=is_explicit,
+                    hf_config=SimpleNamespace(
+                        quantization_config={
+                            "quant_algo": "NVFP4",
+                            "group_size": 16,
+                            "ignore": ignored_layers,
+                        }
+                    ),
+                )
+
+                config = get_quant_config(model_config, LoadConfig(), {})
+
+                self.assertEqual(config.get_name(), "modelopt_fp4")
+                self.assertEqual(config.is_checkpoint_nvfp4_serialized, is_serialized)
+
+    def test_unquantized_modelopt_fp4_preserves_modelopt_workflows(self):
+        model_config = SimpleNamespace(
+            quantization="modelopt_fp4",
+            _is_already_quantized=lambda: False,
+        )
+
+        # Online conversion runs through the regular per-layer weight loaders.
+        online_loader = get_model_loader(LoadConfig(), model_config)
+        self.assertIsInstance(online_loader, DefaultModelLoader)
+        self.assertNotIsInstance(online_loader, ModelOptModelLoader)
+
+        # Explicit ModelOpt checkpoint/export workflows still need its loader.
+        for option in (
+            "modelopt_checkpoint_restore_path",
+            "modelopt_checkpoint_save_path",
+            "modelopt_export_path",
+        ):
+            with self.subTest(option=option):
+                loader = get_model_loader(
+                    LoadConfig(**{option: "/tmp/modelopt"}), model_config
+                )
+                self.assertIsInstance(loader, ModelOptModelLoader)
+
 
 class TestModelOptMixedPrecisionConfig(CustomTestCase):
     def test_minimax_mixed_precision_resolves_runtime_names_and_mxfp8(self):
@@ -734,7 +820,11 @@ class TestModelOptMixedPrecisionConfig(CustomTestCase):
         return_value=True,
     )
     def test_explicit_nvfp4_per_token_activation_false_overrides_env(self, _):
-        config = ModelOptFp4Config(use_per_token_activation=False)
+        config = ModelOptFp4Config(
+            is_checkpoint_nvfp4_serialized=True,
+            group_size=16,
+            use_per_token_activation=False,
+        )
 
         self.assertFalse(config.use_per_token_activation)
 

test/registered/unit/server_args/test_server_args.py

@@ -1,3 +1,4 @@
+import dataclasses
 import importlib
 import json
 import os
@@ -81,6 +82,17 @@ class TestPrepareServerArgs(CustomTestCase):
                 return_hidden_states_mode="lst",
             )
 
+    def test_draft_quantization_explicitness_survives_asdict_round_trip(self):
+        inherited = ServerArgs(model_path="dummy", quantization="modelopt_fp4")
+        inherited._handle_missing_default_values()
+        self.assertEqual(inherited.speculative_draft_model_quantization, "modelopt_fp4")
+        self.assertFalse(inherited._speculative_draft_quantization_explicitly_set)
+
+        reconstructed = ServerArgs(**dataclasses.asdict(inherited))
+        reconstructed._handle_missing_default_values()
+
+        self.assertFalse(reconstructed._speculative_draft_quantization_explicitly_set)
+
     def test_config_nested_dict_args_are_json(self):
         with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
             f.write("mm-process-config:\n  image:\n    resize: 128\n")