From 0da87024d3a02a59d9f302f0df3b30ba2bb6233f Mon Sep 17 00:00:00 2001 From: LinyuanLi Date: Sun, 16 Aug 2026 14:03:17 +0800 Subject: [PATCH] [NPU] Add mxfp4-w4a8 MOE Quantization Support for NPU (#30318) --- docs/docs/advanced_features/quantization.mdx | 3 +- .../ascend-npus/optimization/quantization.mdx | 25 ++++++ .../npu/quantization/moe_methods.py | 86 +++++++++++++++++++ .../quantization/modelslim/modelslim.py | 6 +- .../modelslim/schemes/__init__.py | 2 + .../schemes/modelslim_w4a8_mxfp4_moe.py | 83 ++++++++++++++++++ 6 files changed, 202 insertions(+), 3 deletions(-) create mode 100644 python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a8_mxfp4_moe.py diff --git a/docs/docs/advanced_features/quantization.mdx b/docs/docs/advanced_features/quantization.mdx index 92abe9989..5d9016a72 100644 --- a/docs/docs/advanced_features/quantization.mdx +++ b/docs/docs/advanced_features/quantization.mdx @@ -61,7 +61,7 @@ The following table summarizes quantization method support across NVIDIA and AMD No No Yes (A5) - Ascend NPU only; online W4A8 for Qwen3 dense LLM (MXFP4 weights + MXFP8 activations) on A5 series; offline W4A8_MXFP checkpoints are auto-detected via modelslim + Ascend NPU only; online W4A8 for Qwen3 dense LLM (MXFP4 weights + MXFP8 activations) on A5 series; offline W4A8_MXFP dense and MoE checkpoints are auto-detected via modelslim blockwise_int8 @@ -808,6 +808,7 @@ MindStudio-ModelSlim (msModelSlim) is a model offline quantization compression t - [x] ```W8A8_DYNAMIC``` linear with online quantization of activations - [x] ```W4A4_DYNAMIC``` MOE with online quantization of activations - [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 - [ ] ```W4A8``` linear TBD - [ ] ```W4A16``` linear TBD diff --git a/docs/docs/hardware-platforms/ascend-npus/optimization/quantization.mdx b/docs/docs/hardware-platforms/ascend-npus/optimization/quantization.mdx index 8a7c01667..660035a56 100644 --- a/docs/docs/hardware-platforms/ascend-npus/optimization/quantization.mdx +++ b/docs/docs/hardware-platforms/ascend-npus/optimization/quantization.mdx @@ -68,6 +68,14 @@ SGLang supports **mix-bits** quantization (independently defines and loads each √ x + + MXFP4 W4A8 (ModelSlim) + MoE + x + x + √ + x + MXFP4 W4A4 Linear @@ -418,6 +426,23 @@ python3 -m sglang.launch_server \ > - The packed-FP4 dtype passed to the NPU ops (`dst_type` / `x2_dtype` / `input_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. > - Online and offline share the same kernel path and layout; they differ only in the weight source (RTN at load vs msmodelslim calibration). +**ModelSlim W4A8 MXFP4 for LLM MoE models:** + +SGLang auto-detects offline ModelSlim `W4A8_MXFP` MoE checkpoints from `quant_model_description.json`; do not pass `--quantization`. This path requires Ascend A5 or newer. + +```bash Command +MODEL_PATH=/path/to/w4a8-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 MXFP8 dispatch. SGLang dynamically quantizes each expert input to MXFP8 immediately before grouped matmul. + **MXFP4 W4A4 for LLM dense models (e.g. Qwen3 / Qwen3.5):** LLM dense W4A4 (MXFP4 4-bit weights + 4-bit activations) Linear support was added in [PR #23795](https://github.com/sgl-project/sglang/pull/23795). Requires Ascend A5 series (Ascend 950) or newer — the dual-level online path uses the `DualLevelQuantBatchMatmul` op, which A2/A3 lack. On the Ascend NPU backend `--quantization mxfp4` selects this W4A4 path (on GPU the same flag selects the upstream OCP MXFP4 MoE config instead). 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 33fa8e54f..84099f122 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py @@ -20,6 +20,7 @@ from sglang.srt.hardware_backend.npu.moe.matmul import ( ) from sglang.srt.hardware_backend.npu.moe.quant import HiddenStatesDynamicQuant from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import ( + _get_float4_e2m1fn_x2_dtype, _get_float8_e8m0fnu_dtype, ) @@ -184,6 +185,91 @@ class _NPUMoEMethodBase(FusedMoEMethodBase): return {"bias": [bias]} if bias is not None else {} +# --------------------------------------------------------------------------- +# NPUW4A8MXFP4MoEMethod +# --------------------------------------------------------------------------- +class NPUW4A8MXFP4MoEMethod(_NPUMoEMethodBase): + """ModelSlim W4A8 MoE with packed MXFP4 weights and MXFP8 activations.""" + + def __init__(self): + super().__init__(quant_config=None) + self.matmul = GroupedMatmul() + self.hidden_states_quantizer = HiddenStatesDynamicQuant( + quant_dtype=torch.float8_e4m3fn + ) + + def process_weights_after_loading( + self, layer: torch.nn.Module, weight_prefix: str + ) -> None: + self._validate_weight_prefix(layer, weight_prefix) + + fp4_dtype = _get_float4_e2m1fn_x2_dtype() + if fp4_dtype is None: + raise RuntimeError("NPU W4A8 MXFP MoE requires float4 support.") + + weight = getattr(layer, f"{weight_prefix}_weight") + weight.data = npu_format_cast( + weight.data, + customize_dtype=torch.float8_e4m3fn, + input_dtype=fp4_dtype, + ).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 to MXFP8 immediately before 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 = _get_float4_e2m1fn_x2_dtype() + if fp4_dtype is None: + raise RuntimeError("NPU W4A8 MXFP MoE requires float4 support.") + 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] // 64, 2 + ) + + 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=None, + scale_dtype=None, + per_token_scale=[pertoken_scale], + antiquant_scale=[ + getattr(quant_info, f"{weight_prefix}_weight_scale", None) + ], + x_dtype=torch.float8_e4m3fn, + weight_dtype=fp4_dtype, + per_token_scale_dtype=e8m0_dtype, + ) + + # --------------------------------------------------------------------------- # NPUW4A4Int4DynamicMoEMethod # --------------------------------------------------------------------------- diff --git a/python/sglang/srt/layers/quantization/modelslim/modelslim.py b/python/sglang/srt/layers/quantization/modelslim/modelslim.py index 7a02601b1..65f74a306 100644 --- a/python/sglang/srt/layers/quantization/modelslim/modelslim.py +++ b/python/sglang/srt/layers/quantization/modelslim/modelslim.py @@ -23,6 +23,7 @@ from sglang.srt.layers.quantization.modelslim.schemes import ( ModelSlimW4A4Int4, ModelSlimW4A4Int4MoE, ModelSlimW4A8Int8MoE, + ModelSlimW4A8MXFP4MoE, ModelSlimW8A8Int8, ModelSlimW8A8Int8MoE, ) @@ -275,6 +276,7 @@ class ModelSlimConfig(QuantizationConfig): prefix: str, ): moe_quant_schemes = [ + ("W4A8_MXFP", ModelSlimW4A8MXFP4MoE), ("W4A4_DYNAMIC", ModelSlimW4A4Int4MoE), ("W4A8_DYNAMIC", ModelSlimW4A8Int8MoE), ("W8A8_DYNAMIC", ModelSlimW8A8Int8MoE), @@ -532,8 +534,8 @@ class ModelSlimFusedMoEMethod(FusedMoEMethodBase): w2_weight=layer.w2_weight, w13_weight_scale=layer.w13_weight_scale, w2_weight_scale=layer.w2_weight_scale, - w13_weight_offset=layer.w13_weight_offset, - w2_weight_offset=layer.w2_weight_offset, + w13_weight_offset=getattr(layer, "w13_weight_offset", None), + w2_weight_offset=getattr(layer, "w2_weight_offset", None), w13_scale_bias=getattr(layer, "w13_scale_bias", None), w2_scale_bias=getattr(layer, "w2_scale_bias", None), w13_weight_bias=getattr(layer, "w13_weight_bias", None), diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py b/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py index ca03b17cd..3c9178575 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py @@ -14,6 +14,7 @@ from .modelslim_mxfp8_moe import ModelSlimMXFP8MoEScheme from .modelslim_w4a4_int4 import ModelSlimW4A4Int4 from .modelslim_w4a4_int4_moe import ModelSlimW4A4Int4MoE from .modelslim_w4a8_int8_moe import ModelSlimW4A8Int8MoE +from .modelslim_w4a8_mxfp4_moe import ModelSlimW4A8MXFP4MoE from .modelslim_w8a8_int8 import ModelSlimW8A8Int8 from .modelslim_w8a8_int8_moe import ModelSlimW8A8Int8MoE @@ -24,6 +25,7 @@ __all__ = [ "ModelSlimMXFP4W4A8Scheme", "ModelSlimMXFP4Scheme", "ModelSlimMXFP8MoEScheme", + "ModelSlimW4A8MXFP4MoE", "ModelSlimW8A8Int8", "ModelSlimW4A4Int4", "ModelSlimW4A4Int4MoE", diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a8_mxfp4_moe.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a8_mxfp4_moe.py new file mode 100644 index 000000000..634d8faf7 --- /dev/null +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a8_mxfp4_moe.py @@ -0,0 +1,83 @@ +"""ModelSlim W4A8_MXFP 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 ( + NPUW4A8MXFP4MoEMethod, +) +from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimMoEScheme +from sglang.srt.utils import set_weight_attrs + +W4A8_MXFP4_BLOCK_SIZE = 32 +W4A8_MXFP4_PACK_FACTOR = 2 + +__all__ = ["ModelSlimW4A8MXFP4MoE"] + + +class ModelSlimW4A8MXFP4MoE(ModelSlimMoEScheme): + """Create one ModelSlim W4A8 MXFP 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 = NPUW4A8MXFP4MoEMethod() + + 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 // W4A8_MXFP4_PACK_FACTOR, + 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 + W4A8_MXFP4_BLOCK_SIZE - 1) // W4A8_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)