chore: import upstream snapshot with attribution

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

@@ -0,0 +1,705 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import fnmatch
+import logging
+from typing import TYPE_CHECKING, Any, List, Optional, cast
+
+import torch
+
+from sglang.srt.layers.linear import LinearBase
+from sglang.srt.layers.moe import MoeRunnerConfig
+from sglang.srt.layers.quantization.base_config import (  # noqa: E501
+    FusedMoEMethodBase,
+    LinearMethodBase,
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod
+from sglang.srt.layers.quantization.quark.schemes import (
+    QuarkLinearScheme,
+    QuarkMoEScheme,
+    QuarkW4A4MXFP4,
+    QuarkW4A4MXFp4MoE,
+    QuarkW4A8MXFp4MoE,
+    QuarkW8A8Fp8,
+    QuarkW8A8FP8MoE,
+)
+from sglang.srt.layers.quantization.quark.utils import deep_compare, should_ignore_layer
+from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
+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"]
+
+logger = logging.getLogger(__name__)
+
+_MOE_SHARED_EXPERT_QUANT_LAYER0_BASES: tuple[str, ...] = (
+    "model.layers.0",
+    "model.language_model.layers.0",
+)
+
+_SHARED_EXPERT_BODY_PROJ_SUFFIXES: tuple[str, ...] = (
+    "gate_proj",
+    "up_proj",
+    "gate_up_proj",
+    "down_proj",
+)
+
+
+class QuarkConfig(QuantizationConfig):
+
+    def __init__(
+        self,
+        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)
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+        return [torch.float16, torch.bfloat16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 70
+
+    def get_name(self) -> str:
+        return "quark"
+
+    def apply_weight_name_mapper(self, hf_to_sglang_mapper):
+        mapped = hf_to_sglang_mapper.apply_list(self.exclude_layers)
+        expanded = []
+        for name in mapped:
+            expanded.append(name)
+            if name.startswith("language_model."):
+                expanded.append(name.removeprefix("language_model."))
+        self.exclude_layers = list(dict.fromkeys(expanded))
+
+    def get_quant_method(
+        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,
+            fused_mapping=self.packed_modules_mapping,
+        ):
+            if isinstance(layer, LinearBase):
+                return UnquantizedLinearMethod()
+            elif isinstance(layer, RadixAttention):
+                return QuarkKVCacheMethod(self)
+            return None
+
+        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)
+
+        return None
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> "QuarkConfig":
+        export_config = config.get("export")
+        if export_config is None:
+            raise ValueError(
+                "The export key should be included in "
+                "the configurations of Quark quantized model"
+            )
+
+        kv_cache_group = cast(list[str], export_config.get("kv_cache_group"))
+        pack_method = cast(str, export_config.get("pack_method"))
+
+        # In the export model of quark, the quantization configuration
+        # of kv_cache is stored in layer_quant_config. First, it is
+        # judged whether kv_cache_group exists, and then it is judged
+        # whether layer_quant_config has a quantization configuration
+        # that matches kv_cache.
+        if len(kv_cache_group) == 0:
+            kv_cache_config = None
+        else:
+            kv_cache_set = set(kv_cache_group)
+            layer_quant_config = cast(dict[str, Any], config.get("layer_quant_config"))
+            layer_quant_names = list(layer_quant_config.keys())
+            layer_quant_set = set(layer_quant_names)
+
+            if not kv_cache_set.issubset(layer_quant_set):
+                raise ValueError(
+                    "The Quark quantized model has the "
+                    "kv_cache_group parameter setting, "
+                    "but no kv_cache quantization settings "
+                    "were found in the quantization "
+                    "configuration."
+                )
+
+            q_configs = [
+                cast(dict[str, Any], layer_quant_config.get(name))
+                for name in kv_cache_group
+            ]
+            if not all(deep_compare(q_config, q_configs[0]) for q_config in q_configs):
+                raise ValueError(
+                    "The quantization method used for kv_cache should "
+                    "be the same, but the quantization method for the "
+                    "kv_cache layer in the config is different."
+                )
+            kv_cache_config = q_configs[0].get("output_tensors")
+            if kv_cache_config is None:
+                raise ValueError("The kv_cache quantization configuration is empty.")
+
+            # Since we have already set kv_cache quantization configurations,
+            # we will remove the quantization configuration for the
+            # output_tensors corresponding to the kv_cache layer.
+            for q_config in q_configs:
+                q_config["output_tensors"] = None
+
+            # In case q_proj output is also quantized, remove the configuration
+            # to keep qkv consistency.
+            q_proj_q_config = cast(dict[str, Any], layer_quant_config.get("*q_proj"))
+            if q_proj_q_config is not None:
+                q_proj_q_config["output_tensors"] = None
+
+        return cls(
+            quant_config=config,
+            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()
+
+        if capability_tuple is not None:
+            assert 0 <= capability_tuple[1] < 10
+            capability = capability_tuple[0] * 10 + capability_tuple[1]
+
+            supported = capability >= min_capability
+            if error and not supported:
+                # Pass a single joined message; RuntimeError stringifies
+                # multiple positional args as a tuple repr.
+                raise RuntimeError(
+                    "Quantization scheme is not supported for "
+                    f"the current GPU. Min capability: {min_capability}. "
+                    f"Current capability: {capability}."
+                )
+            return supported
+        else:
+            return False
+
+    def _is_fp8_w8a8(
+        self,
+        weight_quant: Optional[dict[str, Any]],
+        input_quant: Optional[dict[str, Any]],
+    ) -> bool:
+        # Confirm weights and input quantized.
+        if weight_quant is None or input_quant is None:
+            return False
+
+        # Confirm weight scheme is supported
+        is_fp8_dtype = (
+            weight_quant.get("dtype") == "fp8_e4m3"
+            and input_quant.get("dtype") == "fp8_e4m3"
+        )
+        is_static_weight = not weight_quant.get("is_dynamic")
+        is_per_tensor_or_channel_weight = weight_quant.get("qscheme") in [
+            "per_tensor",
+            "per_channel",
+        ]
+
+        if not (is_fp8_dtype and is_static_weight and is_per_tensor_or_channel_weight):
+            return False
+
+        # Dynamic quantization is always supported if weights supported.
+        if input_quant.get("is_dynamic"):
+            return True
+
+        # Confirm activation scheme is supported.
+        is_per_tensor_activation = input_quant.get("qscheme") == "per_tensor"
+        return is_per_tensor_activation
+
+    def _is_mx_fp4(
+        self,
+        weight_quant: Optional[dict[str, Any]],
+        input_quant: Optional[dict[str, Any]],
+    ) -> bool:
+        # Confirm weights and input quantized.
+        if weight_quant is None or input_quant is None:
+            logger.debug(
+                "Quark model is not in MX-FP4 format: "
+                "weight_quant or input_quant not set"
+            )
+            return False
+
+        # Input and weight dtype needs to be fp4.
+        if weight_quant.get("dtype") != "fp4" or input_quant.get("dtype") != "fp4":
+            logger.debug("Quark model is not in MX-FP4 format: dtype not fp4")
+            return False
+
+        # Input and weight qscheme needs to be per group.
+        if (
+            weight_quant.get("qscheme") != "per_group"
+            or input_quant.get("qscheme") != "per_group"
+        ):
+            logger.debug("Quark model is not in MX-FP4 format: not per_group")
+            return False
+
+        # Input and weight group size needs to be 32.
+        if weight_quant.get("group_size") != 32 or input_quant.get("group_size") != 32:
+            logger.debug("Quark model is not in MX-FP4 format: not group_size=32")
+            return False
+
+        # Weights need to use static quantization.
+        if weight_quant.get("is_dynamic") is True:
+            logger.debug("Quark model is not in MX-FP4 format: not weight static")
+            return False
+
+        # Activations need to use dynamic quantization.
+        if input_quant.get("is_dynamic") is False:
+            logger.debug("Quark model is not in MX-FP4 format: not activation dynamic")
+            return False
+
+        # Activations and weight scales need to be in e8m0 format.
+        if (
+            weight_quant.get("scale_format") != "e8m0"
+            or input_quant.get("scale_format") != "e8m0"
+        ):
+            logger.debug("Quark model is not in MX-FP4 format: not scale_format e8m0")
+            return False
+
+        return True
+
+    def _is_mx_w4a8(
+        self,
+        weight_quant: Optional[dict[str, Any]],
+        input_quant: Optional[dict[str, Any]],
+    ) -> bool:
+        if weight_quant is None or input_quant is None:
+            return False
+
+        is_mx_fp4_weight = (
+            weight_quant.get("dtype") == "fp4"
+            and weight_quant.get("qscheme") == "per_group"
+            and weight_quant.get("group_size") == 32
+            and not weight_quant.get("is_dynamic")
+            and weight_quant.get("scale_format") == "e8m0"
+        )
+        is_static_fp8_activation = (
+            input_quant.get("dtype") in ("fp8_e4m3", "fp8_e4m3fn")
+            and input_quant.get("qscheme") == "per_tensor"
+            and not input_quant.get("is_dynamic")
+        )
+        return is_mx_fp4_weight and is_static_fp8_activation
+
+    def _find_matched_config(
+        self, layer_name: str, module: torch.nn.Module
+    ) -> dict[str, Any]:
+
+        proj_name = layer_name.split(".")[-1]
+        if proj_name in self.packed_modules_mapping:
+            shard_proj_names = self.packed_modules_mapping[proj_name]
+
+            # Convert fused_name --> [shard_names]
+            shard_names = [
+                layer_name.replace(proj_name, shard_proj_name)
+                for shard_proj_name in shard_proj_names
+            ]
+            shard_configs = [
+                self._find_matched_config(shard_name, module)
+                for shard_name in shard_names
+            ]
+            if not all(
+                deep_compare(q_config, shard_configs[0]) for q_config in shard_configs
+            ):
+                raise ValueError(
+                    f"Found a different quantization configuration for "
+                    f"{shard_proj_names} in {layer_name}. SGLang "
+                    "requires all to use the same scheme."
+                )
+            return shard_configs[0]
+        else:
+            layer_quant_config = cast(
+                dict[str, Any], self.quant_config.get("layer_quant_config")
+            )
+            for name_pattern in layer_quant_config:
+                if fnmatch.fnmatch(layer_name, name_pattern):
+                    return layer_quant_config[name_pattern]
+
+            layer_type = type(module).__name__
+            layer_type_quant_config = cast(
+                dict[str, Any], self.quant_config.get("layer_type_quant_config")
+            )
+            if layer_type in layer_type_quant_config:
+                return layer_type_quant_config[layer_type]
+
+            global_quant_config = cast(
+                dict[str, Any], self.quant_config.get("global_quant_config")
+            )
+            return global_quant_config
+
+    def _get_scheme_from_config(self, config: dict[str, Any]) -> "QuarkLinearScheme":
+        if config.get("output_tensors") or config.get("bias"):
+            raise NotImplementedError(
+                "Currently, Quark models with output_tensors "
+                "and bias quantized are not supported"
+            )
+        weight_config = cast(dict[str, Any], config.get("weight"))
+        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,
+                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
+            )
+            if is_fp8_w8a8_supported:
+                return QuarkW8A8Fp8(weight_config, input_config)
+
+        raise NotImplementedError(
+            "No quark compatible scheme was found. "
+            f"Weight config: {weight_config}, "
+            f"Input config: {input_config}"
+        )
+
+    def get_linear_scheme(
+        self, layer: torch.nn.Module, layer_name: str
+    ) -> "QuarkLinearScheme":
+
+        layer_quant_config = self._find_matched_config(layer_name, layer)
+
+        # Find the quant_scheme
+        scheme = self._get_scheme_from_config(layer_quant_config)
+
+        # Raise error if device does not support the scheme
+        # (e.g. fp8 needs ada lovelace)
+        self._check_scheme_supported(scheme.get_min_capability())
+
+        return scheme
+
+    def get_moe_scheme(
+        self,
+        module: torch.nn.Module,
+        layer_name: str,
+    ) -> "QuarkMoEScheme":
+        layer_quant_config = self._find_matched_config(layer_name, module)
+
+        if layer_quant_config.get("output_tensors") or layer_quant_config.get("bias"):
+            raise NotImplementedError(
+                "Currently, Quark models with "
+                "output_tensors and bias "
+                "quantized are not supported"
+            )
+        weight_config = layer_quant_config.get("weight")
+        input_config = layer_quant_config.get("input_tensors")
+
+        if self._is_mx_fp4(weight_config, input_config):
+            return QuarkW4A4MXFp4MoE(
+                weight_config,
+                input_config,
+                is_checkpoint_mxfp4_serialized=self.is_prequantized,
+            )
+        elif self._is_mx_w4a8(weight_config, input_config):
+            logger.info_once("Using Quark MXFP4-W/FP8-A MoE scheme")
+            return QuarkW4A8MXFp4MoE(weight_config, input_config)
+        elif self._is_fp8_w8a8(weight_config, input_config):
+            return QuarkW8A8FP8MoE(weight_config, input_config)
+        else:
+            raise RuntimeError("Unsupported FusedMoe scheme")
+
+    def get_scaled_act_names(self) -> List[str]:
+        return []
+
+    def can_fuse_shared_expert(self) -> bool:
+        # Shared-expert body excluded from quant; the gate must not veto fusion.
+        if any(
+            "shared_expert" in layer
+            and "shared_expert_gate" not in layer
+            and not layer.startswith("mtp.")
+            for layer in self.exclude_layers
+        ):
+            return False
+
+        # No per-layer config -> uniform spec, nothing to compare.
+        layer_quant_config = self.quant_config.get("layer_quant_config") or {}
+        if not layer_quant_config:
+            return True
+
+        # Compare routed vs shared specs at layer 0 (stub module needed by
+        # _find_matched_config; an unmatched name -> ValueError -> cannot fuse).
+        lookup_stub = torch.nn.Module()
+        try:
+            for base in _MOE_SHARED_EXPERT_QUANT_LAYER0_BASES:
+                moe_name = f"{base}.mlp.experts"
+                moe_cfg = self._find_matched_config(moe_name, lookup_stub)
+                for suffix in _SHARED_EXPERT_BODY_PROJ_SUFFIXES:
+                    shared_name = f"{base}.mlp.shared_expert.{suffix}"
+                    shared_cfg = self._find_matched_config(shared_name, lookup_stub)
+                    if not deep_compare(moe_cfg, shared_cfg):
+                        return False
+        except ValueError:
+            return False
+
+        return True
+
+
+class QuarkLinearMethod(LinearMethodBase):
+
+    def __init__(self, quantization_config: QuarkConfig):
+        self.quantization_config = quantization_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer.scheme.process_weights_after_loading(layer)
+
+    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,
+    ):
+        """
+        Use the QuarkLinearScheme associated with the layer to create
+        the necessary parameters for the layer. See LinearMethodBase for param
+        details
+        """
+        weight_loader = extra_weight_attrs.get("weight_loader")
+        layer.scheme.create_weights(
+            layer=layer,
+            input_size=input_size,
+            input_size_per_partition=input_size_per_partition,
+            output_partition_sizes=output_partition_sizes,
+            output_size=output_size,
+            params_dtype=params_dtype,
+            weight_loader=weight_loader,
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ):
+        """
+        Use the output of create_weights and the QuarkLinearScheme
+        associated with the layer to apply the forward pass with the
+        layer input.  See LinearMethodBase for param details
+
+        """
+        scheme = layer.scheme
+        if scheme is None:
+            raise ValueError("A scheme must be defined for each layer")
+        return scheme.apply_weights(layer, x, bias=bias)
+
+
+class QuarkFusedMoEMethod(FusedMoEMethodBase):
+
+    def __init__(self, quantization_config: QuarkConfig):
+        self.quantization_config = quantization_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer.scheme.process_weights_after_loading(layer)
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        num_experts: int,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        """
+        Use the QuarkMoEScheme associated with the layer to create
+        the necessary parameters for the layer. See FusedMoEMethodBase for param
+        details
+        """
+        layer.scheme.create_weights(
+            layer=layer,
+            num_experts=num_experts,
+            hidden_size=hidden_size,
+            intermediate_size_per_partition=intermediate_size_per_partition,
+            params_dtype=params_dtype,
+            **extra_weight_attrs,
+        )
+
+    def create_moe_runner(
+        self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
+    ):
+        layer.scheme.create_moe_runner(layer, moe_runner_config)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: "StandardDispatchOutput",
+    ):
+        """
+        Use the output of create_weights and the QuarkMoEScheme
+        associated with the layer to apply the forward pass with the
+        fused MoE layer. See FusedMoEMethodBase for param details
+
+        """
+        scheme = layer.scheme
+        if scheme is None:
+            raise ValueError("A scheme must be defined for each layer")
+        return scheme.apply_weights(layer, dispatch_output)
+
+
+class QuarkKVCacheMethod(BaseKVCacheMethod):
+    """
+    Supports loading kv-cache scaling factors from quark checkpoints.
+    """
+
+    def __init__(self, quant_config: QuarkConfig):
+        self.validate_kv_cache_config(quant_config.kv_cache_config)
+        super().__init__(quant_config)
+
+    @staticmethod
+    def validate_kv_cache_config(kv_cache_config: Optional[dict[str, Any]]):
+        """
+        Validator for the kv cache configuration. Useful for controlling the
+        kv cache quantization schemes, that are being supported in vLLM
+        :param kv_cache_config: the quark kv cache scheme
+        """
+        if kv_cache_config is None:
+            return
+
+        dtype = kv_cache_config.get("dtype")
+        if dtype != "fp8_e4m3":
+            raise NotImplementedError(
+                "Currently supported kv cache quantization is "
+                f"dtype=fp8_e4m3, however received {dtype}"
+            )
+
+        qscheme = kv_cache_config.get("qscheme")
+        if qscheme != "per_tensor":
+            raise NotImplementedError(
+                "Only support per-tensor scaling factor "
+                "for quark KV cache. "
+                f"Expected qscheme: per_tensor, found qscheme: {qscheme}"
+            )

