[NPU] Add mxfp4-w4a4 MOE Quantization Support for NPU (#30319)

python/sglang/srt/hardware_backend/npu/moe/quant.py

@@ -27,16 +27,17 @@ class HiddenStatesDynamicQuant(BaseHiddenStatesQuant):
     """
     Dynamic per‑token quantisation of hidden states.
 
-    ``torch.float8_e4m3fn`` selects the MX (block-scaled) op, whose scale is a
-    ``float8_e8m0fnu`` block scale ``[N, K//64, 2]`` rather than one scalar per
-    token; the int8/int4 dtypes keep the plain per-token op.
+    ``torch.float8_e4m3fn`` selects the MX (block-scaled) op. Set
+    ``use_mx_quant`` for other MX dtypes whose NPU op argument is not represented
+    by the matching torch dtype object; the int8/int4 dtypes keep the plain
+    per-token op.
 
     Returns ``(quantized_hidden_states, per‑token_scale)``.
     """
 
-    def __init__(self, quant_dtype: torch.dtype) -> None:
+    def __init__(self, quant_dtype: torch.dtype, use_mx_quant: bool = False) -> None:
         super().__init__(quant_dtype)
-        if quant_dtype == torch.float8_e4m3fn:
+        if use_mx_quant or quant_dtype == torch.float8_e4m3fn:
             self._op = torch.ops.npu.npu_dynamic_mx_quant
         elif quant_dtype in (torch.int8, torch.quint4x2):
             self._op = torch.ops.npu.npu_dynamic_quant

python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py

@@ -270,6 +270,86 @@ class NPUW4A8MXFP4MoEMethod(_NPUMoEMethodBase):
         )
 
 
+# ---------------------------------------------------------------------------
+#  NPUW4A4MXFP4MoEMethod
+# ---------------------------------------------------------------------------
+class NPUW4A4MXFP4MoEMethod(_NPUMoEMethodBase):
+    """ModelSlim W4A4 MXFP4 MoE with single-level FP4 weights and activations."""
+
+    def __init__(self):
+        super().__init__(quant_config=None)
+        self.matmul = GroupedMatmul()
+        fp4_dtype = _get_float4_e2m1fn_x2_dtype()
+        if fp4_dtype is None:
+            raise RuntimeError("NPU W4A4 MXFP4 MoE requires float4 support.")
+        self.hidden_states_quantizer = HiddenStatesDynamicQuant(
+            quant_dtype=fp4_dtype,
+            use_mx_quant=True,
+        )
+
+    def process_weights_after_loading(
+        self, layer: torch.nn.Module, weight_prefix: str
+    ) -> None:
+        self._validate_weight_prefix(layer, weight_prefix)
+
+        weight = getattr(layer, f"{weight_prefix}_weight")
+        weight.data = npu_format_cast(weight.data).transpose(-1, -2)
+
+        weight_scale = getattr(layer, f"{weight_prefix}_weight_scale")
+        scale = weight_scale.data.reshape(
+            weight_scale.shape[0],
+            weight_scale.shape[1],
+            weight_scale.shape[2] // 2,
+            2,
+        ).transpose(1, 2)
+        weight_scale.data = scale
+
+        # The refactored Ascend dispatchers currently support BF16 and INT8.
+        # Keep dispatch in BF16 and quantize immediately before each GMM.
+        if weight_prefix == "w13":
+            self._set_dispatcher_output_dtype(layer, "bf16")
+
+    def apply(
+        self,
+        quant_info: "AscendQuantInfo",
+        hidden_states: torch.Tensor,
+        expert_tokens: torch.Tensor,
+        pertoken_scale: Optional[torch.Tensor],
+        output_dtype: torch.dtype,
+        weight_prefix: str,
+        group_list_type: int,
+    ) -> torch.Tensor:
+        fp4_dtype = self.hidden_states_quantizer.quant_dtype
+        e8m0_dtype = _require_e8m0_dtype()
+
+        if pertoken_scale is None:
+            hidden_states, pertoken_scale = self.hidden_states_quantizer(hidden_states)
+        elif pertoken_scale is not None:
+            pertoken_scale = pertoken_scale.reshape(
+                hidden_states.shape[0], hidden_states.shape[1] // 32, 2
+            )
+
+        scale_args: Dict[str, Any] = {
+            "scale": [getattr(quant_info, f"{weight_prefix}_weight_scale", None)],
+            "scale_dtype": e8m0_dtype,
+            "per_token_scale": [pertoken_scale],
+            "per_token_scale_dtype": e8m0_dtype,
+            "x_dtype": fp4_dtype,
+            "weight_dtype": fp4_dtype,
+        }
+        scale_args.update(self._get_bias_args(quant_info, weight_prefix))
+        return self.matmul.forward(
+            quant_info,
+            weight_prefix,
+            hidden_states,
+            expert_tokens.to(torch.int64),
+            output_dtype,
+            group_list_type=group_list_type,
+            transposed=True,
+            **scale_args,
+        )
+
+
 # ---------------------------------------------------------------------------
 #  NPUW4A4Int4DynamicMoEMethod
 # ---------------------------------------------------------------------------

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

