diff --git a/docs_new/docs/advanced_features/quantization.mdx b/docs_new/docs/advanced_features/quantization.mdx index 935651e9c..8fa3afc38 100644 --- a/docs_new/docs/advanced_features/quantization.mdx +++ b/docs_new/docs/advanced_features/quantization.mdx @@ -14,7 +14,7 @@ on-the-fly to convert high-precision weights into a lower-precision format. **Note: For better performance, usability and convenience, offline quantization is recommended over online quantization.** -If you use a pre-quantized model, do not add `--quantization` to enable online quantization at the same time. +If you use a pre-quantized model, **do not add `--quantization` to enable online quantization at the same time**. For popular pre-quantized models, please visit [Unsloth](https://huggingface.co/unsloth), [NVIDIA ModelOpt](https://huggingface.co/collections/nvidia/inference-optimized-checkpoints-with-model-optimizer) or [NeuralMagic](https://huggingface.co/collections/neuralmagic) collections on HF for some popular quality validated quantized models. Quantized models must be validated via benchmarks post-quantization @@ -804,6 +804,18 @@ SGLang running on AMD GPUs (CDNA3 or CDNA4 architecture) supports the quantizati Other layers (e.g. projections in the attention layers) have their weights quantized online to float8 directly. +### `quark_mxfp4` online quantization method + +SGLang running on AMD GPUs with hardware FP4 support (CDNA4 architecture, e.g. MI355x) supports the quantization method `--quantization quark_mxfp4`, that will quantize BF16 model weights to MXFP4 at load time, use dynamic MXFP4 quantization for activations and MXFP4 GEMMs instead of BF16 GEMMs. + +Example: + +```bash +sglang serve --model-path Qwen/Qwen3-30B-A3B \ + --tensor-parallel-size 1 \ + --quantization quark_mxfp4 +``` + ## Reference - [GPTQModel](https://github.com/ModelCloud/GPTQModel) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index f01ec9f6e..34798c3d3 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1145,6 +1145,7 @@ class ModelConfig: "mxfp4", "auto-round", "quark_int4fp8_moe", + "quark_mxfp4", ] optimized_quantization_methods = [ "fp8", @@ -1167,6 +1168,7 @@ class ModelConfig: "petit_nvfp4", "quark", "modelslim", + "quark_mxfp4", ] compatible_quantization_methods = { "modelopt_fp8": ["modelopt"], diff --git a/python/sglang/srt/constants.py b/python/sglang/srt/constants.py index 76d3b1cd3..2d47c6337 100644 --- a/python/sglang/srt/constants.py +++ b/python/sglang/srt/constants.py @@ -10,3 +10,5 @@ GPU_MEMORY_ALL_TYPES = [ ] HEALTH_CHECK_RID_PREFIX = "HEALTH_CHECK" + +GIB_BYTES = 1073741824 # 1024**3 diff --git a/python/sglang/srt/layers/quantization/__init__.py b/python/sglang/srt/layers/quantization/__init__.py index 8fdcb737e..ae30f2fb8 100644 --- a/python/sglang/srt/layers/quantization/__init__.py +++ b/python/sglang/srt/layers/quantization/__init__.py @@ -90,6 +90,7 @@ BASE_QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = { "petit_nvfp4": PetitNvFp4Config, "fbgemm_fp8": FBGEMMFp8Config, "quark": QuarkConfig, + "quark_mxfp4": QuarkConfig, "auto-round": AutoRoundConfig, "modelslim": ModelSlimConfig, "quark_int4fp8_moe": QuarkInt4Fp8Config, diff --git a/python/sglang/srt/layers/quantization/quark/quark.py b/python/sglang/srt/layers/quantization/quark/quark.py index 4b0f673a0..33ef806d6 100644 --- a/python/sglang/srt/layers/quantization/quark/quark.py +++ b/python/sglang/srt/layers/quantization/quark/quark.py @@ -29,6 +29,8 @@ from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.utils import get_device_capability if TYPE_CHECKING: + from transformers import PretrainedConfig + from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput __all__ = ["QuarkLinearMethod", "QuarkFusedMoEMethod"] @@ -40,21 +42,47 @@ class QuarkConfig(QuantizationConfig): def __init__( self, - quant_config: dict[str, Any], + quant_config: Optional[dict[str, Any]] = None, + hf_config: "PretrainedConfig | None" = None, kv_cache_group: Optional[list[str]] = None, kv_cache_config: Optional[dict[str, Any]] = None, pack_method: str = "reorder", + is_prequantized: bool = False, + online_scheme: Optional[str] = None, ): super().__init__() if kv_cache_group is None: kv_cache_group = [] + + if online_scheme is not None: + assert not is_prequantized + if online_scheme == "quark_mxfp4": + quant_config = self._create_online_mxfp4_config( + model_type=hf_config.model_type + ) + else: + raise ValueError(f"Unsupported online_scheme: {online_scheme}") + + if quant_config is None: + raise ValueError("Either quant_config or online_scheme must be provided") + self.quant_config = quant_config self.kv_cache_group = kv_cache_group self.kv_cache_config = kv_cache_config self.pack_method = pack_method self.exclude_layers = cast(list[str], self.quant_config.get("exclude", [])) + self.is_prequantized = is_prequantized self.packed_modules_mapping = self.quant_config["packed_modules_mapping"] + self._quantized_layers = set() + + @property + def quantized_layers(self) -> tuple[list[str], int]: + # Extract unique layer types (last part after ".") + layer_types = sorted( + set(name.split(".")[-1] for name in self._quantized_layers) + ) + return layer_types, len(self._quantized_layers) def get_linear_method(self) -> "QuarkLinearMethod": return QuarkLinearMethod(self) @@ -83,6 +111,7 @@ class QuarkConfig(QuantizationConfig): self, layer: torch.nn.Module, prefix: str ) -> Optional["QuantizeMethodBase"]: # Check if the layer is skipped for quantization. + if should_ignore_layer( prefix, ignore=self.exclude_layers, @@ -97,14 +126,17 @@ class QuarkConfig(QuantizationConfig): if isinstance(layer, LinearBase): scheme = self.get_linear_scheme(layer=layer, layer_name=prefix) layer.scheme = scheme + self._quantized_layers.add(prefix) return QuarkLinearMethod(self) if isinstance(layer, RadixAttention): + self._quantized_layers.add(prefix) return QuarkKVCacheMethod(self) from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE if isinstance(layer, FusedMoE): + self._quantized_layers.add(prefix) layer.scheme = self.get_moe_scheme(layer, prefix) return QuarkFusedMoEMethod(self) @@ -175,12 +207,72 @@ class QuarkConfig(QuantizationConfig): kv_cache_group=kv_cache_group, kv_cache_config=kv_cache_config, pack_method=pack_method, + is_prequantized=True, ) @classmethod def get_config_filenames(cls) -> list[str]: return [] + @staticmethod + def _create_online_mxfp4_config(model_type: str) -> dict[str, Any]: + """ + Create a synthetic quant_config for online MXFP4 quantization. + """ + # MOE gate/router is typically implemented as a ReplicatedLinear, and skipped for quantization for accuracy reasons. + # lm_head/embed_tokens is also skipped for accuracy reasons, normally not handled by `QuarkConfig` in any case, but adding them here for safety. + exclude = [ + "re:.*gate$", + "re:.*router", + "re:.*lm_head", + "re:.*embed_tokens", + ] + + if model_type == "qwen3_5_moe": + # Exclusion for accuracy adapted from + # https://huggingface.co/amd/Qwen3.5-397B-A17B-MXFP4/blob/main/config.json + exclude.extend( + [ + "re:.*n_proj_a", + "re:.*in_proj_b", + "re:.*in_proj_qkv", + "re:.*in_proj_z", + "re:.*o_proj", + "re:.*out_proj", + "re:.*qkv_proj", + "re:.*shared_expert", + ] + ) + + return { + "packed_modules_mapping": {}, + "exclude": exclude, + "global_quant_config": { + "weight": { + "dtype": "fp4", + "qscheme": "per_group", + "group_size": 32, + "is_dynamic": False, + "scale_format": "e8m0", + }, + "input_tensors": { + "dtype": "fp4", + "qscheme": "per_group", + "group_size": 32, + "is_dynamic": True, + "scale_format": "e8m0", + }, + "output_tensors": None, + "bias": None, + }, + "layer_quant_config": {}, + "layer_type_quant_config": {}, + "export": { + "kv_cache_group": [], + "pack_method": "reorder", + }, + } + def _check_scheme_supported(self, min_capability: int, error: bool = True) -> bool: capability_tuple = get_device_capability() @@ -337,7 +429,11 @@ class QuarkConfig(QuantizationConfig): input_config = cast(dict[str, Any], config.get("input_tensors")) if self._is_mx_fp4(weight_config, input_config): - return QuarkW4A4MXFP4(weight_config, input_config) + return QuarkW4A4MXFP4( + weight_config, + input_config, + is_checkpoint_mxfp4_serialized=self.is_prequantized, + ) if self._is_fp8_w8a8(weight_config, input_config): is_fp8_w8a8_supported = self._check_scheme_supported( QuarkW8A8Fp8.get_min_capability(), error=False @@ -383,7 +479,11 @@ class QuarkConfig(QuantizationConfig): input_config = layer_quant_config.get("input_tensors") if self._is_mx_fp4(weight_config, input_config): - return QuarkW4A4MXFp4MoE(weight_config, input_config) + return QuarkW4A4MXFp4MoE( + weight_config, + input_config, + is_checkpoint_mxfp4_serialized=self.is_prequantized, + ) elif self._is_fp8_w8a8(weight_config, input_config): return QuarkW8A8FP8MoE(weight_config, input_config) else: diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py index f4510def7..6bcf21477 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 +import logging from typing import Any, Callable, Optional import torch @@ -7,6 +8,7 @@ import torch from sglang.srt.layers.parameter import GroupQuantScaleParameter, PackedvLLMParameter from sglang.srt.layers.quantization.quark.schemes import QuarkLinearScheme from sglang.srt.utils import is_hip +from sglang.srt.utils.common import mxfp_supported _is_hip = is_hip() if _is_hip: @@ -19,6 +21,7 @@ if _is_hip: __all__ = ["QuarkW4A4MXFP4"] +logger = logging.getLogger(__name__) OCP_MX_BLOCK_SIZE = 32 @@ -26,19 +29,35 @@ OCP_MX_BLOCK_SIZE = 32 class QuarkW4A4MXFP4(QuarkLinearScheme): def __init__( - self, weight_quant_spec: dict[str, Any], input_quant_spec: dict[str, Any] + self, + weight_quant_spec: dict[str, Any], + input_quant_spec: dict[str, Any], + is_checkpoint_mxfp4_serialized: bool = True, ): self.out_dtype = torch.get_default_dtype() self.qscheme = "per_group" self.weight_quant_spec = weight_quant_spec self.input_quant_spec = input_quant_spec + self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized + + if not self.is_checkpoint_mxfp4_serialized: + if not mxfp_supported(): + raise NotImplementedError( + "Online MXFP4 quantization requires an AMD ROCm device with " + "FP4 hardware support (gfx95x, e.g. MI355x)." + ) + logger.info_once( + "Using online MXFP4 quantization from a higher precision checkpoint. Beware that this optimization may degrade prediction quality - please validate your model accuracy. More details at https://docs.sglang.io/advanced_features/quantization.html#online-quantization." + ) @classmethod def get_min_capability(cls) -> int: return 70 def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - return + if not self.is_checkpoint_mxfp4_serialized: + assert layer.weight.dtype == torch.uint8 + assert layer.weight_scale.dtype == torch.uint8 def create_weights( self, @@ -49,10 +68,19 @@ class QuarkW4A4MXFP4(QuarkLinearScheme): weight_loader: Callable, **kwargs, ): + self.input_size_per_partition = input_size_per_partition + output_size_per_partition = sum(output_partition_sizes) + self.output_size_per_partition = output_size_per_partition + layer.logical_widths = output_partition_sizes + original_weight_loader = weight_loader + if not self.is_checkpoint_mxfp4_serialized: + weight_loader = self.get_online_mxfp4_weight_loader(layer, weight_loader) + # WEIGHT + # Both serialized and online quantization use packed uint8 format weight = PackedvLLMParameter( data=torch.empty( output_size_per_partition, @@ -76,19 +104,50 @@ class QuarkW4A4MXFP4(QuarkLinearScheme): ), input_dim=1, output_dim=0, - weight_loader=weight_loader, + weight_loader=original_weight_loader, ) layer.register_parameter("weight_scale", weight_scale) + def get_online_mxfp4_weight_loader( + self, + layer, + original_weight_loader: Callable, + ) -> Callable: + """ + Wrap the original weight loader to perform online MXFP4 quantization. + """ + + def online_mxfp4_weight_loader( + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + shard_id: int | str | None = None, + ): + # Materialize on device the loaded weight. + loaded_weight = loaded_weight.to(param.device) + + # Quantize the loaded weight shard immediately. Since MXFP4 uses per-group quantization, there is no need to load all shards (e.g. q_proj, k_proj, v_proj) before doing online quantization. + qweight, weight_scale = dynamic_mxfp4_quant(loaded_weight) + + # Required e.g. for q_proj, k_proj, v_proj. + kwargs = {} + if shard_id is not None: + kwargs["loaded_shard_id"] = shard_id + + # Use the original weight loader to handle the loading logic + # (e.g. qkv sharding, etc.) + original_weight_loader(param, qweight, **kwargs) + + layer.weight_scale.weight_loader(layer.weight_scale, weight_scale, **kwargs) + + return online_mxfp4_weight_loader + def apply_weights( self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: - # This path does not have support for bias currently - assert bias is None, "bias is not supported" - + # Bias will be added after the GEMM if provided three_d = False fused_gemm_split_cat = False x_s = None @@ -152,6 +211,9 @@ class QuarkW4A4MXFP4(QuarkLinearScheme): else: gemm_afp4wfp4(x_q, layer.weight, x_s, layer.weight_scale, self.out_dtype, y) + if bias is not None: + y = y + bias + if fused_gemm_split_cat: return k, v elif three_d: diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py index fd3629cd7..c6890dfd5 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py @@ -16,6 +16,7 @@ from sglang.srt.utils import ( is_hip, set_weight_attrs, ) +from sglang.srt.utils.common import mxfp_supported if TYPE_CHECKING: from sglang.srt.layers.moe.token_dispatcher import ( @@ -35,14 +36,25 @@ if _use_aiter: from aiter.ops.shuffle import shuffle_weight from aiter.utility.fp4_utils import e8m0_shuffle +if _is_hip: + from aiter.ops.triton.quant import dynamic_mxfp4_quant +else: + dynamic_mxfp4_quant = None + OCP_MX_BLOCK_SIZE = 32 class QuarkW4A4MXFp4MoE(QuarkMoEScheme): - def __init__(self, weight_config: dict[str, Any], input_config: dict[str, Any]): + def __init__( + self, + weight_config: dict[str, Any], + input_config: dict[str, Any], + is_checkpoint_mxfp4_serialized: bool = True, + ): self.weight_quant = weight_config self.input_quant = input_config + self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized weight_qscheme = self.weight_quant.get("qscheme") input_qscheme = self.input_quant.get("qscheme") @@ -56,6 +68,18 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme): self.static_input_scales = not self.input_quant.get("is_dynamic") self.with_bias = False + if not self.is_checkpoint_mxfp4_serialized: + if not mxfp_supported(): + raise NotImplementedError( + "Online MXFP4 quantization for MoE layers requires an AMD ROCm " + "device with FP4 hardware support (gfx95x, e.g. MI355x)." + ) + logger.info_once( + "Using online MXFP4 quantization for MoE layers from a higher precision checkpoint. " + "Beware that this optimization may degrade prediction quality - please validate your model accuracy. " + "More details at https://docs.sglang.io/advanced_features/quantization.html#online-quantization." + ) + @classmethod def get_min_capability(cls) -> int: return 70 @@ -90,7 +114,16 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme): params_dtype = torch.uint8 - # WEIGHTS + original_weight_loader = extra_weight_attrs.get("weight_loader") + + if self.is_checkpoint_mxfp4_serialized: + weight_loader = original_weight_loader + else: + weight_loader = self.get_online_weight_loader(layer, original_weight_loader) + + extra_weight_attrs["weight_loader"] = weight_loader + + # WEIGHTS — always uint8 (packed mxfp4), always on device w13_weight = torch.nn.Parameter( torch.empty( num_experts, @@ -101,7 +134,6 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme): requires_grad=False, ) layer.register_parameter("w13_weight", w13_weight) - set_weight_attrs(w13_weight, extra_weight_attrs) w2_weight = torch.nn.Parameter( @@ -114,10 +146,11 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme): requires_grad=False, ) layer.register_parameter("w2_weight", w2_weight) - set_weight_attrs(w2_weight, extra_weight_attrs) # WEIGHT_SCALES + extra_weight_attrs["weight_loader"] = original_weight_loader + w13_weight_scale = torch.nn.Parameter( torch.ones( num_experts, @@ -140,26 +173,60 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme): ), requires_grad=False, ) - set_weight_attrs(w2_weight_scale, extra_weight_attrs) set_weight_attrs(w13_weight_scale, extra_weight_attrs) - + set_weight_attrs(w2_weight_scale, extra_weight_attrs) layer.register_parameter("w13_weight_scale", w13_weight_scale) layer.register_parameter("w2_weight_scale", w2_weight_scale) - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - float_dtype = torch.get_default_dtype() + def get_online_weight_loader(self, layer, original_weight_loader): + """ + Wrap the original weight loader to perform online MXFP4 quantization. + """ + def online_mxfp4_moe_weight_loader( + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: int, + ): + if dynamic_mxfp4_quant is None: + raise NotImplementedError( + "Online MXFP4 quantization for MoE is only supported on AMD GPUs." + ) + + # Materialize on device the loaded weight. + loaded_weight = loaded_weight.to(param.device) + + # Quantize the high-precision shard loaded_weight to MXFP4. + qweight, weight_scale = dynamic_mxfp4_quant(loaded_weight) + + original_weight_loader(param, qweight, weight_name, shard_id, expert_id) + + if "w13" in weight_name: + scale_param = layer.w13_weight_scale + scale_weight_name = "w13_weight_scale" + else: + # w2. + scale_param = layer.w2_weight_scale + scale_weight_name = "w2_weight_scale" + + scale_param.weight_loader( + scale_param, weight_scale, scale_weight_name, shard_id, expert_id + ) + + return online_mxfp4_moe_weight_loader + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Pre-shuffle weight scales s0, s1, _ = layer.w13_weight_scale.shape w13_weight_scale = layer.w13_weight_scale.view(s0 * s1, -1) w13_weight_scale = e8m0_shuffle(w13_weight_scale) - # layer.w13_weight_scale = torch.nn.Parameter(w13_weight_scale, requires_grad=False) layer.w13_weight_scale.data = w13_weight_scale.view(s0, s1, -1) s0, s1, _ = layer.w2_weight_scale.shape w2_weight_scale = layer.w2_weight_scale.view(s0 * s1, -1) w2_weight_scale = e8m0_shuffle(w2_weight_scale) - # layer.w2_weight_scale = torch.nn.Parameter(w2_weight_scale, requires_grad=False) layer.w2_weight_scale.data = w2_weight_scale.view(s0, s1, -1) # Pre-shuffle weight diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index c99c8f54d..12c83090b 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1417,6 +1417,20 @@ class ModelRunner(ModelRunnerKVCacheMixin): f"avail mem={after_avail_memory:.2f} GB, " f"mem usage={self.weight_load_mem_usage:.2f} GB." ) + + # TODO: Make sure all models have `quant_config` attribute, and all online quantization methods register which layers they actually quantize. + if ( + hasattr(self.model, "quant_config") + and hasattr(self.model.quant_config, "quantized_layers") + and self.server_args.quantization is not None + ): + layer_types, quantized_layers_count = ( + self.model.quant_config.quantized_layers + ) + logger.info( + f"Online {self.server_args.quantization} quantization: quantized {quantized_layers_count} layers of types: {layer_types}" + ) + if self.server_args.debug_tensor_dump_output_folder is not None: dump_folder = self.server_args.debug_tensor_dump_output_folder if self.spec_algorithm.is_eagle(): diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 2ee8d3e3b..32ff5bac2 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -37,12 +37,14 @@ import huggingface_hub import numpy as np import torch +from sglang.srt.constants import GIB_BYTES from sglang.srt.model_loader.remote_instance_weight_loader_utils import ( RemoteInstanceWeightLoaderBackend, get_remote_instance_transfer_engine_info_per_rank, register_memory_region, ) from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import get_available_gpu_memory # Try to import accelerate (optional dependency) try: @@ -82,6 +84,7 @@ from sglang.srt.model_loader.utils import ( get_model_architecture, set_default_torch_dtype, ) +from sglang.srt.utils.common import is_cuda_alike # Constants for memory management DEFAULT_GPU_MEMORY_FRACTION_FOR_CALIBRATION = ( @@ -746,8 +749,29 @@ class DefaultModelLoader(BaseModelLoader): @staticmethod def load_weights_and_postprocess(model, weights, target_device): + # Used in tests to verify memory savings when using online quantization. + if is_cuda_alike(): + peak_memory = torch.cuda.max_memory_allocated() + logger.debug( + "Peak GPU memory before loading weights: %s GiB", + f"{peak_memory / GIB_BYTES:.3f}", + ) + memory_start = get_available_gpu_memory( + target_device.type, gpu_id=torch.cuda.current_device() + ) + model.load_weights(weights) + # Used in tests to verify memory savings when using online quantization. + if is_cuda_alike(): + memory_end = get_available_gpu_memory( + target_device.type, gpu_id=torch.cuda.current_device() + ) + logger.debug( + "Memory increase during load_weights: %s GiB", + f"{memory_start - memory_end:.3f}", + ) + for _, module in model.named_modules(): quant_method = getattr(module, "quant_method", None) if quant_method is not None: diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 1690d69d7..e53a3193b 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -294,9 +294,15 @@ def get_quant_config( possible_config_filenames = quant_cls.get_config_filenames() # If the quantization config is not found, use the default config. + # TODO: standardize the handling of online quantization with custom handlenames (mxfp8, quark_mxfp4, etc.) if not possible_config_filenames: if model_config.quantization == "mxfp8": return Fp8Config(use_mxfp8=True, is_checkpoint_fp8_serialized=False) + if model_config.quantization == "quark_mxfp4": + return quant_cls( + online_scheme=model_config.quantization, + hf_config=model_config.hf_config, + ) return quant_cls() config_files = glob.glob(os.path.join(hf_folder, "*.json")) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index c67e2393b..c0abb6fce 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -125,32 +125,36 @@ LOAD_FORMAT_CHOICES = [ "runai_streamer", ] +# TODO: this list should likely contain only methods that support online quantization, or that support using custom quantization classes compatible with a given `quant_method` in config.json. +# Some of the choices here do NOT support online quantization. QUANTIZATION_CHOICES = [ "awq", - "fp8", - "mxfp8", + "fp8", # MOE + linear online quantization. + "mxfp8", # MOE + linear online quantization. "gptq", "marlin", "gptq_marlin", "awq_marlin", "bitsandbytes", "gguf", + # Modelopt has some online quantization support through ModelOptModelLoader. "modelopt", "modelopt_fp8", "modelopt_fp4", "modelopt_mixed", "petit_nvfp4", - "w8a8_int8", - "w8a8_fp8", - "moe_wna16", + "w8a8_int8", # mentioned in quantization.md documentation, supporting compressed-tensors quant_method. + "w8a8_fp8", # mentioned in quantization.md documentation, supporting compressed-tensors quant_method. + "moe_wna16", # custom loading logic for gptq/awq checkpoints (likely untested/unused) "qoq", "w4afp8", - "mxfp4", + "mxfp4", # MOE-only. "auto-round", "compressed-tensors", # for Ktransformers "modelslim", # for NPU "quark", # AMD Quark quantizer (FP8 / MXFP4 / Int4FP8 etc.) "quark_int4fp8_moe", + "quark_mxfp4", # Online MOE + linear quantization. # Apple Silicon MLX backend — on-the-fly quantization of fp16 weights at load # time via mlx.nn.quantize. Only takes effect when SGLANG_USE_MLX=1. "mlx_q4", # 4 bits, group_size=64 (mlx-community default) @@ -158,6 +162,7 @@ QUANTIZATION_CHOICES = [ "unquant", ] + SPECULATIVE_DRAFT_MODEL_QUANTIZATION_CHOICES = QUANTIZATION_CHOICES ATTENTION_BACKEND_CHOICES = [ diff --git a/test/registered/quant/test_quark_mxfp4.py b/test/registered/quant/test_quark_mxfp4.py new file mode 100644 index 000000000..637a3a423 --- /dev/null +++ b/test/registered/quant/test_quark_mxfp4.py @@ -0,0 +1,188 @@ +import io +import re +import unittest + +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci + +register_cuda_ci(est_time=103, stage="base-b", runner_config="1-gpu-small") +register_amd_ci(est_time=106, suite="stage-b-test-1-gpu-small-amd-mi35x") +import time +from types import SimpleNamespace + +import requests + +from sglang.srt.utils import kill_process_tree +from sglang.srt.utils.common import is_cuda_alike, mxfp_supported +from sglang.test.few_shot_gsm8k import run_eval +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + + +class TestOnlineQuantizationMemoryLoad(CustomTestCase): + runner_args = [] + + @classmethod + def setUpClass(cls): + if not mxfp_supported(): + raise unittest.SkipTest( + "online MXFP4 quantization requires an AMD ROCm device with " + "FP4 hardware support (gfx95x, e.g. MI355x)" + ) + cls.base_url = DEFAULT_URL_FOR_TEST + cls.stdout = io.StringIO() + cls.stderr = io.StringIO() + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--quantization", + "quark_mxfp4", + # `context-length` limitation required for Qwen MOE model + # TODO: Remove once https://github.com/sgl-project/sglang/pull/18255 and https://github.com/sgl-project/sglang/pull/18263 are merged. + "--context-length", + "3000", + "--tensor-parallel-size", + cls.tp if hasattr(cls, "tp") else "1", + "--log-level", + "debug", + *cls.runner_args, + ], + return_stdout_stderr=(cls.stdout, cls.stderr), + ) + + url = cls.base_url + "/health" + timeout = DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH + start_time = time.perf_counter() + while True: + try: + response = requests.get(url) + if response.status_code == 200: + print(f"Server {url} is ready") + break + except Exception: + pass + + if time.perf_counter() - start_time > timeout: + raise RuntimeError(f"Server {url} failed to start in {timeout}s") + time.sleep(1) + + # # Extract and display peak GPU memory from logs + combined_output = cls.stdout.getvalue() + cls.stderr.getvalue() + + peak_memory_before_load = cls._extract_peak_memory_before_load(combined_output) + if is_cuda_alike() and not peak_memory_before_load: + raise ValueError("Should have found peak memory") + cls.peak_memory_before_load = float(peak_memory_before_load) + + memory_increase_load_weights = cls._extract_memory_increase_load_weights( + combined_output + ) + if is_cuda_alike() and not memory_increase_load_weights: + raise ValueError("Should have found memory increase in load_weights") + cls.memory_increase_load_weights = float(memory_increase_load_weights) + + @classmethod + def _extract_peak_memory_before_load(cls, log_output): + """Extract peak GPU memory value from log output.""" + # Search for the log message pattern + pattern = r"Peak GPU memory before loading weights:\s+([\d.]+)\s+GiB" + match = re.search(pattern, log_output) + if match: + return match.group(1) + return None + + @classmethod + def _extract_memory_increase_load_weights(cls, log_output): + """Extract memory increase during load_weights call.""" + # Search for the log message pattern + pattern = r"Memory increase during load_weights:\s+([\d.]+)\s+GiB" + match = re.search(pattern, log_output) + if match: + return match.group(1) + return None + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + cls.stdout.close() + cls.stderr.close() + + def _test_peak_memory( + self, threshold, test_start: bool, add_peak_memory_before_load: bool + ): + """Helper method to test peak memory against a threshold.""" + if not is_cuda_alike(): + self.skipTest("not is_cuda_alike") + + # NOTE: We can not simply rely on peak memory after `load_weights` as functions used + # in-between (e.g. FP8->MXFP4 requantization) during weight loading may have a higher peak memory footprint + # than simply the allocated weights. + if add_peak_memory_before_load: + reference_gib = ( + self.memory_increase_load_weights + self.peak_memory_before_load + ) + else: + reference_gib = self.memory_increase_load_weights + + assert reference_gib < threshold + + if test_start: + # Weights initialized on meta device (not for dense BF16->MXFP4) + assert self.peak_memory_before_load < 5 + + def _test_gsm8k(self, accuracy_threshold): + """Helper method to test GSM8K accuracy against a threshold.""" + args = SimpleNamespace( + num_shots=8, + data_path=None, + num_questions=500, + max_new_tokens=512, + parallel=128, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreater(metrics["accuracy"], accuracy_threshold) + + +class TestOnlineQuantizationMemoryLoadDense(TestOnlineQuantizationMemoryLoad): + model = "Qwen/Qwen3-8B" + + def test_peak_memory(self): + # Original Qwen/Qwen3-8B BF16 model: 15.268 GiB + self._test_peak_memory( + threshold=6, test_start=False, add_peak_memory_before_load=True + ) + + def test_gsm8k(self): + # Original Qwen/Qwen3-8B reference accuracy: ~0.92 + self._test_gsm8k(accuracy_threshold=0.85) + + +class TestOnlineQuantizationMemoryLoadMOE(TestOnlineQuantizationMemoryLoad): + # Unfortunately, smaller models as Qwen/Qwen1.5-MoE-A2.7B or ibm-granite/granite-3.0-3b-a800m-base currently crash in AITER: + # - Qwen/Qwen1.5-MoE-A2.7B => K // 2 = 704 as intermediate size, not multiple of 128. + # - ibm-granite/granite-3.0-3b-a800m-base: dtype issue with fp16 in AITER MOE MLP activation + # so using a large model here. + model = "Qwen/Qwen3-30B-A3B-Instruct-2507" + # TODO: test TP>=2 with an other model (Qwen/Qwen3-30B-A3B-Instruct-2507 crashes in this case as 768/2 = 384, and 384/32 = 12 not divisible by BLOCK_SIZE_N=8. in fused_dynamic_mxfp4_quant_moe_sort. + + def test_peak_memory(self): + # Original Qwen/Qwen3-30B-A3B-Instruct-2507 BF16 model: 56.940 GiB + self._test_peak_memory( + threshold=17, test_start=False, add_peak_memory_before_load=True + ) + + def test_gsm8k(self): + # Original Qwen/Qwen3-30B-A3B-Instruct-2507 reference accuracy: 0.94 + self._test_gsm8k(accuracy_threshold=0.89) + + +if __name__ == "__main__": + unittest.main()