From 2cc52d83263281fbd8b9c3a33b9455df0c34c232 Mon Sep 17 00:00:00 2001 From: Daifeng Li <119935962+fengli1702@users.noreply.github.com> Date: Fri, 17 Apr 2026 07:51:32 +0800 Subject: [PATCH] feat: Support MXFP4 quantized dense models on AMD CDNA2/CDNA3 GPUs (#19143) --- 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, 648 insertions(+), 282 deletions(-) create mode 100644 python/sglang/srt/layers/quantization/petit_mxfp4.py create mode 100644 python/sglang/srt/layers/quantization/petit_nvfp4.py diff --git a/python/pyproject_other.toml b/python/pyproject_other.toml index 37de1f458..85591284e 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.2", + "petit_kernel==0.0.3", "wave-lang==3.8.2", ] diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 779551626..adec92367 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -946,6 +946,7 @@ class ModelConfig: "fbgemm_fp8", "w8a8_fp8", "petit_nvfp4", + "petit_mxfp4", "quark", "mxfp4", "auto-round", @@ -970,6 +971,7 @@ class ModelConfig: "qoq", "w4afp8", "petit_nvfp4", + "petit_mxfp4", "quark", "modelslim", ] @@ -978,6 +980,7 @@ 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 8a6b1b06e..83af6acbe 100644 --- a/python/sglang/srt/layers/quantization/__init__.py +++ b/python/sglang/srt/layers/quantization/__init__.py @@ -36,7 +36,8 @@ 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 import PetitNvFp4Config +from sglang.srt.layers.quantization.petit_mxfp4 import PetitMxfp4Config +from sglang.srt.layers.quantization.petit_nvfp4 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 @@ -72,6 +73,7 @@ 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 37b7fbc54..b0183760c 100644 --- a/python/sglang/srt/layers/quantization/petit.py +++ b/python/sglang/srt/layers/quantization/petit.py @@ -1,253 +1,11 @@ -# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/quantization/modelopt.py +"""Backward-compatible import shim. +Use `petit_nvfp4.py` for implementation. Keep this module for existing imports. +""" -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_nvfp4 import ( # noqa: F401 + PetitNvFp4Config, + PetitNvFp4LinearMethod, ) -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 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, - ) +__all__ = ["PetitNvFp4Config", "PetitNvFp4LinearMethod"] diff --git a/python/sglang/srt/layers/quantization/petit_mxfp4.py b/python/sglang/srt/layers/quantization/petit_mxfp4.py new file mode 100644 index 000000000..eb2095777 --- /dev/null +++ b/python/sglang/srt/layers/quantization/petit_mxfp4.py @@ -0,0 +1,209 @@ +# 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 new file mode 100644 index 000000000..fc538e21a --- /dev/null +++ b/python/sglang/srt/layers/quantization/petit_nvfp4.py @@ -0,0 +1,247 @@ +# 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 529869f24..0d674c6d5 100644 --- a/python/sglang/srt/layers/quantization/petit_utils.py +++ b/python/sglang/srt/layers/quantization/petit_utils.py @@ -1,53 +1,51 @@ -from typing import Optional +import logging +from typing import Any, Dict, 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_nvfp4_a16, process_nvfp4_scales, repack_nvfp4 + from petit_kernel import ( + mul_mxfp4_a16, + mul_nvfp4_a16, + process_mxfp4_scales, + process_nvfp4_scales, + repack_mxfp4, + 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 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 _require_petit_kernel() -> None: + if mul_nvfp4_a16 is None: + raise ValueError(_PETIT_INSTALL_ERROR) def _check_petit_nvfp4_supported( quant_method: str, group_size: Optional[int] ) -> tuple[bool, Optional[str]]: - if quant_method != "NVFP4": + if quant_method.upper() != "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) @@ -58,7 +56,108 @@ 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 @@ -72,7 +171,20 @@ def prepare_nvfp4_layer_for_petit(layer: torch.nn.Module) -> None: ) layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) - return + +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) def apply_petit_nvfp4_linear( @@ -84,6 +196,8 @@ 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,) @@ -102,3 +216,35 @@ 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 312b94103..c8e99b0dc 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -112,6 +112,7 @@ QUANTIZATION_CHOICES = [ "modelopt_fp4", "modelopt_mixed", "petit_nvfp4", + "petit_mxfp4", "w8a8_int8", "w8a8_fp8", "moe_wna16",