python/sglang/srt/layers/quantization/quark/schemes/__init__.py

@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from .quark_scheme import QuarkLinearScheme, QuarkMoEScheme
+from .quark_w4a4_mxfp4 import QuarkW4A4MXFP4
+from .quark_w4a4_mxfp4_moe import QuarkW4A4MXFp4MoE
+from .quark_w4a8_mxfp4_moe import QuarkW4A8MXFp4MoE
+from .quark_w8a8_fp8 import QuarkW8A8Fp8
+from .quark_w8a8_fp8_moe import QuarkW8A8FP8MoE
+
+__all__ = [
+    "QuarkLinearScheme",
+    "QuarkMoEScheme",
+    "QuarkW4A4MXFP4",
+    "QuarkW8A8Fp8",
+    "QuarkW4A4MXFp4MoE",
+    "QuarkW4A8MXFp4MoE",
+    "QuarkW8A8FP8MoE",
+]

python/sglang/srt/layers/quantization/quark/schemes/quark_scheme.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.layers.moe import MoeRunnerConfig
+from sglang.srt.layers.quantization.base_scheme import BaseLinearScheme, BaseMoEScheme
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
+
+__all__ = ["QuarkLinearScheme", "QuarkMoEScheme"]
+
+
+class QuarkLinearScheme(BaseLinearScheme):
+    """
+    Abstract class used to describe the weight creation and forward pass
+    of different quantization schemes supported by Quark.
+    """
+
+    @classmethod
+    @abstractmethod
+    def get_min_capability(cls) -> int:
+        """
+        Get minimum device capability.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def create_weights(self, *args, **kwargs):
+        """
+        Weight creation for the particular scheme. Inputs to this function
+
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        """
+        Called after weight loading is complete for any cleanup that
+        needs to occur.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def apply_weights(
+        self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor]
+    ):
+        """
+        Run the forward pass for the particular scheme. This is where
+        scheme-specific dequant/quant steps/kernels should be applied.
+
+        :param layer: torch.nn.Module with the registered weights and
+            other parameters relevant to the particular scheme.
+        :param x: input to the layer
+        :param bias: bias parameter
+
+        """
+        raise NotImplementedError
+
+
+class QuarkMoEScheme(BaseMoEScheme):
+    """
+    Abstract class used to describe the weight creation and forward pass
+    of different quantization schemes supported by Quark.
+    """
+
+    @classmethod
+    @abstractmethod
+    def get_min_capability(cls) -> int:
+        """
+        Get minimum device capability.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def create_weights(self, *args, **kwargs):
+        """
+        Weight creation for the particular scheme. Inputs to this function
+
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def create_moe_runner(
+        self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
+    ):
+        raise NotImplementedError
+
+    @abstractmethod
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        """
+        Called after weight loading is complete for any cleanup that
+        needs to occur.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: "StandardDispatchOutput",
+    ):
+        """
+        Run the forward pass for the particular scheme. This is where
+        scheme-specific dequant/quant steps/kernels should be applied.
+
+        :param layer: torch.nn.Module with the registered weights and
+            other parameters relevant to the particular scheme.
+        :param x: input to the layer
+        :param bias: bias parameter
+
+        """
+        raise NotImplementedError

