[AMD][MXFP4] Online MXFP4 quantization 1/N - dense and MOE models w. original BF16 weight (#18005)

Co-authored-by: Bowen Bao <bowenbao@amd.com>
Co-authored-by: Colin Zeng <Colin.Zeng@amd.com>

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)

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

python/sglang/srt/constants.py

@@ -10,3 +10,5 @@ GPU_MEMORY_ALL_TYPES = [
 ]
 
 HEALTH_CHECK_RID_PREFIX = "HEALTH_CHECK"
+
+GIB_BYTES = 1073741824  # 1024**3

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,

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:

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
+        # Bias will be added after the GEMM if provided
-        assert bias is None, "bias is not supported"
-
         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:

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:
+    def get_online_weight_loader(self, layer, original_weight_loader):
-        float_dtype = torch.get_default_dtype()
+        """
+        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

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():

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:

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

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",
+    "fp8",  # MOE + linear online quantization.
-    "mxfp8",
+    "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_int8",  # mentioned in quantization.md documentation, supporting compressed-tensors quant_method.
-    "w8a8_fp8",
+    "w8a8_fp8",  # mentioned in quantization.md documentation, supporting compressed-tensors quant_method.
-    "moe_wna16",
+    "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 = [

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()