chore: import upstream snapshot with attribution

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)