@@ -22,6 +22,7 @@ from sglang.srt.layers.quantization.modelslim.schemes import (
     ModelSlimMXFP8Scheme,
     ModelSlimW4A4Int4,
     ModelSlimW4A4Int4MoE,
+    ModelSlimW4A4MXFP4MoE,
     ModelSlimW4A8Int8MoE,
     ModelSlimW4A8MXFP4MoE,
     ModelSlimW8A8Int8,
@@ -336,6 +337,7 @@ class ModelSlimConfig(QuantizationConfig):
         prefix: str,
     ):
         moe_quant_schemes = [
+            ("W4A4_MXFP4", ModelSlimW4A4MXFP4MoE),
             ("W4A8_MXFP", ModelSlimW4A8MXFP4MoE),
             ("W4A4_DYNAMIC", ModelSlimW4A4Int4MoE),
             ("W4A8_DYNAMIC", ModelSlimW4A8Int8MoE),

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

@@ -13,6 +13,7 @@ from .modelslim_mxfp4 import ModelSlimMXFP4Scheme
 from .modelslim_mxfp8_moe import ModelSlimMXFP8MoEScheme
 from .modelslim_w4a4_int4 import ModelSlimW4A4Int4
 from .modelslim_w4a4_int4_moe import ModelSlimW4A4Int4MoE
+from .modelslim_w4a4_mxfp4_moe import ModelSlimW4A4MXFP4MoE
 from .modelslim_w4a8_int8_moe import ModelSlimW4A8Int8MoE
 from .modelslim_w4a8_mxfp4_moe import ModelSlimW4A8MXFP4MoE
 from .modelslim_w8a8_int8 import ModelSlimW8A8Int8
@@ -25,6 +26,7 @@ __all__ = [
     "ModelSlimMXFP4W4A8Scheme",
     "ModelSlimMXFP4Scheme",
     "ModelSlimMXFP8MoEScheme",
+    "ModelSlimW4A4MXFP4MoE",
     "ModelSlimW4A8MXFP4MoE",
     "ModelSlimW8A8Int8",
     "ModelSlimW4A4Int4",

python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_mxfp4_moe.py

@@ -0,0 +1,82 @@
+"""ModelSlim W4A4_MXFP4 MoE scheme for Ascend NPU."""
+
+from __future__ import annotations
+
+from typing import Any, Dict
+
+import torch
+
+from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
+    NPUW4A4MXFP4MoEMethod,
+)
+from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimMoEScheme
+from sglang.srt.utils import set_weight_attrs
+
+MXFP4_BLOCK_SIZE = 32
+
+__all__ = ["ModelSlimW4A4MXFP4MoE"]
+
+
+class ModelSlimW4A4MXFP4MoE(ModelSlimMoEScheme):
+    """Create one ModelSlim MXFP4 expert-weight group (w13 or w2)."""
+
+    def __init__(
+        self,
+        quant_config: Dict[str, Any],
+        weight_prefix: str,
+    ) -> None:
+        if weight_prefix not in ("w13", "w2"):
+            raise ValueError(
+                f"weight_prefix must be 'w13' or 'w2', got '{weight_prefix}'"
+            )
+        self.quant_config = quant_config
+        self.weight_prefix = weight_prefix
+        self.kernel = NPUW4A4MXFP4MoEMethod()
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        num_experts: int,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        **extra_weight_attrs,
+    ) -> None:
+        from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
+
+        extra_weight_attrs.update(
+            {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value}
+        )
+
+        if self.weight_prefix == "w13":
+            output_size = 2 * intermediate_size_per_partition
+            input_size = hidden_size
+        else:
+            output_size = hidden_size
+            input_size = intermediate_size_per_partition
+
+        weight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                output_size,
+                input_size // 2,
+                dtype=torch.uint8,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter(f"{self.weight_prefix}_weight", weight)
+        set_weight_attrs(weight, extra_weight_attrs)
+
+        weight_scale = torch.nn.Parameter(
+            torch.zeros(
+                num_experts,
+                output_size,
+                (input_size + MXFP4_BLOCK_SIZE - 1) // MXFP4_BLOCK_SIZE,
+                dtype=torch.uint8,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter(f"{self.weight_prefix}_weight_scale", weight_scale)
+        set_weight_attrs(weight_scale, extra_weight_attrs)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        self.kernel.process_weights_after_loading(layer, self.weight_prefix)