[Quant] Support NVFP4_AWQ checkpoints in ModelOpt FP4 path (#31825)

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

@@ -1178,6 +1178,7 @@ class ModelOptFp4Config(ModelOptQuantConfig):
         exclude_modules: List[str] = None,
         packed_modules_mapping: Optional[Dict[str, List[str]]] = None,
         use_per_token_activation: Optional[bool] = None,
+        is_awq: bool = False,
     ) -> None:
         super().__init__(kv_cache_quant_algo, exclude_modules, packed_modules_mapping)
         self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized
@@ -1186,6 +1187,7 @@ class ModelOptFp4Config(ModelOptQuantConfig):
                 "Detected nvfp4 checkpoint. Please note that the "
                 "format is experimental and subject to change."
             )
+        self.is_awq = is_awq
         self.group_size = group_size
         self.use_per_token_activation = (
             envs.SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION.get()
@@ -1290,6 +1292,9 @@ class ModelOptFp4Config(ModelOptQuantConfig):
                     first_group = next(iter(config_groups.values()), {})
                     weights_config = first_group.get("weights", {})
                     group_size = weights_config.get("group_size")
+            # NVFP4 (incl. NVFP4_AWQ) always uses group_size 16
+            if group_size is None and quant_method and "NVFP4" in quant_method:
+                group_size = 16
 
             exclude_modules = config.get("ignore", [])
         else:
@@ -1308,10 +1313,10 @@ class ModelOptFp4Config(ModelOptQuantConfig):
                     "Expected either flat format (config.json) or nested format (hf_quant_config.json)."
                 )
 
-        if quant_method not in ["FP8", "NVFP4"]:
+        if quant_method not in ["FP8", "NVFP4", "NVFP4_AWQ"]:
             raise ValueError(
-                "ModelOpt currently only supports: FP8, NVFP4"
+                "ModelOpt currently only supports: FP8, NVFP4, NVFP4_AWQ "
-                " quantizations in sglang. Please check the "
+                "quantizations in sglang. Please check the "
                 "quantization config for your model's configuration."
             )
         is_checkpoint_nvfp4_serialized = "NVFP4" in quant_method
@@ -1332,6 +1337,7 @@ class ModelOptFp4Config(ModelOptQuantConfig):
             group_size,
             exclude_modules,
             config.get("packed_modules_mapping"),
+            is_awq="AWQ" in quant_method,
         )
 
     def get_quant_method(self, layer: torch.nn.Module, prefix: str):
@@ -1457,6 +1463,18 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
         )
         layer.register_parameter("input_scale", input_scale)
 
+        # NVFP4_AWQ: per-input-channel activation pre-scale baked into the weights
+        # offline. Length == input_size_per_partition; shards along the input dim
+        # (input_dim=0) so it splits correctly on row-parallel linears.
+        if self.quant_config.is_awq:
+            pre_quant_scale = ModelWeightParameter(
+                data=torch.ones(input_size_per_partition, dtype=params_dtype),
+                input_dim=0,
+                output_dim=0,
+                weight_loader=weight_loader,
+            )
+            layer.register_parameter("pre_quant_scale", pre_quant_scale)
+
         weight_scale_2 = _make_per_tensor_scale_parameter(
             (len(output_partition_sizes),),
             weight_loader=weight_loader,
@@ -1676,6 +1694,9 @@ class ModelOptFp4LinearMethod(LinearMethodBase):
             x_m = x_fp4.shape[0]
             output_dtype = layer.params_dtype
         else:
+            # NVFP4_AWQ: apply the per-input-channel pre_quant_scale.
+            if self.quant_config.is_awq:
+                x = x * layer.pre_quant_scale
             x_fp4, x_scale_interleaved = fp4_quantize(x, layer.input_scale_inv)
             x_m, _ = x.shape
             output_dtype = x.dtype

test/registered/unit/model_loader/test_modelopt_loader.py

@@ -456,6 +456,7 @@ class TestParseQuantHfConfig(CustomTestCase):
         ({"quant_method": "modelopt", "quant_algo": "FP8"}, "modelopt_fp8"),
         ({"quant_method": "modelopt", "quant_algo": "FP4"}, "modelopt_fp4"),
         ({"quant_method": "modelopt", "quant_algo": "NVFP4"}, "modelopt_fp4"),
+        ({"quant_algo": "NVFP4_AWQ"}, "modelopt_fp4"),
         ({"quant_method": "modelopt", "quant_algo": "MIXED_PRECISION"}, "w4afp8"),
         ({"quant_algo": "FP8"}, "modelopt_fp8"),
         ({"quant_algo": "FP4"}, "modelopt_fp4"),
@@ -493,6 +494,18 @@ class TestParseQuantHfConfig(CustomTestCase):
                 result = self.model_config._parse_quant_hf_config()
                 self.assertEqual(result["quant_method"], expected)
 
+    def test_awq_flat_config_defaults_group_size(self):
+        """NVFP4_AWQ flat config.json omits group_size; from_config must default it to 16."""
+        cfg = ModelOptFp4Config.from_config(
+            {
+                "quant_algo": "NVFP4_AWQ",
+                "ignore": ["lm_head"],
+                "quant_method": "modelopt",
+            }
+        )
+        self.assertEqual(cfg.group_size, 16)
+        self.assertTrue(cfg.is_awq)
+
     def test_non_modelopt_quant_method_unchanged(self):
         """Non-modelopt quant_method (e.g. 'gptq') must NOT enter the modelopt path."""
         self.model_config.hf_config.quantization_config = {