From f0f014816719bd3844de5cfe4a110770c7cca137 Mon Sep 17 00:00:00 2001 From: YC Yen-Ching Tseng Date: Fri, 17 Apr 2026 12:53:25 +0800 Subject: [PATCH] Revert "feat: Support MXFP4 quantized dense models on AMD CDNA2/CDNA3 GPUs (#19143)" (#23031) --- python/pyproject_other.toml | 2 +- python/sglang/srt/configs/model_config.py | 3 - .../srt/layers/quantization/__init__.py | 4 +- .../sglang/srt/layers/quantization/petit.py | 256 +++++++++++++++++- .../srt/layers/quantization/petit_mxfp4.py | 209 -------------- .../srt/layers/quantization/petit_nvfp4.py | 247 ----------------- .../srt/layers/quantization/petit_utils.py | 208 +++----------- python/sglang/srt/server_args.py | 1 - 8 files changed, 282 insertions(+), 648 deletions(-) delete mode 100644 python/sglang/srt/layers/quantization/petit_mxfp4.py delete mode 100644 python/sglang/srt/layers/quantization/petit_nvfp4.py diff --git a/python/pyproject_other.toml b/python/pyproject_other.toml index 85591284e..37de1f458 100755 --- a/python/pyproject_other.toml +++ b/python/pyproject_other.toml @@ -93,7 +93,7 @@ tracing = [ srt_hip = [ "sglang[runtime_common]", "torch", - "petit_kernel==0.0.3", + "petit_kernel==0.0.2", "wave-lang==3.8.2", ] diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index adec92367..779551626 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -946,7 +946,6 @@ class ModelConfig: "fbgemm_fp8", "w8a8_fp8", "petit_nvfp4", - "petit_mxfp4", "quark", "mxfp4", "auto-round", @@ -971,7 +970,6 @@ class ModelConfig: "qoq", "w4afp8", "petit_nvfp4", - "petit_mxfp4", "quark", "modelslim", ] @@ -980,7 +978,6 @@ class ModelConfig: "modelopt_fp4": ["modelopt"], "modelopt_mixed": ["modelopt"], "petit_nvfp4": ["modelopt"], - "petit_mxfp4": ["mxfp4", "quark"], "w8a8_int8": ["compressed-tensors", "compressed_tensors"], "w8a8_fp8": ["compressed-tensors", "compressed_tensors"], } diff --git a/python/sglang/srt/layers/quantization/__init__.py b/python/sglang/srt/layers/quantization/__init__.py index 83af6acbe..8a6b1b06e 100644 --- a/python/sglang/srt/layers/quantization/__init__.py +++ b/python/sglang/srt/layers/quantization/__init__.py @@ -36,8 +36,7 @@ from sglang.srt.layers.quantization.modelopt_quant import ( from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config from sglang.srt.layers.quantization.mxfp4 import Mxfp4Config -from sglang.srt.layers.quantization.petit_mxfp4 import PetitMxfp4Config -from sglang.srt.layers.quantization.petit_nvfp4 import PetitNvFp4Config +from sglang.srt.layers.quantization.petit import PetitNvFp4Config from sglang.srt.layers.quantization.qoq import QoQConfig from sglang.srt.layers.quantization.quark.quark import QuarkConfig from sglang.srt.layers.quantization.quark_int4fp8_moe import QuarkInt4Fp8Config @@ -73,7 +72,6 @@ BASE_QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = { "qoq": QoQConfig, "w4afp8": W4AFp8Config, "petit_nvfp4": PetitNvFp4Config, - "petit_mxfp4": PetitMxfp4Config, "fbgemm_fp8": FBGEMMFp8Config, "quark": QuarkConfig, "auto-round": AutoRoundConfig, diff --git a/python/sglang/srt/layers/quantization/petit.py b/python/sglang/srt/layers/quantization/petit.py index b0183760c..37b7fbc54 100644 --- a/python/sglang/srt/layers/quantization/petit.py +++ b/python/sglang/srt/layers/quantization/petit.py @@ -1,11 +1,253 @@ -"""Backward-compatible import shim. +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/modelopt.py -Use `petit_nvfp4.py` for implementation. Keep this module for existing imports. -""" -from sglang.srt.layers.quantization.petit_nvfp4 import ( # noqa: F401 - PetitNvFp4Config, - PetitNvFp4LinearMethod, +import logging +from typing import Any, Dict, List, Optional + +import regex as re +import torch +from torch.nn.parameter import Parameter + +from sglang.srt.layers.linear import LinearBase +from sglang.srt.layers.parameter import ModelWeightParameter, PerTensorScaleParameter +from sglang.srt.layers.quantization.base_config import ( + LinearMethodBase, + QuantizationConfig, + QuantizeMethodBase, ) +from sglang.srt.layers.quantization.petit_utils import ( + apply_petit_nvfp4_linear, + prepare_nvfp4_layer_for_petit, + verify_petit_nvfp4_supported, +) +from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod +from sglang.srt.layers.quantization.utils import is_layer_skipped +from sglang.srt.utils import is_hip -__all__ = ["PetitNvFp4Config", "PetitNvFp4LinearMethod"] +_is_hip = is_hip() + +# Initialize logger for the module +logger = logging.getLogger(__name__) + + +# Configuration class to support the NVFP4 quantized model generated by the ModelOpt quantization tool +class PetitNvFp4Config(QuantizationConfig): + """Config class for Petit FP4.""" + + def __init__( + self, + is_checkpoint_nvfp4_serialized: bool = False, + kv_cache_quant_algo: str = None, + group_size: int = None, + exclude_modules: List[str] = None, + ) -> None: + self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized + if is_checkpoint_nvfp4_serialized: + logger.warning( + "Detected nvfp4 checkpoint. Please note that the " + "format is experimental and subject to change." + ) + self.group_size = group_size + self.kv_cache_quant_algo = kv_cache_quant_algo + self.exclude_modules = exclude_modules + + @classmethod + def get_name(cls) -> str: + return "petit_nvfp4" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + # Petit supports the gfx90a and gfx942 GPUs + return 90 + + @classmethod + def get_config_filenames(cls) -> List[str]: + return ["hf_quant_config.json"] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> "PetitNvFp4Config": + quant_config = cls.get_from_keys(config, ["quantization"]) + quant_method = quant_config["quant_algo"] + group_size = quant_config.get("group_size", None) + verify_petit_nvfp4_supported(quant_method, group_size) + + is_checkpoint_nvfp4_serialized = "NVFP4" in quant_method + kv_cache_quant_algo = quant_config["kv_cache_quant_algo"] + if not kv_cache_quant_algo: + kv_cache_quant_algo = "auto" + exclude_modules = quant_config.get("exclude_modules", None) + if not (group_size and kv_cache_quant_algo and (exclude_modules is not None)): + logger.warning( + f"group_size: {group_size}," + f"kv_cache_quant_algo: {kv_cache_quant_algo}," + f"exclude_modules: {exclude_modules}" + ) + raise ValueError( + "NVFP4 quantization requires group size and " + "kv_cache_quant_algo specified in " + "hf_quant_config.json" + ) + return cls( + is_checkpoint_nvfp4_serialized, + kv_cache_quant_algo, + group_size, + exclude_modules, + ) + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant) -> Optional[str]: + can_convert = cls.is_petit_nvfp4_compatible(hf_quant_cfg) + if can_convert: + return cls.get_name() + return None + + @classmethod + def is_petit_nvfp4_compatible(cls, quant_config: Dict[str, Any]) -> bool: + quant_method = quant_config.get("quant_method", "").lower() + return _is_hip and quant_method == "modelopt" + + def is_layer_excluded(self, prefix: str, exclude_modules: list): + for pattern in exclude_modules: + regex_str = pattern.replace(".", r"\.").replace("*", r".*") + if re.fullmatch(regex_str, prefix): + return True + return False + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional["QuantizeMethodBase"]: + if isinstance(layer, LinearBase): + if is_layer_skipped(prefix, self.exclude_modules) or self.is_layer_excluded( + prefix, self.exclude_modules + ): + return UnquantizedLinearMethod() + return PetitNvFp4LinearMethod(self) + return None + + def get_scaled_act_names(self) -> List[str]: + return [] + + +class PetitNvFp4LinearMethod(LinearMethodBase): + """Linear method for NVFP4. + Supports loading NVFP4 checkpoints with the following structure: + + |Tensor Name | datatype | shape | + |----------------------------------------------------| + |input_scale | torch.float32 | scalar | + |weight | NVFP4(SE2M1) | [1, X, y/2] | + |weight_scale | FP8-E4M3 | [X, Y] | + |weight_scale_2 | torch.float32 | scalar | + + The weights are quantized per block of 16 elements. + Args: quant_config: The ModelOpt quantization config. + """ + + def __init__(self, quant_config: PetitNvFp4Config): + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + del input_size, output_size + if not self.quant_config.is_checkpoint_nvfp4_serialized: + raise ValueError( + "NVFP4 quantization was selected, " + " dynamic quantization is not supported." + ) + + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + layer.logical_widths = output_partition_sizes + + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + if input_size_per_partition % 16 != 0: + raise ValueError( + "Unsupported model when in features size is " "not multiple of 16" + ) + + weight_dtype = ( + torch.float8_e4m3fn + if self.quant_config.is_checkpoint_nvfp4_serialized + else params_dtype + ) + + weight = ModelWeightParameter( + data=torch.empty( + # 2 fp4 data is packed in one uint8 in the input dimension + output_size_per_partition, + input_size_per_partition // 2, + dtype=torch.uint8, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + input_scale = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + + layer.register_parameter("input_scale", input_scale) + + weight_scale_2 = PerTensorScaleParameter( + data=torch.empty(len(output_partition_sizes), dtype=torch.float32), + weight_loader=weight_loader, + ) + layer.register_parameter("weight_scale_2", weight_scale_2) + + weight_scale = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition // self.quant_config.group_size, + dtype=weight_dtype, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + + layer.register_parameter("weight_scale", weight_scale) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + input_scale_2 = layer.input_scale.max().to(torch.float32) + weight_scale_2 = layer.weight_scale_2.max().to(torch.float32) + layer.input_scale = Parameter(input_scale_2, requires_grad=False) + layer.weight_scale_2 = Parameter(weight_scale_2, requires_grad=False) + layer.alpha = Parameter( + layer.input_scale * layer.weight_scale_2, requires_grad=False + ) + + prepare_nvfp4_layer_for_petit(layer) + del layer.input_scale + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return apply_petit_nvfp4_linear( + input=x, + weight=layer.weight, + weight_scale=layer.weight_scale, + weight_scale_2=layer.weight_scale_2, + size_n=layer.output_size_per_partition, + size_k=layer.input_size_per_partition, + bias=bias, + ) diff --git a/python/sglang/srt/layers/quantization/petit_mxfp4.py b/python/sglang/srt/layers/quantization/petit_mxfp4.py deleted file mode 100644 index eb2095777..000000000 --- a/python/sglang/srt/layers/quantization/petit_mxfp4.py +++ /dev/null @@ -1,209 +0,0 @@ -# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/modelopt.py - -import logging -from typing import Any, Dict, List, Optional - -import torch -from torch.nn.parameter import Parameter - -from sglang.srt.layers.linear import LinearBase -from sglang.srt.layers.parameter import ModelWeightParameter -from sglang.srt.layers.quantization.base_config import ( - LinearMethodBase, - QuantizationConfig, - QuantizeMethodBase, -) -from sglang.srt.layers.quantization.petit_utils import ( - apply_petit_mxfp4_linear, - is_quark_mxfp4_compatible_config, - prepare_mxfp4_layer_for_petit, - verify_petit_mxfp4_supported, -) -from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod -from sglang.srt.layers.quantization.utils import is_layer_skipped -from sglang.srt.utils import is_hip - -_is_hip = is_hip() -logger = logging.getLogger(__name__) - - -class PetitMxfp4Config(QuantizationConfig): - """Config class for Petit MXFP4 linear inference on ROCm.""" - - def __init__( - self, - is_checkpoint_mxfp4_serialized: bool = False, - group_size: int = 32, - exclude_modules: Optional[List[str]] = None, - ) -> None: - self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized - self.group_size = group_size - self.exclude_modules = exclude_modules or [] - if is_checkpoint_mxfp4_serialized: - logger.warning( - "Detected mxfp4 checkpoint for petit kernel path. " - "This format is experimental and subject to change." - ) - - @classmethod - def get_name(cls) -> str: - return "petit_mxfp4" - - @classmethod - def get_supported_act_dtypes(cls) -> List[torch.dtype]: - # Petit MXFP4 kernel currently supports BF16 activations on ROCm. - return [torch.bfloat16] - - @classmethod - def get_min_capability(cls) -> int: - return 90 - - @classmethod - def get_config_filenames(cls) -> List[str]: - return ["hf_quant_config.json"] - - @classmethod - def from_config(cls, config: Dict[str, Any]) -> "PetitMxfp4Config": - quant_section = config.get("quantization", config) - quant_method = ( - quant_section.get("quant_algo") - or quant_section.get("quant_method") - or config.get("quant_method") - or "" - ) - group_size = quant_section.get("group_size", 32) - verify_petit_mxfp4_supported(quant_method, group_size, quant_config=config) - - exclude_modules = quant_section.get("exclude_modules", []) - quant_method_lower = str(quant_method).lower() - is_checkpoint_mxfp4_serialized = "mxfp4" in quant_method_lower or ( - quant_method_lower == "quark" and is_quark_mxfp4_compatible_config(config) - ) - return cls( - is_checkpoint_mxfp4_serialized=is_checkpoint_mxfp4_serialized, - group_size=group_size, - exclude_modules=exclude_modules, - ) - - @classmethod - def override_quantization_method(cls, hf_quant_cfg, user_quant) -> Optional[str]: - # Keep legacy MXFP4 flows unless user explicitly opts into petit_mxfp4. - if str(user_quant).lower() != cls.get_name(): - return None - if cls.is_petit_mxfp4_compatible(hf_quant_cfg): - return cls.get_name() - return None - - @classmethod - def is_petit_mxfp4_compatible(cls, quant_config: Dict[str, Any]) -> bool: - if not _is_hip or not quant_config: - return False - - quant_method = str(quant_config.get("quant_method", "")).lower() - quant_algo = str(quant_config.get("quant_algo", "")).lower() - nested_quant = quant_config.get("quantization") - if isinstance(nested_quant, dict): - quant_algo = str(nested_quant.get("quant_algo", quant_algo)).lower() - if "mxfp4" in quant_method or "mxfp4" in quant_algo: - return True - if quant_method == "quark": - return is_quark_mxfp4_compatible_config(quant_config) - return False - - def get_quant_method( - self, layer: torch.nn.Module, prefix: str - ) -> Optional["QuantizeMethodBase"]: - if isinstance(layer, LinearBase): - if is_layer_skipped(prefix, self.exclude_modules): - return UnquantizedLinearMethod() - return PetitMxfp4LinearMethod(self) - return None - - def get_scaled_act_names(self) -> List[str]: - return [] - - -class PetitMxfp4LinearMethod(LinearMethodBase): - """Linear method for MXFP4 weights + Petit kernel execution.""" - - def __init__(self, quant_config: PetitMxfp4Config): - self.quant_config = quant_config - - def create_weights( - self, - layer: torch.nn.Module, - input_size_per_partition: int, - output_partition_sizes: List[int], - input_size: int, - output_size: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ) -> None: - del input_size, output_size, params_dtype - - if not self.quant_config.is_checkpoint_mxfp4_serialized: - raise ValueError( - "MXFP4 quantization was selected, but dynamic quantization " - "is not supported for petit_mxfp4." - ) - - output_size_per_partition = sum(output_partition_sizes) - weight_loader = extra_weight_attrs.get("weight_loader") - layer.logical_widths = output_partition_sizes - layer.input_size_per_partition = input_size_per_partition - layer.output_size_per_partition = output_size_per_partition - - if input_size_per_partition % self.quant_config.group_size != 0: - raise ValueError( - "Unsupported model when in features size is not divisible by " - f"group_size={self.quant_config.group_size}." - ) - if input_size_per_partition % 2 != 0: - raise ValueError("MXFP4 packed weights require even K dimension.") - - weight = ModelWeightParameter( - data=torch.empty( - output_size_per_partition, - input_size_per_partition // 2, - dtype=torch.uint8, - ), - input_dim=1, - output_dim=0, - weight_loader=weight_loader, - ) - layer.register_parameter("weight", weight) - - weight_scale = ModelWeightParameter( - data=torch.empty( - output_size_per_partition, - input_size_per_partition // self.quant_config.group_size, - dtype=torch.uint8, - ), - input_dim=1, - output_dim=0, - weight_loader=weight_loader, - ) - layer.register_parameter("weight_scale", weight_scale) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - prepare_mxfp4_layer_for_petit(layer) - layer.weight_scale_2 = Parameter( - torch.ones(1, device=layer.weight.device, dtype=torch.float32), - requires_grad=False, - ) - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - return apply_petit_mxfp4_linear( - input=x, - weight=layer.weight, - weight_scale=layer.weight_scale, - global_scale=layer.weight_scale_2, - size_n=layer.output_size_per_partition, - size_k=layer.input_size_per_partition, - bias=bias, - ) diff --git a/python/sglang/srt/layers/quantization/petit_nvfp4.py b/python/sglang/srt/layers/quantization/petit_nvfp4.py deleted file mode 100644 index fc538e21a..000000000 --- a/python/sglang/srt/layers/quantization/petit_nvfp4.py +++ /dev/null @@ -1,247 +0,0 @@ -# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/modelopt.py - - -import logging -from typing import Any, Dict, List, Optional - -import torch -from torch.nn.parameter import Parameter - -from sglang.srt.layers.linear import LinearBase -from sglang.srt.layers.parameter import ModelWeightParameter, PerTensorScaleParameter -from sglang.srt.layers.quantization.base_config import ( - LinearMethodBase, - QuantizationConfig, - QuantizeMethodBase, -) -from sglang.srt.layers.quantization.petit_utils import ( - apply_petit_nvfp4_linear, - prepare_nvfp4_layer_for_petit, - verify_petit_nvfp4_supported, -) -from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod -from sglang.srt.layers.quantization.utils import is_layer_skipped -from sglang.srt.utils import is_hip - -_is_hip = is_hip() - -# Initialize logger for the module -logger = logging.getLogger(__name__) - - -# Configuration class to support the NVFP4 quantized model generated by the ModelOpt quantization tool -class PetitNvFp4Config(QuantizationConfig): - """Config class for Petit NVFP4. - - This config is intentionally NVFP4-only. MXFP4 is handled by - `petit_mxfp4` (see `petit_mxfp4.py`). - """ - - def __init__( - self, - is_checkpoint_nvfp4_serialized: bool = False, - kv_cache_quant_algo: str = None, - group_size: int = None, - exclude_modules: List[str] = None, - ) -> None: - self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized - - if is_checkpoint_nvfp4_serialized: - logger.warning( - "Detected nvfp4 checkpoint. Please note that the " - "format is experimental and subject to change." - ) - self.group_size = group_size - self.kv_cache_quant_algo = kv_cache_quant_algo - self.exclude_modules = exclude_modules - - @classmethod - def get_name(cls) -> str: - return "petit_nvfp4" - - @classmethod - def get_supported_act_dtypes(cls) -> List[torch.dtype]: - return [torch.bfloat16, torch.half] - - @classmethod - def get_min_capability(cls) -> int: - # Petit supports the gfx90a and gfx942 GPUs - return 90 - - @classmethod - def get_config_filenames(cls) -> List[str]: - return ["hf_quant_config.json"] - - @classmethod - def from_config(cls, config: Dict[str, Any]) -> "PetitNvFp4Config": - quant_config = cls.get_from_keys(config, ["quantization"]) - quant_method = str(quant_config["quant_algo"]) - quant_method_upper = quant_method.upper() - group_size = quant_config.get("group_size", None) - verify_petit_nvfp4_supported(quant_method, group_size) - - is_checkpoint_nvfp4_serialized = "NVFP4" in quant_method_upper - kv_cache_quant_algo = quant_config["kv_cache_quant_algo"] - if not kv_cache_quant_algo: - kv_cache_quant_algo = "auto" - exclude_modules = quant_config.get("exclude_modules", None) - if not (group_size and kv_cache_quant_algo and (exclude_modules is not None)): - logger.warning( - f"group_size: {group_size}," - f"kv_cache_quant_algo: {kv_cache_quant_algo}," - f"exclude_modules: {exclude_modules}" - ) - raise ValueError( - "NVFP4 quantization requires group size and kv_cache_quant_algo " - "specified in hf_quant_config.json" - ) - return cls( - is_checkpoint_nvfp4_serialized, - kv_cache_quant_algo, - group_size, - exclude_modules, - ) - - @classmethod - def override_quantization_method(cls, hf_quant_cfg, user_quant) -> Optional[str]: - can_convert = cls.is_petit_nvfp4_compatible(hf_quant_cfg) - if can_convert: - return cls.get_name() - return None - - @classmethod - def is_petit_nvfp4_compatible(cls, quant_config: Dict[str, Any]) -> bool: - quant_method = quant_config.get("quant_method", "").lower() - return _is_hip and quant_method == "modelopt" - - def get_quant_method( - self, layer: torch.nn.Module, prefix: str - ) -> Optional["QuantizeMethodBase"]: - if isinstance(layer, LinearBase): - if is_layer_skipped(prefix, self.exclude_modules): - return UnquantizedLinearMethod() - return PetitNvFp4LinearMethod(self) - return None - - def get_scaled_act_names(self) -> List[str]: - return [] - - -class PetitNvFp4LinearMethod(LinearMethodBase): - """Linear method for NVFP4. - - For NVFP4: - |Tensor Name | datatype | shape | - |----------------------------------------------------| - |input_scale | torch.float32 | scalar | - |weight | NVFP4(SE2M1) | [1, X, y/2] | - |weight_scale | FP8-E4M3 | [X, Y] | - |weight_scale_2 | torch.float32 | scalar | - - Args: quant_config: The Petit quantization config. - """ - - def __init__(self, quant_config: PetitNvFp4Config): - self.quant_config = quant_config - - def create_weights( - self, - layer: torch.nn.Module, - input_size_per_partition: int, - output_partition_sizes: List[int], - input_size: int, - output_size: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - del input_size, output_size - if not self.quant_config.is_checkpoint_nvfp4_serialized: - raise ValueError( - "NVFP4 quantization was selected, " - " dynamic quantization is not supported." - ) - - output_size_per_partition = sum(output_partition_sizes) - weight_loader = extra_weight_attrs.get("weight_loader") - - layer.logical_widths = output_partition_sizes - - layer.input_size_per_partition = input_size_per_partition - layer.output_size_per_partition = output_size_per_partition - if input_size_per_partition % 16 != 0: - raise ValueError( - "Unsupported model when in features size is " "not multiple of 16" - ) - - weight_dtype = ( - torch.float8_e4m3fn - if self.quant_config.is_checkpoint_nvfp4_serialized - else params_dtype - ) - - weight = ModelWeightParameter( - data=torch.empty( - # 2 fp4 data is packed in one uint8 in the input dimension - output_size_per_partition, - input_size_per_partition // 2, - dtype=torch.uint8, - ), - input_dim=1, - output_dim=0, - weight_loader=weight_loader, - ) - layer.register_parameter("weight", weight) - - input_scale = PerTensorScaleParameter( - data=torch.empty(len(output_partition_sizes), dtype=torch.float32), - weight_loader=weight_loader, - ) - - layer.register_parameter("input_scale", input_scale) - - weight_scale_2 = PerTensorScaleParameter( - data=torch.empty(len(output_partition_sizes), dtype=torch.float32), - weight_loader=weight_loader, - ) - layer.register_parameter("weight_scale_2", weight_scale_2) - - weight_scale = ModelWeightParameter( - data=torch.empty( - output_size_per_partition, - input_size_per_partition // self.quant_config.group_size, - dtype=weight_dtype, - ), - input_dim=1, - output_dim=0, - weight_loader=weight_loader, - ) - - layer.register_parameter("weight_scale", weight_scale) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - input_scale_2 = layer.input_scale.max().to(torch.float32) - weight_scale_2 = layer.weight_scale_2.max().to(torch.float32) - layer.input_scale = Parameter(input_scale_2, requires_grad=False) - layer.weight_scale_2 = Parameter(weight_scale_2, requires_grad=False) - layer.alpha = Parameter( - layer.input_scale * layer.weight_scale_2, requires_grad=False - ) - - prepare_nvfp4_layer_for_petit(layer) - del layer.input_scale - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - return apply_petit_nvfp4_linear( - input=x, - weight=layer.weight, - weight_scale=layer.weight_scale, - weight_scale_2=layer.weight_scale_2, - size_n=layer.output_size_per_partition, - size_k=layer.input_size_per_partition, - bias=bias, - ) diff --git a/python/sglang/srt/layers/quantization/petit_utils.py b/python/sglang/srt/layers/quantization/petit_utils.py index 0d674c6d5..529869f24 100644 --- a/python/sglang/srt/layers/quantization/petit_utils.py +++ b/python/sglang/srt/layers/quantization/petit_utils.py @@ -1,51 +1,53 @@ -import logging -from typing import Any, Dict, Optional +from typing import Optional import torch -_PETIT_INSTALL_ERROR = ( - "Petit is not installed. Please install it with `pip install petit-kernel`." -) -logger = logging.getLogger(__name__) - try: - from petit_kernel import ( - mul_mxfp4_a16, - mul_nvfp4_a16, - process_mxfp4_scales, - process_nvfp4_scales, - repack_mxfp4, - repack_nvfp4, - ) + from petit_kernel import mul_nvfp4_a16, process_nvfp4_scales, repack_nvfp4 except ImportError: - mul_mxfp4_a16 = None - mul_nvfp4_a16 = None - process_mxfp4_scales = None - process_nvfp4_scales = None - repack_mxfp4 = None - repack_nvfp4 = None + def _check_petit_nvfp4_supported( + quant_method: str, group_size: Optional[int] + ) -> tuple[bool, Optional[str]]: + return ( + False, + "Petit is not installed. Please install it with `pip install petit-kernel`.", + ) -def _require_petit_kernel() -> None: - if mul_nvfp4_a16 is None: - raise ValueError(_PETIT_INSTALL_ERROR) + def prepare_nvfp4_layer_for_petit(layer: torch.nn.Module) -> None: + raise ValueError( + "Petit is not installed. Please install it with `pip install petit-kernel`." + ) + + def apply_petit_nvfp4_linear( + input: torch.Tensor, + weight: torch.Tensor, + weight_scale: torch.Tensor, + weight_scale_2: torch.Tensor, + size_n: int, + size_k: int, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + raise ValueError( + "Petit is not installed. Please install it with `pip install petit-kernel`." + ) def _check_petit_nvfp4_supported( quant_method: str, group_size: Optional[int] ) -> tuple[bool, Optional[str]]: - if quant_method.upper() != "NVFP4": + if quant_method != "NVFP4": return ( False, - "Petit currently only supports: NVFP4 " - "quantizations in sglang. Please check the " + "Petit currently only supports: NVFP4" + " quantizations in sglang. Please check the " "`hf_quant_config.json` file for your model's " "quant configuration.", ) if group_size is not None and group_size != 16: return ( False, - "Petit currently only supports: group_size=16 quantizations.", + "Petit currently only supports: group_size=16" " quantizations.", ) return (True, None) @@ -56,108 +58,7 @@ def verify_petit_nvfp4_supported(quant_method: str, group_size: Optional[int]) - raise ValueError(error_msg) -def _check_petit_mxfp4_supported( - quant_method: str, - group_size: Optional[int], - quant_config: Optional[Dict[str, Any]] = None, -) -> tuple[bool, Optional[str]]: - quant_method_lower = quant_method.lower() - is_mxfp4_method = "mxfp4" in quant_method_lower - is_quark_method = quant_method_lower == "quark" - - if not is_mxfp4_method and not is_quark_method: - return ( - False, - "Petit MXFP4 currently only supports MXFP4 or Quark-MXFP4 quantizations " - "in sglang. Please check the model quantization config.", - ) - - if is_quark_method and quant_config is not None: - if not is_quark_mxfp4_compatible_config(quant_config): - return ( - False, - "Detected quant_method=quark, but the quark quantization config " - "does not look like MXFP4 weights " - "(fp4/per_group/group_size=32/e8m0).", - ) - - if group_size is not None and group_size != 32: - return ( - False, - "Petit MXFP4 currently only supports: group_size=32 quantizations.", - ) - return (True, None) - - -def verify_petit_mxfp4_supported( - quant_method: str, - group_size: Optional[int], - quant_config: Optional[Dict[str, Any]] = None, -) -> None: - supported, error_msg = _check_petit_mxfp4_supported( - quant_method, group_size, quant_config - ) - if not supported: - raise ValueError(error_msg) - - -def _is_quark_mxfp4_layer_quant_config(config: Dict[str, Any]) -> bool: - weight_quant = config.get("weight") - input_quant = config.get("input_tensors") - if not isinstance(weight_quant, dict): - return False - - if isinstance(input_quant, dict): - if hasattr(logger, "warning_once"): - logger.warning_once( - "Quark input_tensors quant config is ignored for petit_mxfp4 " - "(kernel path is w4a16). Only weight quant config is validated." - ) - else: - logger.warning( - "Quark input_tensors quant config is ignored for petit_mxfp4 " - "(kernel path is w4a16). Only weight quant config is validated." - ) - - return ( - weight_quant.get("dtype") == "fp4" - and weight_quant.get("qscheme") == "per_group" - and weight_quant.get("group_size") == 32 - and weight_quant.get("is_dynamic") is False - and weight_quant.get("scale_format") == "e8m0" - ) - - -def is_quark_mxfp4_compatible_config(quant_config: Dict[str, Any]) -> bool: - """Best-effort detection for Quark MXFP4 dense configs. - - Some checkpoints only expose `quant_method=quark` in config.json without - full layer quant metadata. In that case we return True and defer validation - to weight-loading/runtime checks when user explicitly selects petit_mxfp4. - """ - candidates: list[Dict[str, Any]] = [] - - global_quant = quant_config.get("global_quant_config") - if isinstance(global_quant, dict): - candidates.append(global_quant) - - layer_quant = quant_config.get("layer_quant_config") - if isinstance(layer_quant, dict): - candidates.extend(v for v in layer_quant.values() if isinstance(v, dict)) - - layer_type_quant = quant_config.get("layer_type_quant_config") - if isinstance(layer_type_quant, dict): - candidates.extend(v for v in layer_type_quant.values() if isinstance(v, dict)) - - if not candidates: - return True - - return all(_is_quark_mxfp4_layer_quant_config(cfg) for cfg in candidates) - - def prepare_nvfp4_layer_for_petit(layer: torch.nn.Module) -> None: - _require_petit_kernel() - # Repack weights to petit format part_size_n = layer.output_size_per_partition part_size_k = layer.input_size_per_partition @@ -171,20 +72,7 @@ def prepare_nvfp4_layer_for_petit(layer: torch.nn.Module) -> None: ) layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) - -def prepare_mxfp4_layer_for_petit(layer: torch.nn.Module) -> None: - _require_petit_kernel() - - part_size_n = layer.output_size_per_partition - part_size_k = layer.input_size_per_partition - qweight = layer.weight.view(torch.int32).contiguous() - petit_qweight = repack_mxfp4(qweight, size_n=part_size_n, size_k=part_size_k) - layer.weight = torch.nn.Parameter(petit_qweight, requires_grad=False) - - weight_scale = process_mxfp4_scales( - scales=layer.weight_scale, size_k=part_size_k, size_n=part_size_n - ) - layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) + return def apply_petit_nvfp4_linear( @@ -196,8 +84,6 @@ def apply_petit_nvfp4_linear( size_k: int, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: - _require_petit_kernel() - reshaped_x = input.reshape(-1, input.shape[-1]) out_shape = input.shape[:-1] + (size_n,) @@ -216,35 +102,3 @@ def apply_petit_nvfp4_linear( output.add_(bias) # In-place add return output.reshape(out_shape) - - -def apply_petit_mxfp4_linear( - input: torch.Tensor, - weight: torch.Tensor, - weight_scale: torch.Tensor, - size_n: int, - size_k: int, - bias: Optional[torch.Tensor] = None, - global_scale: Optional[torch.Tensor] = None, -) -> torch.Tensor: - _require_petit_kernel() - - reshaped_x = input.reshape(-1, input.shape[-1]) - out_shape = input.shape[:-1] + (size_n,) - if global_scale is None: - global_scale = torch.ones(1, device=reshaped_x.device, dtype=torch.float32) - - output = mul_mxfp4_a16( - a=reshaped_x, - b=weight, - s=weight_scale, - global_scale=global_scale, - size_m=reshaped_x.size(0), - size_n=size_n, - size_k=size_k, - solution_id=-1, - ) - if bias is not None: - output.add_(bias) - - return output.reshape(out_shape) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 769200fa5..cb58cfd8b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -112,7 +112,6 @@ QUANTIZATION_CHOICES = [ "modelopt_fp4", "modelopt_mixed", "petit_nvfp4", - "petit_mxfp4", "w8a8_int8", "w8a8_fp8", "moe_wna16",