This commit is contained in:
YC Yen-Ching Tseng
@@ -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",
|
||||
]
|
||||
|
||||
|
||||
@@ -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"],
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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)
|
||||
|
||||
@@ -112,7 +112,6 @@ QUANTIZATION_CHOICES = [
|
||||
"modelopt_fp4",
|
||||
"modelopt_mixed",
|
||||
"petit_nvfp4",
|
||||
"petit_mxfp4",
|
||||
"w8a8_int8",
|
||||
"w8a8_fp8",
|
||||
"moe_wna16",
|
||||
|
||||
Reference in New Issue
Block a user