python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py

@@ -0,0 +1,351 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+from typing import Any, Callable, Optional
+
+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 direct_register_custom_op, mxfp_supported
+
+_is_hip = is_hip()
+if _is_hip:
+    from aiter.ops.triton.gemm.fused.fused_gemm_afp4wfp4_split_cat import (
+        fused_gemm_afp4wfp4_split_cat as _fused_gemm_afp4wfp4_split_cat_orig,
+    )
+    from aiter.ops.triton.gemm_afp4wfp4 import gemm_afp4wfp4 as _gemm_afp4wfp4_orig
+    from aiter.ops.triton.gemm_afp4wfp4_pre_quant_atomic import (
+        gemm_afp4wfp4_pre_quant as _gemm_afp4wfp4_pre_quant_orig,
+    )
+    from aiter.ops.triton.quant import dynamic_mxfp4_quant as _dynamic_mxfp4_quant_orig
+
+    def _aiter_gemm_afp4wfp4(
+        x: torch.Tensor,
+        w: torch.Tensor,
+        x_scales: torch.Tensor,
+        w_scales: torch.Tensor,
+        y: torch.Tensor,
+    ) -> None:
+        _gemm_afp4wfp4_orig(x, w, x_scales, w_scales, y.dtype, y)
+
+    def _aiter_gemm_afp4wfp4_fake(
+        x: torch.Tensor,
+        w: torch.Tensor,
+        x_scales: torch.Tensor,
+        w_scales: torch.Tensor,
+        y: torch.Tensor,
+    ) -> None:
+        return None
+
+    direct_register_custom_op(
+        op_name="aiter_gemm_afp4wfp4",
+        op_func=_aiter_gemm_afp4wfp4,
+        mutates_args=["y"],
+        fake_impl=_aiter_gemm_afp4wfp4_fake,
+    )
+
+    def gemm_afp4wfp4(x, w, x_scales, w_scales, dtype, y):
+        torch.ops.sglang.aiter_gemm_afp4wfp4(x, w, x_scales, w_scales, y)
+
+    def _aiter_gemm_afp4wfp4_pre_quant(
+        x: torch.Tensor,
+        w: torch.Tensor,
+        w_scales: torch.Tensor,
+        y: torch.Tensor,
+    ) -> None:
+        _gemm_afp4wfp4_pre_quant_orig(x, w, w_scales, y.dtype, y)
+
+    def _aiter_gemm_afp4wfp4_pre_quant_fake(
+        x: torch.Tensor,
+        w: torch.Tensor,
+        w_scales: torch.Tensor,
+        y: torch.Tensor,
+    ) -> None:
+        return None
+
+    direct_register_custom_op(
+        op_name="aiter_gemm_afp4wfp4_pre_quant",
+        op_func=_aiter_gemm_afp4wfp4_pre_quant,
+        mutates_args=["y"],
+        fake_impl=_aiter_gemm_afp4wfp4_pre_quant_fake,
+    )
+
+    def gemm_afp4wfp4_pre_quant(x, w, w_scales, dtype, y):
+        torch.ops.sglang.aiter_gemm_afp4wfp4_pre_quant(x, w, w_scales, y)
+
+    def _aiter_dynamic_mxfp4_quant(
+        x: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return _dynamic_mxfp4_quant_orig(x)
+
+    def _aiter_dynamic_mxfp4_quant_fake(
+        x: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        M, N = x.shape
+        x_fp4 = torch.empty((M, N // 2), dtype=torch.uint8, device=x.device)
+        blockscale = torch.empty(
+            (M, (N + 31) // 32), dtype=torch.uint8, device=x.device
+        )
+        return x_fp4, blockscale
+
+    direct_register_custom_op(
+        op_name="aiter_dynamic_mxfp4_quant",
+        op_func=_aiter_dynamic_mxfp4_quant,
+        mutates_args=[],
+        fake_impl=_aiter_dynamic_mxfp4_quant_fake,
+    )
+
+    def dynamic_mxfp4_quant(x):
+        return torch.ops.sglang.aiter_dynamic_mxfp4_quant(x)
+
+    def _aiter_fused_gemm_split_cat(
+        x: torch.Tensor,
+        w: torch.Tensor,
+        y: torch.Tensor,
+        x_scale: torch.Tensor,
+        w_scale: torch.Tensor,
+        S1: int,
+        S2: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return _fused_gemm_afp4wfp4_split_cat_orig(
+            x=x,
+            w=w,
+            y=y,
+            x_scale=x_scale,
+            w_scale=w_scale,
+            S1=S1,
+            S2=S2,
+            dtype=y.dtype,
+        )
+
+    def _aiter_fused_gemm_split_cat_fake(
+        x: torch.Tensor,
+        w: torch.Tensor,
+        y: torch.Tensor,
+        x_scale: torch.Tensor,
+        w_scale: torch.Tensor,
+        S1: int,
+        S2: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        M = x.shape[0]
+        D = y.shape[1]
+        S3 = y.shape[2]
+        c1 = torch.empty((M, D, S1 + S3), dtype=y.dtype, device=x.device)
+        c2 = torch.empty((M, D, S2), dtype=y.dtype, device=x.device)
+        return c1, c2
+
+    direct_register_custom_op(
+        op_name="aiter_fused_gemm_split_cat",
+        op_func=_aiter_fused_gemm_split_cat,
+        mutates_args=[],
+        fake_impl=_aiter_fused_gemm_split_cat_fake,
+    )
+
+    def fused_gemm_afp4wfp4_split_cat(x, w, y, x_scale, w_scale, S1, S2, dtype):
+        return torch.ops.sglang.aiter_fused_gemm_split_cat(
+            x, w, y, x_scale, w_scale, S1, S2
+        )
+
+
+__all__ = ["QuarkW4A4MXFP4"]
+logger = logging.getLogger(__name__)
+
+OCP_MX_BLOCK_SIZE = 32
+
+
+class QuarkW4A4MXFP4(QuarkLinearScheme):
+
+    def __init__(
+        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:
+        if not self.is_checkpoint_mxfp4_serialized:
+            assert layer.weight.dtype == torch.uint8
+            assert layer.weight_scale.dtype == torch.uint8
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        output_partition_sizes: list[int],
+        input_size_per_partition: int,
+        params_dtype: torch.dtype,
+        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,
+                input_size_per_partition // 2,
+                dtype=torch.uint8,
+            ),
+            input_dim=1,
+            output_dim=0,
+            packed_dim=1,
+            packed_factor=2,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+        # WEIGHT SCALE
+        weight_scale = GroupQuantScaleParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition // OCP_MX_BLOCK_SIZE,
+                dtype=torch.uint8,
+            ),
+            input_dim=1,
+            output_dim=0,
+            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:
+        # Bias will be added after the GEMM if provided
+        three_d = False
+        fused_gemm_split_cat = False
+        x_s = None
+        y = None
+
+        if isinstance(x, tuple):
+            assert len(x) in [
+                2,
+                3,
+                5,
+            ], "For tuple input, only (x, x_s), (x, x_s, y), or (x, y, S1, S2, out_dtype) formats are accepted"
+            if len(x) == 2:
+                x, x_s = x
+            elif len(x) == 3:
+                x, x_s, y = x
+            elif len(x) == 5:
+                x, y, S1, S2, out_dtype = x
+                fused_gemm_split_cat = True
+
+        use_fused_quant_gemm = (
+            not fused_gemm_split_cat
+            and x_s is None
+            and y is not None
+            and layer.weight.shape[0] == y.shape[1]
+        )
+
+        if x.dim() == 3:
+            three_d = True
+            x = x.view(-1, x.shape[-1])
+            output_shape = [*x.shape[:-1], layer.weight.shape[0]]
+
+        # use_fused_quant_gemm = true, x_q is a bf16/fp16 num
+        # x_s is not None = true, x_q is uint8 num
+        if use_fused_quant_gemm or x_s is not None:
+            x_q = x
+        else:
+            x_q, x_s = dynamic_mxfp4_quant(x)
+
+        if y is None:
+            y = torch.empty(
+                x_q.shape[0],
+                layer.weight.shape[0],
+                device=x_q.device,
+                dtype=self.out_dtype,
+            )
+
+        if use_fused_quant_gemm:
+            gemm_afp4wfp4_pre_quant(x_q, layer.weight, layer.weight_scale, y.dtype, y)
+            y = y.to(x.dtype)
+        elif fused_gemm_split_cat:
+            k, v = fused_gemm_afp4wfp4_split_cat(
+                x=x_q,
+                w=layer.weight,
+                y=y,
+                x_scale=x_s,
+                w_scale=layer.weight_scale,
+                S1=S1,
+                S2=S2,
+                dtype=out_dtype,
+            )
+        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:
+            return y.view(*output_shape)
+        else:
+            return y

python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py

@@ -0,0 +1,295 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
+from sglang.srt.layers.moe.utils import get_moe_weight_sizes
+from sglang.srt.layers.quantization.quark.schemes import QuarkMoEScheme
+from sglang.srt.utils import (
+    get_bool_env_var,
+    is_gfx95_supported,
+    is_hip,
+    set_weight_attrs,
+)
+from sglang.srt.utils.common import mxfp_supported
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.token_dispatcher import (
+        CombineInput,
+        StandardDispatchOutput,
+    )
+
+logger = logging.getLogger(__name__)
+
+_is_shuffle_moe_mxfp4 = is_gfx95_supported()
+
+__all__ = ["QuarkW4A4MXFp4MoE"]
+
+_is_hip = is_hip()
+_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
+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],
+        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")
+        if not (weight_qscheme == "per_group" and input_qscheme == "per_group"):
+            raise ValueError(
+                "For MX(FP4) Fused MoE layers, only per-group scales "
+                "for weights and activations are supported. Found "
+                f"{weight_qscheme}, {input_qscheme}"
+            )  # noqa E501
+
+        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
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        num_experts: int,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+
+        from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
+
+        w13_up_dim, w2_down_dim, weight_padded = get_moe_weight_sizes(
+            intermediate_size_per_partition,
+            is_aiter_moe=_use_aiter,
+            is_concat=True,
+            is_packed=True,
+        )
+
+        # Add the quantization method used (per tensor/grouped/channel)
+        # to ensure the weight scales are loaded in properly
+        extra_weight_attrs.update(
+            {
+                "quant_method": FusedMoeWeightScaleSupported.BLOCK.value,
+                "weight_padded": weight_padded,
+            },
+        )
+
+        params_dtype = torch.uint8
+
+        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,
+                w13_up_dim,
+                hidden_size // 2,
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight", w13_weight)
+        set_weight_attrs(w13_weight, extra_weight_attrs)
+
+        w2_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                hidden_size,
+                w2_down_dim,
+                dtype=params_dtype,
+            ),
+            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,
+                w13_up_dim,
+                hidden_size // OCP_MX_BLOCK_SIZE,
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+
+        # 1. w2 scale is floor division of inter_dim by blockscale.
+        # 2. w2 scale needs to scale up just as w2.
+        # We combine 1. and 2. to keep the integer precision.
+        w2_weight_scale = torch.nn.Parameter(
+            torch.ones(
+                num_experts,
+                hidden_size,
+                (w2_down_dim * 2) // OCP_MX_BLOCK_SIZE,
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+        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 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.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.data = w2_weight_scale.view(s0, s1, -1)
+
+        # Pre-shuffle weight
+        if _is_shuffle_moe_mxfp4:
+            layer.w13_weight.data = shuffle_weight(
+                layer.w13_weight.contiguous(), (16, 16)
+            )
+            layer.w2_weight.data = shuffle_weight(
+                layer.w2_weight.contiguous(), (16, 16)
+            )
+            layer.w13_weight.is_shuffled = True
+            layer.w2_weight.is_shuffled = True
+
+        if hasattr(layer, "dispatcher"):
+            # Weights are stored as torch.uint8 but semantically MXFP4
+            layer.dispatcher.set_quant_config({"weight_dtype": torch.float4_e2m1fn_x2})
+
+    def create_moe_runner(
+        self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
+    ):
+        from sglang.srt.layers.moe.utils import (
+            get_moe_a2a_backend,
+            get_moe_runner_backend,
+        )
+
+        self.moe_runner_config = moe_runner_config
+        moe_runner_backend = get_moe_runner_backend()
+        if moe_runner_backend.is_auto() and get_moe_a2a_backend().supports_aiter():
+            moe_runner_backend = MoeRunnerBackend.AITER
+
+        if moe_runner_backend.is_aiter():
+            self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
+        else:
+            # TODO(cwan): refactor other backends
+            pass
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: StandardDispatchOutput,
+    ) -> CombineInput:
+        from sglang.srt.layers.moe.moe_runner.aiter import (
+            AiterMoeQuantInfo,
+            AiterQuantType,
+        )
+
+        if hasattr(torch, "float4_e2m1fn_x2"):
+            w13_weight = layer.w13_weight.view(torch.float4_e2m1fn_x2)
+            w2_weight = layer.w2_weight.view(torch.float4_e2m1fn_x2)
+        else:
+            w13_weight = layer.w13_weight
+            w2_weight = layer.w2_weight
+
+        if hasattr(layer.w13_weight, "is_shuffled"):
+            w13_weight.is_shuffled = True
+            w2_weight.is_shuffled = True
+
+        quant_info = AiterMoeQuantInfo(
+            w13_weight=w13_weight,
+            w2_weight=w2_weight,
+            quant_type=AiterQuantType.PER_1X32,
+            w13_scale=layer.w13_weight_scale,
+            w2_scale=layer.w2_weight_scale,
+            expert_mask=layer.dispatcher.expert_mask_gpu,
+        )
+        return self.runner.run(dispatch_output, quant_info)

python/sglang/srt/layers/quantization/quark/schemes/quark_w4a8_mxfp4_moe.py

@@ -0,0 +1,407 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+from dataclasses import replace
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
+from sglang.srt.layers.moe.utils import get_moe_weight_sizes
+from sglang.srt.layers.quantization.quark.schemes import QuarkMoEScheme
+from sglang.srt.layers.quantization.utils import all_close_1d
+from sglang.srt.utils import (
+    get_bool_env_var,
+    is_gfx95_supported,
+    is_hip,
+    round_up,
+    set_weight_attrs,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.token_dispatcher import (
+        CombineInput,
+        StandardDispatchOutput,
+    )
+
+logger = logging.getLogger(__name__)
+
+_is_shuffle_moe_mxfp4 = is_gfx95_supported()
+
+__all__ = ["QuarkW4A8MXFp4MoE"]
+
+_is_hip = is_hip()
+_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
+if _use_aiter:
+    from aiter.ops.shuffle import (
+        shuffle_scale,
+        shuffle_scale_a16w4,
+        shuffle_weight,
+        shuffle_weight_a16w4,
+    )
+
+OCP_MX_BLOCK_SIZE = 32
+
+
+class QuarkW4A8MXFp4MoE(QuarkMoEScheme):
+    """Quark MoE scheme for MXFP4 weights with static FP8 activations."""
+
+    def __init__(self, weight_config: dict[str, Any], input_config: dict[str, Any]):
+        self.weight_quant = weight_config
+        self.input_quant = input_config
+
+        weight_qscheme = self.weight_quant.get("qscheme")
+        input_qscheme = self.input_quant.get("qscheme")
+        weight_dtype = self.weight_quant.get("dtype")
+        input_dtype = self.input_quant.get("dtype")
+
+        if not (
+            weight_dtype == "fp4"
+            and weight_qscheme == "per_group"
+            and self.weight_quant.get("group_size") == OCP_MX_BLOCK_SIZE
+            and not self.weight_quant.get("is_dynamic")
+            and self.weight_quant.get("scale_format") == "e8m0"
+        ):
+            raise ValueError(
+                "For W4A8 MXFP4-FP8 Fused MoE layers, weights must be "
+                "static per-group FP4 with group_size=32 and e8m0 scales. "
+                f"Found {self.weight_quant}."
+            )
+
+        if not (
+            input_dtype in ("fp8_e4m3", "fp8_e4m3fn")
+            and input_qscheme == "per_tensor"
+            and not self.input_quant.get("is_dynamic")
+        ):
+            raise ValueError(
+                "For W4A8 MXFP4-FP8 Fused MoE layers, activations must be "
+                "static per-tensor fp8_e4m3/fp8_e4m3fn. "
+                f"Found {self.input_quant}."
+            )
+
+        self.with_bias = False
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 70
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        num_experts: int,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
+
+        self.num_experts = num_experts
+        self.with_bias = extra_weight_attrs.get("with_bias", False)
+        if _use_aiter:
+            intermediate_size_per_partition_after_pad = round_up(
+                intermediate_size_per_partition, 256
+            )
+            hidden_size = round_up(hidden_size, 256)
+            self.hidden_pad = hidden_size - layer.hidden_size
+            self.intermediate_pad = (
+                intermediate_size_per_partition_after_pad
+                - layer.intermediate_size_per_partition
+            )
+        else:
+            intermediate_size_per_partition_after_pad = intermediate_size_per_partition
+            self.hidden_pad = 0
+            self.intermediate_pad = 0
+
+        w13_up_dim, w2_down_dim, weight_padded = get_moe_weight_sizes(
+            intermediate_size_per_partition_after_pad,
+            is_aiter_moe=_use_aiter,
+            is_concat=True,
+            is_packed=True,
+        )
+        self.intermediate_size_per_partition = intermediate_size_per_partition_after_pad
+        self.hidden_size = hidden_size
+
+        # Add the quantization method used (per tensor/grouped/channel)
+        # to ensure the weight scales are loaded in properly.
+        extra_weight_attrs.update(
+            {
+                "quant_method": FusedMoeWeightScaleSupported.BLOCK.value,
+                "weight_padded": weight_padded,
+            },
+        )
+
+        weight_dtype = torch.uint8
+
+        # WEIGHTS
+        # MXFP4 weights are stored as uint8, with two FP4 values packed per
+        # byte. The AITER path later views these buffers as float4_e2m1fn_x2.
+        # Use ``zeros`` (not ``empty``) so the alignment padding (hidden
+        # 2880->3072, intermediate 2880->3072 for GPT-OSS) dequantizes to
+        # 0.0 if it ever reaches the matmul. The current AITER kernel
+        # skips the padded tail via ``n_pad_zeros`` / ``k_pad_zeros`` so
+        # this is defensive, but it matches ``Mxfp4MoEMethod``'s
+        # convention for the same kernel.
+        w13_weight = torch.nn.Parameter(
+            torch.zeros(
+                num_experts,
+                w13_up_dim,
+                hidden_size // 2,
+                dtype=weight_dtype,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight", w13_weight)
+        set_weight_attrs(w13_weight, extra_weight_attrs)
+
+        w2_weight = torch.nn.Parameter(
+            torch.zeros(
+                num_experts,
+                hidden_size,
+                w2_down_dim,
+                dtype=weight_dtype,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_weight", w2_weight)
+        set_weight_attrs(w2_weight, extra_weight_attrs)
+
+        w13_weight_bias = torch.nn.Parameter(
+            torch.zeros(
+                num_experts,
+                w13_up_dim,
+                dtype=torch.float32,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight_bias", w13_weight_bias)
+        set_weight_attrs(w13_weight_bias, extra_weight_attrs)
+
+        w2_weight_bias = torch.nn.Parameter(
+            torch.zeros(num_experts, hidden_size, dtype=torch.float32),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_weight_bias", w2_weight_bias)
+        set_weight_attrs(w2_weight_bias, extra_weight_attrs)
+
+        # WEIGHT_SCALES
+        # MXFP4 uses one e8m0 scale per 32-value block. These scales are
+        # loaded as uint8 and shuffled after loading for the kernel layout.
+        w13_weight_scale = torch.nn.Parameter(
+            torch.ones(
+                num_experts,
+                w13_up_dim,
+                hidden_size // OCP_MX_BLOCK_SIZE,
+                dtype=weight_dtype,
+            ),
+            requires_grad=False,
+        )
+        # 1. w2 scale is floor division of inter_dim by blockscale.
+        # 2. w2 scale needs to scale up just as w2.
+        # We combine 1. and 2. to keep the integer precision.
+        w2_weight_scale = torch.nn.Parameter(
+            torch.ones(
+                num_experts,
+                hidden_size,
+                (w2_down_dim * 2) // OCP_MX_BLOCK_SIZE,
+                dtype=weight_dtype,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(w2_weight_scale, extra_weight_attrs)
+        set_weight_attrs(w13_weight_scale, extra_weight_attrs)
+
+        layer.register_parameter("w13_weight_scale", w13_weight_scale)
+        layer.register_parameter("w2_weight_scale", w2_weight_scale)
+
+        # Add the quantization method used (per tensor/grouped/channel)
+        # to ensure the activation scales are loaded in properly.
+        extra_weight_attrs.update(
+            {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
+        )
+
+        # INPUT_SCALES
+        # W4A8 checkpoints carry static per-tensor FP8 activation scales for
+        # gate_up_proj and down_proj. These are separate from the MXFP4 weight
+        # block scales above.
+        w13_input_scale = torch.nn.Parameter(
+            torch.ones(num_experts, dtype=torch.float32),
+            requires_grad=False,
+        )
+        w2_input_scale = torch.nn.Parameter(
+            torch.ones(num_experts, dtype=torch.float32),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_input_scale", w13_input_scale)
+        layer.register_parameter("w2_input_scale", w2_input_scale)
+        set_weight_attrs(w13_input_scale, extra_weight_attrs)
+        set_weight_attrs(w2_input_scale, extra_weight_attrs)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        # Mirror native MXFP4 post-load shuffling. The default
+        # `SGLANG_USE_AITER_MOE_GU_ITLV=1` path uses the gate-up-aware
+        # a16w4 layout; the `=0` fallback keeps the separated gate/up layout.
+        # The Quark loader (`_load_quark_experts_weights` in
+        # `python/sglang/srt/models/gpt_oss.py`) already writes the
+        # SEPARATED-layout `[g0..g_{N-1}, u0..u_{N-1}]` buffer per expert,
+        # which is exactly the starting state the native path is in after
+        # its post-load `.view(e, n//2, 2, k).permute(0, 2, 1, 3)` step.
+        if envs.SGLANG_USE_AITER_MOE_GU_ITLV.get():
+            if _is_shuffle_moe_mxfp4:
+                layer.w13_weight.data = shuffle_weight_a16w4(
+                    layer.w13_weight.contiguous(), 16, True
+                )
+                layer.w2_weight.data = shuffle_weight_a16w4(
+                    layer.w2_weight.contiguous(), 16, False
+                )
+                layer.w13_weight.is_shuffled = True
+                layer.w2_weight.is_shuffled = True
+            shuffled_w13_scale = shuffle_scale_a16w4(
+                layer.w13_weight_scale.view(-1, layer.w13_weight_scale.shape[-1]),
+                self.num_experts,
+                True,
+            )
+            shuffled_w2_scale = shuffle_scale_a16w4(
+                layer.w2_weight_scale.view(-1, layer.w2_weight_scale.shape[-1]),
+                self.num_experts,
+                False,
+            )
+        else:
+            if _is_shuffle_moe_mxfp4:
+                layer.w13_weight.data = shuffle_weight(
+                    layer.w13_weight.contiguous(),
+                    is_guinterleave=False,
+                    gate_up=True,
+                )
+                layer.w2_weight.data = shuffle_weight(
+                    layer.w2_weight.contiguous(),
+                    is_guinterleave=False,
+                    gate_up=False,
+                )
+                layer.w13_weight.is_shuffled = True
+                layer.w2_weight.is_shuffled = True
+            shuffled_w13_scale = shuffle_scale(
+                layer.w13_weight_scale.view(-1, layer.w13_weight_scale.shape[-1]),
+                experts_cnt=self.num_experts,
+                is_guinterleave=False,
+                gate_up=True,
+            )
+            shuffled_w2_scale = shuffle_scale(
+                layer.w2_weight_scale.view(-1, layer.w2_weight_scale.shape[-1]),
+                experts_cnt=self.num_experts,
+                is_guinterleave=False,
+                gate_up=False,
+            )
+
+        layer.w13_weight_scale = torch.nn.Parameter(
+            shuffled_w13_scale, requires_grad=False
+        )
+        layer.w2_weight_scale = torch.nn.Parameter(
+            shuffled_w2_scale, requires_grad=False
+        )
+
+        # Static FP8 MoE kernels consume a single activation scale. Use the
+        # maximum if expert-local checkpoint scales differ.
+        if layer.w13_input_scale is None or layer.w2_input_scale is None:
+            raise ValueError("W4A8 MXFP4-FP8 MoE requires static input scales.")
+        if not all_close_1d(layer.w13_input_scale) or not all_close_1d(
+            layer.w2_input_scale
+        ):
+            logger.warning(
+                "Found input_scales that are not equal for W4A8 MXFP4-FP8 "
+                "MoE layer. Using the maximum across experts for each layer."
+            )
+        layer.w13_input_scale = torch.nn.Parameter(
+            layer.w13_input_scale.max().to(torch.float32), requires_grad=False
+        )
+        layer.w2_input_scale = torch.nn.Parameter(
+            layer.w2_input_scale.max().to(torch.float32), requires_grad=False
+        )
+
+        if hasattr(layer, "dispatcher"):
+            # Weights are stored as torch.uint8 but semantically MXFP4
+            layer.dispatcher.set_quant_config({"weight_dtype": torch.float4_e2m1fn_x2})
+
+    def create_moe_runner(
+        self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
+    ):
+        from sglang.srt.layers.moe.utils import (
+            get_moe_a2a_backend,
+            get_moe_runner_backend,
+        )
+
+        self.moe_runner_config = moe_runner_config
+        moe_runner_backend = get_moe_runner_backend()
+        if _use_aiter and get_moe_a2a_backend().supports_aiter():
+            moe_runner_backend = MoeRunnerBackend.AITER
+
+        if moe_runner_backend.is_aiter():
+            # MXFP4 hard-codes Swiglu in the AITER kernel path.
+            self.runner = MoeRunner(
+                moe_runner_backend, replace(moe_runner_config, activation="swiglu")
+            )
+        else:
+            raise NotImplementedError(
+                "QuarkW4A8MXFp4MoE is currently only supported with AITER."
+            )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: StandardDispatchOutput,
+    ) -> CombineInput:
+        from sglang.srt.layers.moe.moe_runner.aiter import (
+            AiterMoeQuantInfo,
+            AiterQuantType,
+        )
+
+        if hasattr(torch, "float4_e2m1fn_x2"):
+            w13_weight = layer.w13_weight.view(torch.float4_e2m1fn_x2)
+            w2_weight = layer.w2_weight.view(torch.float4_e2m1fn_x2)
+        else:
+            w13_weight = layer.w13_weight
+            w2_weight = layer.w2_weight
+
+        if hasattr(layer.w13_weight, "is_shuffled"):
+            w13_weight.is_shuffled = True
+            w2_weight.is_shuffled = True
+
+        x_padded = torch.nn.functional.pad(
+            dispatch_output.hidden_states,
+            (0, self.hidden_pad),
+            mode="constant",
+            value=0.0,
+        )
+        quant_info = AiterMoeQuantInfo(
+            w13_weight=w13_weight,
+            w2_weight=w2_weight,
+            quant_type=AiterQuantType.PER_1X32,
+            w13_scale=layer.w13_weight_scale,
+            w2_scale=layer.w2_weight_scale,
+            a13_scale=layer.w13_input_scale,
+            a2_scale=layer.w2_input_scale,
+            b13=layer.w13_weight_bias,
+            b2=layer.w2_weight_bias,
+            expert_mask=layer.dispatcher.expert_mask_gpu,
+            doweight_stage1=self.moe_runner_config.apply_router_weight_on_input,
+            hidden_pad=self.hidden_pad,
+            intermediate_pad=self.intermediate_pad,
+            # gpt-oss populates `gemm1_clamp_limit` (renamed in
+            # `models/gpt_oss.py` from `config.swiglu_limit`); DSv4 populates
+            # `swiglu_limit` directly. Accept either so the AITER `gate_mode`
+            # + `swiglu_limit` dispatch block in `moe_runner/aiter.py` (gated
+            # on `quant_info.swiglu_limit > 0`) is actually entered for both
+            # families. Mirrors the same fix PR #27201 applied to the native
+            # `Mxfp4MoEMethod.apply` path.
+            swiglu_limit=(
+                self.moe_runner_config.gemm1_clamp_limit
+                or self.moe_runner_config.swiglu_limit
+                or 0.0
+            ),
+        )
+        return self.runner.run(
+            dispatch_output._replace(hidden_states=x_padded), quant_info
+        )

python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py

@@ -0,0 +1,186 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any, Callable, Optional, cast
+
+import torch
+from torch.nn import Parameter
+
+from sglang.srt.layers.parameter import (
+    ChannelQuantScaleParameter,
+    ModelWeightParameter,
+    PerTensorScaleParameter,
+)
+from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
+from sglang.srt.layers.quantization.fp8_utils import (
+    apply_fp8_linear,
+    cutlass_fp8_supported,
+    normalize_e4m3fn_to_e4m3fnuz,
+)
+from sglang.srt.layers.quantization.quark.schemes import QuarkLinearScheme
+from sglang.srt.layers.quantization.utils import requantize_with_max_scale
+from sglang.srt.utils import get_bool_env_var, is_hip, set_weight_attrs
+
+__all__ = ["QuarkW8A8Fp8"]
+
+_is_fp8_fnuz = is_fp8_fnuz()
+_is_hip = is_hip()
+_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
+if _use_aiter:
+    from aiter.ops.shuffle import shuffle_weight
+
+
+class QuarkW8A8Fp8(QuarkLinearScheme):
+
+    def __init__(
+        self, weight_config: dict[str, Any], input_config: Optional[dict[str, Any]]
+    ):
+        self.cutlass_fp8_supported = cutlass_fp8_supported()
+        self.weight_qscheme = cast(str, weight_config.get("qscheme"))
+        self.is_static_input_scheme: bool = False
+        self.input_qscheme: Optional[str] = None
+        if input_config is not None:
+            self.is_static_input_scheme = not cast(bool, input_config.get("is_dynamic"))
+            self.input_qscheme = cast(str, input_config.get("qscheme"))
+
+        self.per_token = (
+            not self.is_static_input_scheme and self.input_qscheme == "per_channel"
+        )
+        self.out_dtype = torch.get_default_dtype()
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        # lovelace and up
+        return 89
+
+    def process_weights_after_loading(self, layer) -> None:
+        # If per tensor, when we have a fused module (e.g. QKV) with per
+        # tensor scales (thus N scales being passed to the kernel),
+        # requantize so we can always run per tensor
+        if self.weight_qscheme == "per_tensor":
+            if _is_fp8_fnuz:
+                input_scale = getattr(layer, "input_scale", None)
+                weight, max_w_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz(
+                    weight=layer.weight,
+                    weight_scale=layer.weight_scale,
+                    input_scale=input_scale,
+                )
+                if input_scale is not None:
+                    layer.input_scale = Parameter(input_scale, requires_grad=False)
+            else:
+                max_w_scale = layer.weight_scale
+                weight = layer.weight
+
+            max_w_scale, weight = requantize_with_max_scale(
+                weight=weight,
+                weight_scale=max_w_scale,
+                logical_widths=layer.logical_widths,
+            )
+
+            layer.weight = Parameter(weight.t(), requires_grad=False)
+            layer.weight_scale = Parameter(max_w_scale, requires_grad=False)
+
+        # If channelwise, scales are already lined up, so just transpose.
+        elif self.weight_qscheme == "per_channel":
+            weight = layer.weight
+
+            if _is_fp8_fnuz:
+                input_scale = getattr(layer, "input_scale", None)
+                weight, weight_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz(
+                    weight=weight,
+                    weight_scale=layer.weight_scale,
+                    input_scale=input_scale,
+                )
+                if input_scale is not None:
+                    layer.input_scale = Parameter(input_scale, requires_grad=False)
+            else:
+                weight_scale = layer.weight_scale.data
+            if self.per_token:
+                weight_scale = weight_scale.view(-1, 1)
+            if _use_aiter:
+                layer.weight = Parameter(
+                    shuffle_weight(weight, (16, 16)).t(), requires_grad=False
+                )
+            else:
+                layer.weight = Parameter(weight.t(), requires_grad=False)
+            # required by torch.compile to be torch.nn.Parameter
+            layer.weight_scale = Parameter(weight_scale, requires_grad=False)
+
+        else:
+            raise ValueError(f"Unknown quantization scheme {self.weight_qscheme}")
+
+        # INPUT SCALE
+        if self.is_static_input_scheme:
+            layer.input_scale = Parameter(layer.input_scale.max(), requires_grad=False)
+        else:
+            layer.input_scale = None
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        output_partition_sizes: list[int],
+        input_size_per_partition: int,
+        params_dtype: torch.dtype,
+        weight_loader: Callable,
+        **kwargs,
+    ):
+        output_size_per_partition = sum(output_partition_sizes)
+        layer.logical_widths = output_partition_sizes
+
+        # WEIGHT
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition,
+                dtype=torch.float8_e4m3fn,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+        # WEIGHT SCALE
+        if self.weight_qscheme == "per_channel":
+            weight_scale = ChannelQuantScaleParameter(
+                data=torch.empty((sum(output_partition_sizes)), dtype=torch.float32),
+                output_dim=0,
+                weight_loader=weight_loader,
+            )
+        else:
+            assert self.weight_qscheme == "per_tensor"
+            weight_scale = PerTensorScaleParameter(
+                data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
+                weight_loader=weight_loader,
+            )
+            set_weight_attrs(weight_scale, {"needs_scalar_to_array": True})
+
+        # min requirement for fp8 kernels
+        weight_scale[:] = torch.finfo(torch.float32).min
+        layer.register_parameter("weight_scale", weight_scale)
+
+        # INPUT SCALE
+        if self.is_static_input_scheme:
+            input_scale = PerTensorScaleParameter(
+                data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
+                weight_loader=weight_loader,
+            )
+            input_scale[:] = torch.finfo(torch.float32).min
+            set_weight_attrs(input_scale, {"needs_scalar_to_array": True})
+            layer.register_parameter("input_scale", input_scale)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        return apply_fp8_linear(
+            x,
+            layer.weight,
+            layer.weight_scale,
+            input_scale=layer.input_scale,
+            bias=bias,
+            cutlass_fp8_supported=self.cutlass_fp8_supported,
+            use_per_token_if_dynamic=self.per_token,
+        )

python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8_moe.py

@@ -0,0 +1,312 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
+from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
+from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz, scaled_fp8_quant
+from sglang.srt.layers.quantization.fp8_utils import normalize_e4m3fn_to_e4m3fnuz
+from sglang.srt.layers.quantization.quark.schemes import QuarkMoEScheme
+from sglang.srt.layers.quantization.utils import all_close_1d, per_tensor_dequantize
+from sglang.srt.utils import get_bool_env_var, is_hip, set_weight_attrs
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.token_dispatcher import (
+        CombineInput,
+        StandardDispatchOutput,
+    )
+
+logger = logging.getLogger(__name__)
+
+__all__ = ["QuarkW8A8FP8MoE"]
+
+_is_fp8_fnuz = is_fp8_fnuz()
+_is_hip = is_hip()
+_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
+if _use_aiter:
+    from aiter.ops.shuffle import shuffle_weight
+
+    from sglang.srt.layers.moe.rocm_moe_utils import rocm_fused_experts_tkw1
+
+
+class QuarkW8A8FP8MoE(QuarkMoEScheme):
+
+    def __init__(self, weight_config: dict[str, Any], input_config: dict[str, Any]):
+        self.is_static_input_scheme: bool = False
+        self.input_qscheme = None
+
+        if input_config is not None:
+            self.is_static_input_scheme = not input_config.get("is_dynamic")
+            self.input_qscheme = input_config.get("qscheme")
+
+        self.input_per_token = (
+            not self.is_static_input_scheme and self.input_qscheme == "per_channel"
+        )
+        self.weight_qscheme = weight_config.get("qscheme")
+        self.is_weight_per_channel = self.weight_qscheme == "per_channel"
+        self.out_dtype = torch.get_default_dtype()
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        # lovelace and up
+        return 89
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        num_experts: int,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
+
+        params_dtype = torch.float8_e4m3fn
+
+        # WEIGHTS
+        w13_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                2 * intermediate_size_per_partition,
+                hidden_size,
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight", w13_weight)
+        set_weight_attrs(w13_weight, extra_weight_attrs)
+
+        w2_weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                hidden_size,
+                intermediate_size_per_partition,
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_weight", w2_weight)
+        set_weight_attrs(w2_weight, extra_weight_attrs)
+
+        # WEIGHT_SCALES
+        # per-tensor quantization
+        if self.weight_qscheme == "per_tensor":
+            # Allocate 2 scales for w1 and w3 respectively.
+            # They will be combined to a single scale after weight loading.
+            w13_weight_scale = torch.nn.Parameter(
+                torch.ones(num_experts, 2, dtype=torch.float32), requires_grad=False
+            )
+            w2_weight_scale = torch.nn.Parameter(
+                torch.ones(num_experts, dtype=torch.float32), requires_grad=False
+            )
+            weight_quant_method = FusedMoeWeightScaleSupported.TENSOR.value
+        elif self.weight_qscheme == "per_channel":
+            w13_weight_scale = torch.nn.Parameter(
+                torch.ones(
+                    num_experts,
+                    2 * intermediate_size_per_partition,
+                    dtype=torch.float32,
+                ),
+                requires_grad=False,
+            )
+            w2_weight_scale = torch.nn.Parameter(
+                torch.ones(num_experts, hidden_size, dtype=torch.float32),
+                requires_grad=False,
+            )
+            weight_quant_method = FusedMoeWeightScaleSupported.CHANNEL.value
+        else:
+            raise ValueError(
+                f"Unsupported weight quantization strategy: {self.weight_qscheme}."
+            )
+
+        layer.register_parameter("w13_weight_scale", w13_weight_scale)
+        layer.register_parameter("w2_weight_scale", w2_weight_scale)
+        # Add the quantization method used (per tensor/grouped/channel)
+        # to ensure the weight scales are loaded in properly
+        extra_weight_attrs.update({"quant_method": weight_quant_method})
+        set_weight_attrs(w13_weight_scale, extra_weight_attrs)
+        set_weight_attrs(w2_weight_scale, extra_weight_attrs)
+
+        # INPUT_SCALES
+        if self.is_static_input_scheme:
+            assert (
+                self.input_qscheme == "per_tensor"
+            ), "Only per-tensor quantization is supported for static input scales"
+            w13_input_scale = torch.nn.Parameter(
+                torch.ones(num_experts, dtype=torch.float32), requires_grad=False
+            )
+            layer.register_parameter("w13_input_scale", w13_input_scale)
+            set_weight_attrs(w13_input_scale, extra_weight_attrs)
+
+            w2_input_scale = torch.nn.Parameter(
+                torch.ones(num_experts, dtype=torch.float32), requires_grad=False
+            )
+            layer.register_parameter("w2_input_scale", w2_input_scale)
+            set_weight_attrs(w2_input_scale, extra_weight_attrs)
+        else:
+            layer.w13_input_scale = None
+            layer.w2_input_scale = None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        # Fp8 moe kernels require a single activation scale.
+        # We take the max of all the scales in case they differ.
+        if self.is_static_input_scheme:
+            if layer.w13_input_scale is None or layer.w2_input_scale is None:
+                raise ValueError(
+                    "QuantConfig has static quantization, but found "
+                    "activation scales are None."
+                )
+            if not all_close_1d(layer.w13_input_scale) or not all_close_1d(
+                layer.w2_input_scale
+            ):
+                logger.warning(
+                    "Found input_scales that are not equal for "
+                    "fp8 MoE layer. Using the maximum across experts "
+                    "for each layer."
+                )
+            layer.w13_input_scale = torch.nn.Parameter(
+                layer.w13_input_scale.max(), requires_grad=False
+            )
+            layer.w2_input_scale = torch.nn.Parameter(
+                layer.w2_input_scale.max(), requires_grad=False
+            )
+
+        if _is_fp8_fnuz:
+            # Normalize the weights and scales
+            w13_weight, w13_weight_scale, w13_input_scale = (
+                normalize_e4m3fn_to_e4m3fnuz(
+                    layer.w13_weight, layer.w13_weight_scale, layer.w13_input_scale
+                )
+            )
+            w2_weight, w2_weight_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz(
+                layer.w2_weight, layer.w2_weight_scale, layer.w2_input_scale
+            )
+            # Reset the parameter
+            layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
+            layer.w13_weight_scale = torch.nn.Parameter(
+                w13_weight_scale, requires_grad=False
+            )
+            if w13_input_scale is not None:
+                layer.w13_input_scale = torch.nn.Parameter(
+                    w13_input_scale, requires_grad=False
+                )
+            layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
+            layer.w2_weight_scale = torch.nn.Parameter(
+                w2_weight_scale, requires_grad=False
+            )
+            if w2_input_scale is not None:
+                layer.w2_input_scale = torch.nn.Parameter(
+                    w2_input_scale, requires_grad=False
+                )
+        if self.weight_qscheme == "per_tensor":
+            # Fp8 moe kernel needs single weight scale for w13 per expert.
+            # We take the max then dequant and requant each expert.
+            assert layer.w13_weight_scale is not None
+            shard_size = layer.intermediate_size_per_partition
+            max_w13_scales = layer.w13_weight_scale.max(dim=1).values
+            for expert_id in range(layer.num_local_experts):
+                start = 0
+                for shard_id in range(2):
+                    dq_weight = per_tensor_dequantize(
+                        layer.w13_weight[expert_id][start : start + shard_size, :],
+                        layer.w13_weight_scale[expert_id][shard_id],
+                    )
+                    (
+                        layer.w13_weight[expert_id][start : start + shard_size, :],
+                        _,
+                    ) = scaled_fp8_quant(dq_weight, max_w13_scales[expert_id])
+
+                    start += shard_size
+
+            layer.w13_weight_scale = torch.nn.Parameter(
+                max_w13_scales, requires_grad=False
+            )
+        elif self.weight_qscheme == "per_channel":
+            layer.w13_weight_scale = torch.nn.Parameter(
+                layer.w13_weight_scale.unsqueeze(-1), requires_grad=False
+            )
+            layer.w2_weight_scale = torch.nn.Parameter(
+                layer.w2_weight_scale.unsqueeze(-1), requires_grad=False
+            )
+        else:
+            raise ValueError(
+                f"Unsupported weight quantization strategy: {self.weight_qscheme}."
+            )
+
+        if (
+            _use_aiter
+            and self.is_weight_per_channel
+            and self.moe_runner_config.apply_router_weight_on_input
+        ):
+            with torch.no_grad():
+                # Pre-shuffle weights
+                layer.w13_weight = torch.nn.Parameter(
+                    shuffle_weight(layer.w13_weight.data, (16, 16)),
+                    requires_grad=False,
+                )
+                torch.cuda.empty_cache()
+                layer.w2_weight = torch.nn.Parameter(
+                    shuffle_weight(layer.w2_weight.data, (16, 16)),
+                    requires_grad=False,
+                )
+                torch.cuda.empty_cache()
+
+    def create_moe_runner(
+        self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
+    ):
+        self.moe_runner_config = moe_runner_config
+        self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: StandardDispatchOutput,
+    ) -> CombineInput:
+
+        from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
+
+        x = dispatch_output.hidden_states
+        topk_output = dispatch_output.topk_output
+
+        moe_runner_config = self.moe_runner_config
+
+        if (
+            _use_aiter
+            and self.is_weight_per_channel
+            and moe_runner_config.apply_router_weight_on_input
+        ):
+            topk_weights, topk_ids, _ = topk_output
+            output = rocm_fused_experts_tkw1(
+                hidden_states=x,
+                w1=layer.w13_weight,
+                w2=layer.w2_weight,
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                activation=moe_runner_config.activation,
+                apply_router_weight_on_input=moe_runner_config.apply_router_weight_on_input,
+                use_fp8_w8a8=True,
+                per_channel_quant=self.is_weight_per_channel,
+                w1_scale=layer.w13_weight_scale,
+                w2_scale=layer.w2_weight_scale,
+                a1_scale=layer.w13_input_scale,
+                a2_scale=layer.w2_input_scale,
+            )
+            return StandardCombineInput(hidden_states=output)
+        else:
+            quant_info = TritonMoeQuantInfo(
+                w13_weight=layer.w13_weight,
+                w2_weight=layer.w2_weight,
+                use_fp8_w8a8=True,
+                per_channel_quant=self.is_weight_per_channel,
+                w13_scale=layer.w13_weight_scale,
+                w2_scale=layer.w2_weight_scale,
+                a13_scale=layer.w13_input_scale,
+                a2_scale=layer.w2_input_scale,
+            )
+            return self.runner.run(dispatch_output, quant_info)

python/sglang/srt/layers/quantization/quark/utils.py

@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+from collections.abc import Iterable, Mapping
+from types import MappingProxyType
+from typing import Any, Optional
+
+import torch
+
+try:
+    from aiter.ops.triton.quant import dynamic_mxfp4_quant
+except ImportError:
+
+    def raise_aiter_import_error(*args, **kwargs):
+        raise ImportError(
+            "Failed to import aiter. Make sure AITER is installed and accessible."
+        )
+
+    dynamic_mxfp4_quant = raise_aiter_import_error
+from torch import nn
+
+
+def deep_compare(dict1: Any, dict2: Any) -> bool:
+    if type(dict1) is not type(dict2):
+        return False
+    if isinstance(dict1, dict):
+        if dict1.keys() != dict2.keys():
+            return False
+        return all(deep_compare(dict1[k], dict2[k]) for k in dict1)
+    elif isinstance(dict1, list):
+        return set(dict1) == set(dict2)
+    else:
+        return dict1 == dict2
+
+
+def should_ignore_layer(
+    layer_name: Optional[str],
+    ignore: Iterable[str],
+    fused_mapping: Mapping[str, list[str]] = MappingProxyType({}),
+) -> bool:
+    if layer_name is None:
+        return False
+
+    # layer_name = model.layers.0.self_attn.qkv_proj
+    # proj_name = qkv_proj
+    proj_name = layer_name.split(".")[-1]
+
+    # Fused layers like gate_up_proj or qkv_proj will not be fused
+    # in the safetensors checkpoint. So, we convert the name
+    # from the fused version to unfused + check to make sure that
+    # each shard of the fused layer has the same scheme.
+    if proj_name in fused_mapping:
+        shard_proj_names = fused_mapping[proj_name]
+
+        # Convert fused_name --> [shard_names]
+        shard_names = [
+            layer_name.replace(proj_name, shard_proj_name)
+            for shard_proj_name in shard_proj_names
+        ]
+
+        # Layer should be ignored if shards are ignored.
+        should_ignore_layer = None
+        for shard_name in shard_names:
+            should_ignore_shard = check_equal_or_regex_match(
+                layer_name=shard_name, targets=ignore
+            )
+
+            # If shard_idx=0, set layer ignore to match shard.
+            if should_ignore_layer is None:
+                should_ignore_layer = should_ignore_shard
+
+            # If shard_idx=1+ confirm scheme matches prior shards.
+            elif should_ignore_shard != should_ignore_layer:
+                raise ValueError(
+                    f"Found different quantization schemes for "
+                    f"{shard_proj_names} in {layer_name}. SGLang "
+                    "requires all to use the same scheme."
+                )
+
+    # Unfused layers like down_proj and o_proj will match
+    # the safetensors checkpoint already.
+    else:
+        should_ignore_layer = check_equal_or_regex_match(
+            layer_name=layer_name, targets=ignore
+        )
+
+    assert should_ignore_layer is not None
+
+    return should_ignore_layer
+
+
+def check_equal_or_regex_match(layer_name: str, targets: Iterable[str]) -> bool:
+    """
+    Checks whether a layer_name is exactly equal or a regex match for
+    if target starts with 're:' to any target in list.
+    """
+    for target in targets:
+        if _is_equal_or_regex_match(layer_name, target):
+            return True
+    return False
+
+
+def _is_equal_or_regex_match(
+    value: str, target: str, check_contains: bool = False
+) -> bool:
+    """
+    Checks whether a value is exactly equal or a regex match for target
+    if target starts with 're:'. If check_contains is set to True,
+    additionally checks if the target string is contained within the value.
+    """
+
+    if target.startswith("re:"):
+        pattern = target[3:]
+        if re.match(pattern, value):
+            return True
+    elif check_contains:
+        if target.lower() in value.lower():
+            return True
+    elif target == value:
+        return True
+    return False
+
+
+# utility for tensor dims > 2 cases
+def b_dynamic_mxfp4_quant(x):
+    h, b, d = x.shape
+    x, x_scales = dynamic_mxfp4_quant(x.reshape(-1, d))
+    return x.view(h, b, d // 2), x_scales.view(h, b, d // 32)
+
+
+def mxfp4_to_f32(x, is_3d):
+    # 2 because we pack fp4 in uint8.
+    x = x.repeat_interleave(2, dim=-1)
+    if is_3d:
+        x[..., ::2] = x[..., ::2] & 0xF
+        x[..., 1::2] = x[..., 1::2] >> 4
+    else:
+        x[:, ::2] = x[:, ::2] & 0xF
+        x[:, 1::2] = x[:, 1::2] >> 4
+
+    mxfp4_list = [
+        0.0,
+        0.5,
+        1.0,
+        1.5,
+        2.0,
+        3.0,
+        4.0,
+        6.0,
+        -0.0,
+        -0.5,
+        -1.0,
+        -1.5,
+        -2.0,
+        -3.0,
+        -4.0,
+        -6.0,
+    ]
+    mxfp4_in_f32 = torch.tensor(mxfp4_list, dtype=torch.float32, device="cuda")
+    return mxfp4_in_f32[x.long()]
+
+
+def e8m0_to_f32(x):
+    # Per OCP MX-format v1.0: encoded 0..254 -> 2^(x-127); encoded 255 -> NaN.
+    # Detect the sentinel on the raw integer encoding, not on the float result
+    # (in float32, 2^128 overflows to +inf, so the old `x_f32 == 128` predicate
+    # both missed x=255 and wrongly NaN'd legitimate scale 128.0 at x=134).
+    x_f32 = 2 ** ((x.to(torch.float32)) - 127)
+    x_f32[x == 255] = float("nan")
+    return x_f32
+
+
+def quark_post_load_weights(self_attn: nn.Module, w: torch.Tensor, quant_format: str):
+    if "mxfp4" in quant_format:
+        # when dtype is bf16, the processing flow is to dynamic quantize bf16 tensor to uint8 tensor
+        # do w_kc (bf16) first to get the w_kc(uint8) w_s_kc(uint8)
+        # and w_vc repeating the same procedure of w_kc to get  w_vc(uint8) w_s_vc(uint8)
+        if w.dtype == torch.bfloat16:
+            w_kc, w_vc = w.unflatten(
+                0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim)
+            ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)
+            w_kc, w_s_kc = b_dynamic_mxfp4_quant(w_kc.transpose(-2, -1))
+            w_kc = w_kc.transpose(-2, -1)
+            w_s_kc = w_s_kc.transpose(-2, -1)
+            w_vc, w_s_vc = b_dynamic_mxfp4_quant(w_vc)
+            w_s_kc = w_s_kc.transpose(1, 2).contiguous().transpose(1, 2)
+            w_s_vc = w_s_vc.contiguous().transpose(1, 2)
+        elif w.dtype == torch.uint8:  # static quant for mxfp4
+            # when dtype is uint8, it means the w has been quantized to mxfp4 format
+            # but we must separate it to w_kc and w_vc.
+            # The quantized tensor size is only half of original tensor size
+            # and the scaling factor is 1/32, the transpose behavior will be not correct
+            # need to upcast it to fp32 to separate w to w_kc and w_vc
+            # to ensure the following transpose behavior is correct
+            # and then do mxfp4 quant again
+            w = mxfp4_to_f32(w, True).to(torch.bfloat16)
+            w_scales = self_attn.kv_b_proj.weight_scale.repeat_interleave(32, dim=-1)
+            w_scales = e8m0_to_f32(w_scales).to(torch.bfloat16)
+            w = w * w_scales
+            w_kc, w_vc = w.unflatten(
+                0, (-1, (self_attn.qk_nope_head_dim + self_attn.v_head_dim))
+            ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1)
+            w_kc, w_s_kc = b_dynamic_mxfp4_quant(w_kc.transpose(-2, -1))
+            w_kc = w_kc.transpose(-2, -1)
+            w_s_kc = w_s_kc.transpose(-2, -1)
+            w_vc, w_s_vc = b_dynamic_mxfp4_quant(w_vc)
+            w_s_kc = w_s_kc.transpose(1, 2).contiguous().transpose(1, 2)
+            w_s_vc = w_s_vc.contiguous().transpose(1, 2)
+
+        return w_kc, w_s_kc, w_vc, w_s_vc

python/sglang/srt/layers/quantization/quark/weights.py

@@ -0,0 +1,243 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+import re
+
+import torch
+
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.utils import is_cuda
+
+_is_cuda = is_cuda()
+
+
+def load_gptoss_weight_quark(
+    model,
+    weights,
+    *,
+    is_nextn: bool,
+    weight_name_mapping,
+) -> None:
+    # Regex matching `model.layers.{L}.mlp.experts.{N}.{gate_up_proj|down_proj}.{suffix}`
+    # used by the AMD Quark GPT-OSS per-expert checkpoint layout.
+    quark_expert_pat = re.compile(
+        r"^(.*\.mlp\.experts)\.(\d+)\.(gate_up_proj|down_proj)\."
+        r"(weight|weight_scale|input_scale|bias)$"
+    )
+    quark_experts_weights = []
+    normal_weights = []
+
+    for name, weight in weights:
+        if quark_expert_pat.match(name) is not None:
+            quark_experts_weights.append((name, weight))
+        else:
+            normal_weights.append((name, weight))
+
+    quark_loaded = _load_gptoss_quark_expert_weights(
+        model, quark_experts_weights, quark_expert_pat
+    )
+    model._load_normal_weights(
+        normal_weights,
+        is_nextn=is_nextn,
+        weight_name_mapping=weight_name_mapping,
+        other_loaded_param_names=quark_loaded,
+    )
+
+
+def _load_gptoss_quark_expert_weights(model, weights, quark_expert_pat):
+    """GPT-OSS per-expert style loader for Quark MoE tensors into padded fused buffers.
+
+    Quark stores each expert separately:
+        experts.{N}.gate_up_proj.{weight,weight_scale,input_scale,bias}
+        experts.{N}.down_proj.{weight,weight_scale,input_scale,bias}
+
+    We mirror the static MXFP4 expert loader: slice the checkpoint along
+    the TP-sharded dimension (intermediate axis) and copy into a window
+    of the padded ``w13_*`` / ``w2_*`` parameters allocated by
+    :class:`QuarkW4A8MXFp4MoE`. Down-proj bias is loaded only on
+    ``moe_tp_rank == 0`` to avoid double-counting after all-reduce.
+    """
+    params_dict = dict(model.named_parameters())
+    loaded_params: set[str] = set()
+    mxfp4_block = 32
+
+    moe_tp_rank = get_parallel().moe_tp_rank
+    moe_tp_size = get_parallel().moe_tp_size
+    moe_ep_rank = get_parallel().moe_ep_rank
+    moe_ep_size = get_parallel().moe_ep_size
+
+    intermediate_size = model.config.intermediate_size
+    assert (
+        intermediate_size % mxfp4_block == 0
+    ), f"{intermediate_size=} must be divisible by {mxfp4_block=}"
+    intermediate_size_block = intermediate_size // mxfp4_block
+
+    per_rank_intermediate_size_block = math.ceil(intermediate_size_block / moe_tp_size)
+
+    per_rank_intermediate_size = per_rank_intermediate_size_block * mxfp4_block
+
+    # Calculate common slicing bounds for current rank
+    assert model.config.num_local_experts % moe_ep_size == 0
+    moe_num_local_experts = model.config.num_local_experts // moe_ep_size
+
+    moe_tp_rank_start = moe_tp_rank * per_rank_intermediate_size
+    moe_tp_rank_end = min(
+        (moe_tp_rank + 1) * per_rank_intermediate_size, intermediate_size
+    )
+
+    moe_ep_rank_start = moe_ep_rank * moe_num_local_experts
+    moe_ep_rank_end = (moe_ep_rank + 1) * moe_num_local_experts
+
+    for name, weight in weights:
+        # Quark stores experts separately as
+        # `experts.{N}.{gate_up_proj|down_proj}.{suffix}`; pull the
+        # expert id out of the name (mxfp4 has it as axis 0 instead).
+        m = quark_expert_pat.match(name)
+        if m is None:
+            continue
+        prefix, expert_str, proj, suffix = m.groups()
+        global_expert_id = int(expert_str)
+        if global_expert_id < moe_ep_rank_start or global_expert_id >= moe_ep_rank_end:
+            continue
+        local_expert_id = global_expert_id - moe_ep_rank_start
+
+        if _is_cuda:
+            weight = weight.cuda()
+
+        dispatch_key = f"{proj}.{suffix}"
+
+        if dispatch_key == "gate_up_proj.weight":
+            # Handle MLP gate and up projection weights
+            new_name = f"{prefix}.w13_weight"
+
+            # De-interleave gate/up rows ([g0,u0,g1,u1,...] -> [g..., u...])
+            # then slice the TP window. Each half is written into its own
+            # slot of the padded fused buffer; the gap between halves is
+            # pre-zeroed by `create_weights` and must not be overwritten.
+            narrow_gate = weight[0::2][moe_tp_rank_start:moe_tp_rank_end].contiguous()
+            narrow_up = weight[1::2][moe_tp_rank_start:moe_tp_rank_end].contiguous()
+
+            param = params_dict[new_name]
+            intermediate_pad = param.data.shape[1] // 2
+            g0, g1 = narrow_gate.shape
+            u0, u1 = narrow_up.shape
+            param.data[local_expert_id, :g0, :g1].copy_(
+                narrow_gate.to(param.data.dtype)
+            )
+            param.data[
+                local_expert_id,
+                intermediate_pad : intermediate_pad + u0,
+                :u1,
+            ].copy_(narrow_up.to(param.data.dtype))
+            loaded_params.add(new_name)
+
+        elif dispatch_key == "down_proj.weight":
+            # Handle MLP down projection weights
+            # packed FP4 -> halve the TP bound on the contracting K dim
+            new_name = f"{prefix}.w2_weight"
+
+            narrow_weight = weight[
+                ...,
+                moe_tp_rank_start // 2 : moe_tp_rank_end // 2,
+            ]
+
+            param = params_dict[new_name]
+            d0, d1 = narrow_weight.shape
+            param.data[local_expert_id, :d0, :d1].copy_(
+                narrow_weight.to(param.data.dtype)
+            )
+            loaded_params.add(new_name)
+
+        elif dispatch_key == "gate_up_proj.weight_scale":
+            # Handle MLP gate and up projection weight scales
+            new_name = f"{prefix}.w13_weight_scale"
+
+            narrow_gate = weight[0::2][moe_tp_rank_start:moe_tp_rank_end].contiguous()
+            narrow_up = weight[1::2][moe_tp_rank_start:moe_tp_rank_end].contiguous()
+
+            param = params_dict[new_name]
+            intermediate_pad = param.data.shape[1] // 2
+            g0, g1 = narrow_gate.shape
+            u0, u1 = narrow_up.shape
+            param.data[local_expert_id, :g0, :g1].copy_(
+                narrow_gate.to(param.data.dtype)
+            )
+            param.data[
+                local_expert_id,
+                intermediate_pad : intermediate_pad + u0,
+                :u1,
+            ].copy_(narrow_up.to(param.data.dtype))
+            loaded_params.add(new_name)
+
+        elif dispatch_key == "down_proj.weight_scale":
+            # Handle MLP down projection weight scales
+            # 32 fp4 values per block -> slice by mxfp4_block
+            new_name = f"{prefix}.w2_weight_scale"
+
+            narrow_weight = weight[
+                ...,
+                moe_tp_rank_start // mxfp4_block : moe_tp_rank_end // mxfp4_block,
+            ]
+
+            param = params_dict[new_name]
+            d0, d1 = narrow_weight.shape
+            param.data[local_expert_id, :d0, :d1].copy_(
+                narrow_weight.to(param.data.dtype)
+            )
+            loaded_params.add(new_name)
+
+        elif dispatch_key == "gate_up_proj.bias":
+            # Handle MLP gate and up projection biases
+            new_name = f"{prefix}.w13_weight_bias"
+
+            narrow_gate = weight[0::2][moe_tp_rank_start:moe_tp_rank_end].contiguous()
+            narrow_up = weight[1::2][moe_tp_rank_start:moe_tp_rank_end].contiguous()
+
+            param = params_dict[new_name]
+            intermediate_pad = param.data.shape[1] // 2
+            param.data[local_expert_id, : narrow_gate.shape[0]].copy_(
+                narrow_gate.to(param.data.dtype)
+            )
+            param.data[
+                local_expert_id,
+                intermediate_pad : intermediate_pad + narrow_up.shape[0],
+            ].copy_(narrow_up.to(param.data.dtype))
+            loaded_params.add(new_name)
+
+        elif dispatch_key == "down_proj.bias":
+            # Handle MLP down projection bias
+            # Only TP rank 0 owns the bias; others zero out so the
+            # post-MoE all-reduce sums to the correct value once.
+            narrow_weight = weight
+            if moe_tp_rank != 0:
+                narrow_weight = torch.zeros_like(narrow_weight)
+
+            new_name = f"{prefix}.w2_weight_bias"
+            param = params_dict[new_name]
+            d0 = narrow_weight.shape[0]
+            param.data[local_expert_id, :d0].copy_(narrow_weight.to(param.data.dtype))
+            loaded_params.add(new_name)
+
+        elif dispatch_key == "gate_up_proj.input_scale":
+            # Handle MLP gate/up FP8 activation scale (per-tensor scalar)
+            new_name = f"{prefix}.w13_input_scale"
+            if new_name not in params_dict:
+                # Scheme didn't allocate the parameter (e.g. W4A16); skip.
+                continue
+
+            param = params_dict[new_name]
+            param.data[local_expert_id].copy_(weight.to(param.data.dtype).reshape(()))
+            loaded_params.add(new_name)
+
+        elif dispatch_key == "down_proj.input_scale":
+            # Handle MLP down FP8 activation scale (per-tensor scalar)
+            new_name = f"{prefix}.w2_input_scale"
+            if new_name not in params_dict:
+                # Scheme didn't allocate the parameter (e.g. W4A16); skip.
+                continue
+
+            param = params_dict[new_name]
+            param.data[local_expert_id].copy_(weight.to(param.data.dtype).reshape(()))
+            loaded_params.add(new_name)
+
+    return loaded_params