From 9485c083bbc8d48636a84fc41441a12bdaa93573 Mon Sep 17 00:00:00 2001 From: LinyuanLi Date: Wed, 19 Aug 2026 00:06:05 +0800 Subject: [PATCH] [NPU] Add mxfp4-w4a4 MOE Quantization Support for NPU (#30319) --- docs/docs/advanced_features/quantization.mdx | 3 +- .../ascend-npus/optimization/quantization.mdx | 25 ++++++ .../srt/hardware_backend/npu/moe/quant.py | 11 +-- .../npu/quantization/moe_methods.py | 80 ++++++++++++++++++ .../quantization/modelslim/modelslim.py | 2 + .../modelslim/schemes/__init__.py | 2 + .../schemes/modelslim_w4a4_mxfp4_moe.py | 82 +++++++++++++++++++ 7 files changed, 199 insertions(+), 6 deletions(-) create mode 100644 python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_mxfp4_moe.py diff --git a/docs/docs/advanced_features/quantization.mdx b/docs/docs/advanced_features/quantization.mdx index 5d9016a72..cd32e5055 100644 --- a/docs/docs/advanced_features/quantization.mdx +++ b/docs/docs/advanced_features/quantization.mdx @@ -47,7 +47,7 @@ The following table summarizes quantization method support across NVIDIA and AMD Yes Yes Yes (A5) - On GPU: requires CDNA3/CDNA4 with MXFP support (uses Aiter). On Ascend NPU (A5): W4A4 MXFP4 for Qwen3 dense LLM (MXFP4 weights + activations) — online uses dual-level MXFP4, offline W4A4_MXFP4 checkpoints (single-level) are auto-detected via modelslim + On GPU: requires CDNA3/CDNA4 with MXFP support (uses Aiter). On Ascend NPU (A5): W4A4 MXFP4 for Qwen3 dense and MoE LLMs (MXFP4 weights + activations) — dense models support online dual-level MXFP4; offline W4A4_MXFP4 dense and MoE checkpoints (single-level) are auto-detected via modelslim mxfp8 @@ -807,6 +807,7 @@ MindStudio-ModelSlim (msModelSlim) is a model offline quantization compression t - [x] ```W8A8``` linear with offline quantization of activations - [x] ```W8A8_DYNAMIC``` linear with online quantization of activations - [x] ```W4A4_DYNAMIC``` MOE with online quantization of activations + - [x] ```W4A4_MXFP4``` MOE with dynamic MXFP4 activation quantization - [x] ```W4A8_DYNAMIC``` MOE with online quantization of activations - [x] ```W4A8_MXFP``` MOE with dynamic MXFP8 activation quantization - [x] ```W8A8_DYNAMIC``` MOE with online quantization of activations diff --git a/docs/docs/hardware-platforms/ascend-npus/optimization/quantization.mdx b/docs/docs/hardware-platforms/ascend-npus/optimization/quantization.mdx index 660035a56..17048ebf9 100644 --- a/docs/docs/hardware-platforms/ascend-npus/optimization/quantization.mdx +++ b/docs/docs/hardware-platforms/ascend-npus/optimization/quantization.mdx @@ -84,6 +84,14 @@ SGLang supports **mix-bits** quantization (independently defines and loads each WIP x + + MXFP4 W4A4 (ModelSlim) + MoE + x + x + √ + x + W4A4 dynamic MoE @@ -466,6 +474,23 @@ python3 -m sglang.launch_server \ > - As with W4A8, the packed-FP4 dtype passed to the NPU ops (`dst_type` / `x2_dtype`) must be resolved from `torch_npu.float4_e2m1fn_x2` (an int enum), not the `torch.float4_e2m1fn_x2` dtype object, which recent op-plugin builds reject. > - Validated end-to-end on Ascend A5 hardware. +**ModelSlim W4A4 MXFP4 for LLM MoE models:** + +SGLang auto-detects offline ModelSlim `W4A4_MXFP4` MoE checkpoints from `quant_model_description.json`; do not pass `--quantization`. This path requires Ascend 950 products or newer. + +```bash Command +MODEL_PATH=/path/to/w4a4-mxfp4-moe-model +python3 -m sglang.launch_server \ + --model-path "$MODEL_PATH" \ + --device npu --attention-backend ascend \ + --host 0.0.0.0 --port 30000 \ + --tp-size 1 +``` + +> **Implementation Notes:** +> - ModelSlim supplies packed MXFP4 `w13` and `w2` expert weights with UE8M0 block scales (block size 32). +> - Ascend TP and DeepEP dispatch activations as BF16; this path does not request MXFP4 dispatch. SGLang dynamically quantizes each expert input to MXFP4 immediately before grouped matmul. + ## Diffusion Model Quantization on Ascend NPU SGLang-Diffusion supports MXFP8 online and offline quantization for diffusion models (such as Wan2.2) on Ascend NPUs. MXFP8 requires Ascend 950 Products; the ModelSlim W8A8/W4A4 schemes work on A2/A3. diff --git a/python/sglang/srt/hardware_backend/npu/moe/quant.py b/python/sglang/srt/hardware_backend/npu/moe/quant.py index afb0bb6b9..7c59c3bfc 100644 --- a/python/sglang/srt/hardware_backend/npu/moe/quant.py +++ b/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 diff --git a/python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py b/python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py index 84099f122..6cfe9fb32 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py +++ b/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 # --------------------------------------------------------------------------- diff --git a/python/sglang/srt/layers/quantization/modelslim/modelslim.py b/python/sglang/srt/layers/quantization/modelslim/modelslim.py index d90b3d7fb..83ed583b0 100644 --- a/python/sglang/srt/layers/quantization/modelslim/modelslim.py +++ b/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), diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py b/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py index 3c9178575..ca19ec562 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py +++ b/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", diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_mxfp4_moe.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_mxfp4_moe.py new file mode 100644 index 000000000..87a111ea7 --- /dev/null +++ b/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)