diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index cee4efc39..fb42a72a7 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/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" - " quantizations in sglang. Please check the " + "ModelOpt currently only supports: FP8, NVFP4, NVFP4_AWQ " + "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 diff --git a/test/registered/unit/model_loader/test_modelopt_loader.py b/test/registered/unit/model_loader/test_modelopt_loader.py index 7cfe97f7e..04bee16b5 100644 --- a/test/registered/unit/model_loader/test_modelopt_loader.py +++ b/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 = {