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

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

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()