Revert "feat: Support MXFP4 quantized dense models on AMD CDNA2/CDNA3 GPUs (#19143)" (#23031)

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",
 ]
 

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"],
         }

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 import PetitNvFp4Config
-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
@@ -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,

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
+import logging
-    PetitNvFp4Config,
+from typing import Any, Dict, List, Optional
-    PetitNvFp4LinearMethod,
+
+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,
+        )

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

@@ -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,
-        )

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

@@ -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,
-        )

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

@@ -1,51 +1,53 @@
-import logging
+from typing import Optional
-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 (
+    from petit_kernel import mul_nvfp4_a16, process_nvfp4_scales, repack_nvfp4
-        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 _require_petit_kernel() -> None:
+    def prepare_nvfp4_layer_for_petit(layer: torch.nn.Module) -> None:
-    if mul_nvfp4_a16 is None:
+        raise ValueError(
-        raise ValueError(_PETIT_INSTALL_ERROR)
+            "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 "
+            "Petit currently only supports: NVFP4"
-            "quantizations in sglang. Please check the "
+            " 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)
 
-
+    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(
@@ -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)

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",