From bbd2a3fe4a267b1e5a2a49792a2273e2e519d881 Mon Sep 17 00:00:00 2001 From: Junlin Wu Date: Fri, 17 Jul 2026 14:06:30 +0800 Subject: [PATCH] :sparkles: [llm][npu][quant] Add W4A4 MXFP4 quantization support for Qwen3 Dense on Ascend NPU (#23795) --- .../docs/advanced_features/quantization.mdx | 4 +- .../ascend-npus/ascend_npu_quantization.mdx | 31 ++ .../npu/quantization/linear_method_npu.py | 275 ++++++++++++++++++ .../srt/layers/quantization/__init__.py | 5 + .../quantization/modelslim/modelslim.py | 2 + .../modelslim/schemes/__init__.py | 2 + .../modelslim/schemes/modelslim_mxfp4.py | 96 ++++++ .../srt/layers/quantization/npu_mxfp4_w4a4.py | 140 +++++++++ 8 files changed, 553 insertions(+), 2 deletions(-) create mode 100644 python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp4.py create mode 100644 python/sglang/srt/layers/quantization/npu_mxfp4_w4a4.py diff --git a/docs_new/docs/advanced_features/quantization.mdx b/docs_new/docs/advanced_features/quantization.mdx index f913a0d3e..fd2f6df24 100644 --- a/docs_new/docs/advanced_features/quantization.mdx +++ b/docs_new/docs/advanced_features/quantization.mdx @@ -46,8 +46,8 @@ The following table summarizes quantization method support across NVIDIA and AMD mxfp4 Yes Yes - WIP - Requires CDNA3/CDNA4 with MXFP support; uses Aiter + 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 mxfp8 diff --git a/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_quantization.mdx b/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_quantization.mdx index f4ea12d64..cd8ab03fe 100644 --- a/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_quantization.mdx +++ b/docs_new/docs/hardware-platforms/ascend-npus/ascend_npu_quantization.mdx @@ -68,6 +68,14 @@ SGLang supports **mix-bits** quantization (independently defines and loads each √ x + + MXFP4 W4A4 + Linear + x + x + WIP + x + W4A4 dynamic MoE @@ -377,6 +385,29 @@ 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). +**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). + +- Online W4A4 quantization (BF16/FP16 weights → dual-level MXFP4 at load time): + +```bash Command +python3 -m sglang.launch_server \ + --model-path Qwen/Qwen3-8B \ + --quantization mxfp4 \ + --device npu --attention-backend ascend \ + --host 0.0.0.0 --port 30000 \ + --mem-fraction-static 0.8 --tp-size 1 +``` + +- Offline W4A4 quantization (msmodelslim pre-quantized weights, `W4A4_MXFP4` scheme; no `--quantization` flag needed — auto-detected from `quant_model_description.json`). + +> **Implementation Notes:** +> - **Online** (`NPUDualLevelMXFP4LinearMethod`) uses **dual-level** MXFP4: both weights and activations are quantized with a fine FP8 (E4M3) L0 block scale plus a coarser L1 scale via `npu_dynamic_dual_level_mx_quant`, and the matmul runs via `npu_dual_level_quant_matmul` (weight in FRACTAL_NZ). Dual-level captures per-block dynamic range far better than a single UE8M0 (power-of-2) scale, which is what made an earlier single-level RTN online path degenerate (greedy decoding could loop without emitting EOS). +> - **Offline** (`ModelSlimMXFP4Scheme` → `NPUSingleLevelMXFP4OfflineLinearMethod`) is **single-level**: msmodelslim's `W4A4_MXFP4` checkpoint ships single-level UE8M0 block scales (block_size = 32), so the matmul runs via `npu_quant_matmul(..., x1_dtype=x2_dtype=torch_npu.float4_e2m1fn_x2, group_sizes=[1, 1, 32])`. The online and offline paths therefore use different matmul kernels — they no longer share the matmul path. +> - 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. + ## 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 A5; the ModelSlim W8A8/W4A4 schemes work on A2/A3. diff --git a/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py b/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py index ffd25114d..166220a97 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py @@ -625,3 +625,278 @@ class NPUMXFP4W4A8OfflineLinearMethod(_NPULinearMethodBase): # Restore original shape (replace last dim with output features). output_shape = list(input_shape[:-1]) + [output.shape[-1]] return output.reshape(output_shape) + + +class NPUSingleLevelMXFP4LinearMethod(_NPULinearMethodBase): + """Ascend NPU W4A4 online quantization: single-level MXFP4. + + True W4(weight) A4(activation): both weights and activations are quantised to + single-level MXFP4 (``float4_e2m1fn_x2``), unlike the W4A8 path which keeps FP8 + activations. All NPU ops go through ``torch.ops.npu.*`` (no top-level + ``torch_npu``) and the fp4 dtype comes from ``_get_float4_e2m1fn_x2_dtype()``. + + Weight quantization (process_weights_after_loading): + BF16/FP16 weight → npu_dynamic_mx_quant(dst=float4_e2m1fn_x2) + → (packed FP4 [out, in//2], UE8M0 block scale) → transpose [in//2, out] + + Inference (apply): + BF16/FP16 activation → npu_dynamic_mx_quant(dst=float4_e2m1fn_x2) (A4) + → npu_quant_matmul(x1_dtype = x2_dtype = float4_e2m1fn_x2, + group_sizes=[1, 1, MXFP4_BLOCK_SIZE]) + + Triggered by ``--quantization mxfp4`` on Ascend NPU. Hardware: Ascend 950 (A5) + with a recent torch_npu exposing ``float4_e2m1fn_x2``. + """ + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes, + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + """Register an unquantized (``params_dtype``) weight placeholder. + + The checkpoint still holds full-precision BF16/FP16 weights: the loader + fills this buffer, then ``process_weights_after_loading`` quantizes it to + MXFP4 in place. + """ + from sglang.srt.layers.parameter import ModelWeightParameter + + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + layer.logical_widths = output_partition_sizes + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.orig_dtype = params_dtype + + # Load weights in original dtype; quantise to MXFP4 in + # process_weights_after_loading. + weight = ModelWeightParameter( + data=torch.empty( + output_size_per_partition, + input_size_per_partition, + dtype=params_dtype, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight", weight) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Online single-level MXFP4 weight quant. All NPU ops go through + # torch.ops.npu.* (no torch_npu); the fp4 dtype comes from the shared + # _get_float4_e2m1fn_x2_dtype() helper (the torch_npu int enum). + fp4_dtype = _get_float4_e2m1fn_x2_dtype() + + weight_fp = layer.weight.data + if weight_fp.dtype not in (torch.float16, torch.bfloat16): + weight_fp = weight_fp.to(torch.bfloat16) + # Move to NPU if needed (cpu offload may have put it on CPU). + if not weight_fp.is_npu: + weight_fp = weight_fp.to(f"npu:{torch.npu.current_device()}") + + # BF16 -> packed FP4 (float4_e2m1fn_x2, [out, in//2]) + UE8M0 block scale. + qw, w_scale = torch.ops.npu.npu_dynamic_mx_quant( + weight_fp, dst_type=fp4_dtype, round_mode="round" + ) + # Pre-transpose the weight to [in//2, out] for npu_quant_matmul; use + # .data= to preserve the non-contiguous transpose view (npu_quant_matmul + # reads strides directly — .contiguous() would reorder data and break + # block-scale alignment). + layer.weight = Parameter(qw, requires_grad=False) + layer.weight.data = layer.weight.data.transpose(0, 1) + + # weight_scale -> [in//64, out, 2] (3D), matching the offline W4A4 path, + # the W4A8 path and vllm-ascend's W4A4_MXFP4 layout. npu_dynamic_mx_quant + # already returns the scale as [out, in//64, 2] (3D) on current builds; + # older builds may return [out, in//32] (2D) — reshape those first so the + # transpose always yields the 3D layout npu_quant_matmul requires. + if w_scale.dim() == 2: + n, k = w_scale.shape + w_scale = w_scale.reshape(n, k // 2, 2) + layer.weight_scale = Parameter(w_scale.transpose(-3, -2), requires_grad=False) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + e8m0_dtype = _get_float8_e8m0fnu_dtype() + fp4_dtype = _get_float4_e2m1fn_x2_dtype() + + original_dtype = x.dtype + if original_dtype not in (torch.float16, torch.bfloat16): + x = x.to(torch.bfloat16) + original_dtype = torch.bfloat16 + + # Flatten to 2D [tokens, hidden] for npu_dynamic_mx_quant. + input_shape = x.shape + x_2d = x.reshape(-1, x.shape[-1]) + + # Dynamic single-level MXFP4 activation quantisation (A4 — FP4). + qx, input_scale = torch.ops.npu.npu_dynamic_mx_quant( + x_2d, dst_type=fp4_dtype, round_mode="round" + ) + + # Single-level MXFP4 matmul (weight & scale already transposed at load + # time): x1_dtype = x2_dtype = fp4, group_sizes=[1, 1, block]. + output = torch.ops.npu.npu_quant_matmul( + qx, + layer.weight, + layer.weight_scale, + scale_dtype=e8m0_dtype, + pertoken_scale=input_scale, + pertoken_scale_dtype=e8m0_dtype, + bias=bias.to(torch.float32) if bias is not None else None, + output_dtype=original_dtype, + x1_dtype=fp4_dtype, + x2_dtype=fp4_dtype, + group_sizes=[1, 1, MXFP4_BLOCK_SIZE], + ) + + # Restore original shape (replace last dim with output features). + output_shape = list(input_shape[:-1]) + [output.shape[-1]] + return output.reshape(output_shape) + + +class NPUSingleLevelMXFP4OfflineLinearMethod(NPUSingleLevelMXFP4LinearMethod): + """Ascend NPU offline W4A4 (ModelSlim ``W4A4_MXFP4``): fp8-container FP4 weights. + + Kernel for the offline ``ModelSlimMXFP4Scheme`` (delegated as ``self.kernel``). + The msmodelslim ``W4A4_MXFP4`` checkpoint stores weights as **fp4-in-fp8 + container** (``float8_e4m3fn`` [out, in], one FP4 value per byte) plus UE8M0 + block scales (``uint8`` [out, in//32]). The weight is re-packed to + ``float4_e2m1fn_x2`` (two FP4 per byte) and the scale reshaped to 3D; it then + shares the online :class:`NPUSingleLevelMXFP4LinearMethod` matmul (``apply``) + exactly — only the weight source differs (msmodelslim checkpoint vs online RTN). + Mirrors vllm-ascend's single-level W4A4 MXFP4 layout. + """ + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Re-pack fp8-container FP4 to float4_e2m1fn_x2 and pre-transpose to match + # the online path's layout. All NPU ops go through torch.ops.npu.* (no + # torch_npu); the fp4 dtype must be the torch_npu enum (helper). + fp4_dtype = _get_float4_e2m1fn_x2_dtype() + + weight = layer.weight.data + if not weight.is_npu: + weight = weight.to(f"npu:{torch.npu.current_device()}") + # fp8 container -> float4_e2m1fn_x2 (2 FP4 per byte): [out, in] -> [out, in//2]. + weight_fp4 = torch.ops.npu.npu_dtype_cast(weight, fp4_dtype) + # Transpose to [in//2, out]; no .contiguous() (preserve the strided view so + # the block-scale mapping stays intact). + layer.weight = Parameter(weight_fp4.transpose(0, 1), requires_grad=False) + + weight_scale = layer.weight_scale.data + if not weight_scale.is_npu: + weight_scale = weight_scale.to(f"npu:{torch.npu.current_device()}") + # npu_quant_matmul with float4_e2m1fn_x2 requires x2Scale to be 3D: + # [out, in/32] -> [out, in/64, 2] -> transpose to [in/64, out, 2]. + n_dim, k_dim = weight_scale.shape + layer.weight_scale = Parameter( + weight_scale.reshape(n_dim, k_dim // 2, 2).transpose(0, 1), + requires_grad=False, + ) + + +class NPUDualLevelMXFP4LinearMethod(NPUSingleLevelMXFP4LinearMethod): + """Ascend NPU W4A4 online quantization: dual-level MXFP4 (higher accuracy). + + This is the sole online ``--quantization mxfp4`` linear path. Instead of a single + UE8M0 (power-of-2) block scale, dual-level MX quant produces a finer L0 (FP8 E4M3) + block scale plus a coarser L1 scale, so per-block dynamic range is captured far + more accurately — this fixed the online-RTN degradation that made single-level + decoding loop (never emitting EOS) under greedy sampling. (The single-level + :class:`NPUSingleLevelMXFP4LinearMethod` is retained only as the offline path's + base — msmodelslim checkpoints ship single-level UE8M0 scales.) + + All NPU ops go through ``torch.ops.npu.*`` (no top-level ``torch_npu``). Only + ``create_weights`` (the BF16/FP16 placeholder) is shared with the single-level + base; weight post-processing and the matmul are fully dual-level. + + Weight quantization (process_weights_after_loading): + BF16/FP16 weight → npu_dynamic_dual_level_mx_quant + → (packed FP4 weight, L0 scale, L1 scale); weight cast to FRACTAL_NZ, + L0 scale transposed to [in//l0_block, out]. + + Inference (apply): + BF16/FP16 activation → npu_dynamic_dual_level_mx_quant (A4, dual-level) + → npu_dual_level_quant_matmul(act, weight, act_l0, w_l0, act_l1, w_l1) + + Reference: Diffusion ``NPUMXFP4DiffusionLinearMethod`` / MindIE-SD + ``W4A4MXFP4DualQuantLinear``. Hardware: Ascend 950 (A5) only — the + ``DualLevelQuantBatchMatmul`` op is unavailable on A2/A3. + """ + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + weight_fp = layer.weight.data + if weight_fp.dtype not in (torch.float16, torch.bfloat16): + weight_fp = weight_fp.to(torch.bfloat16) + # Move to NPU if needed (cpu offload may have put it on CPU). + if not weight_fp.is_npu: + weight_fp = weight_fp.to(f"npu:{torch.npu.current_device()}") + + # Dual-level MXFP4 weight quant: packed FP4 weight + L0 (fine, FP8 E4M3) + # and L1 (coarse) block scales. + qw, w_l0_scale, w_l1_scale = torch.ops.npu.npu_dynamic_dual_level_mx_quant( + weight_fp, smooth_scale=None + ) + + # npu_dual_level_quant_matmul requires the weight (x2) in FRACTAL_NZ. + # View the packed FP4 as int8 first (npu_format_cast takes int dtypes). + qw_nz = npu_format_cast( + qw.view(torch.int8), + NPUACLFormat.ACL_FORMAT_FRACTAL_NZ, + customize_dtype=torch.int8, + ) + + # L0 scale -> [in//l0_block, out] (op returns [out, in//l0_block, 1]). + w_l0_scale = w_l0_scale.squeeze(-1).transpose(0, 1).contiguous() + + layer.weight = Parameter(qw_nz, requires_grad=False) + layer.weight_l0_scale = Parameter(w_l0_scale, requires_grad=False) + layer.weight_l1_scale = Parameter(w_l1_scale, requires_grad=False) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + original_dtype = x.dtype + if original_dtype not in (torch.float16, torch.bfloat16): + x = x.to(torch.bfloat16) + original_dtype = torch.bfloat16 + + # Flatten to 2D [tokens, hidden] for the quant operators. + input_shape = x.shape + x_2d = x.reshape(-1, x.shape[-1]) + + # Dynamic dual-level MXFP4 activation quant (A4): packed FP4 + L0/L1 scales. + qx, act_l0_scale, act_l1_scale = torch.ops.npu.npu_dynamic_dual_level_mx_quant( + x_2d, smooth_scale=None + ) + + # Dual-level matmul. Arg order (act, weight, act_l0, w_l0, act_l1, w_l1); + # the weight is NOT transposed here (unlike the single-level path). + output = torch.ops.npu.npu_dual_level_quant_matmul( + qx, + layer.weight, + act_l0_scale, + layer.weight_l0_scale, + act_l1_scale, + layer.weight_l1_scale, + bias=bias.to(torch.float32) if bias is not None else None, + output_dtype=original_dtype, + ) + + # Restore original shape (replace last dim with output features). + output_shape = list(input_shape[:-1]) + [output.shape[-1]] + return output.reshape(output_shape) diff --git a/python/sglang/srt/layers/quantization/__init__.py b/python/sglang/srt/layers/quantization/__init__.py index 5d66ff35f..26a90ca34 100644 --- a/python/sglang/srt/layers/quantization/__init__.py +++ b/python/sglang/srt/layers/quantization/__init__.py @@ -46,6 +46,7 @@ from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config from sglang.srt.layers.quantization.mxfp4 import Mxfp4Config from sglang.srt.layers.quantization.npu_mxfp4 import Mxfp4W4A8Config +from sglang.srt.layers.quantization.npu_mxfp4_w4a4 import Mxfp4W4A4Config from sglang.srt.layers.quantization.nvfp4_online import NvFp4OnlineConfig from sglang.srt.layers.quantization.petit import PetitNvFp4Config from sglang.srt.layers.quantization.qoq import QoQConfig @@ -117,6 +118,10 @@ if is_npu(): BASE_QUANTIZATION_METHODS.update( { "gptq": GPTQAscendConfig, + # On NPU, `mxfp4` means single-level W4A4 MXFP4 for dense LLM (the + # upstream `Mxfp4Config` OCP-MoE path is only registered on + # cpu/cuda/hip above, so there is no collision here). + "mxfp4": Mxfp4W4A4Config, } ) diff --git a/python/sglang/srt/layers/quantization/modelslim/modelslim.py b/python/sglang/srt/layers/quantization/modelslim/modelslim.py index 7af3f9171..e51d58c31 100644 --- a/python/sglang/srt/layers/quantization/modelslim/modelslim.py +++ b/python/sglang/srt/layers/quantization/modelslim/modelslim.py @@ -16,6 +16,7 @@ from sglang.srt.layers.quantization.base_config import ( QuantizationConfig, ) from sglang.srt.layers.quantization.modelslim.schemes import ( + ModelSlimMXFP4Scheme, ModelSlimMXFP4W4A8Scheme, ModelSlimMXFP8Scheme, ModelSlimW4A4Int4, @@ -208,6 +209,7 @@ class ModelSlimConfig(QuantizationConfig): ("W8A8_DYNAMIC", ModelSlimW8A8Int8), ("W8A8_MXFP8", ModelSlimMXFP8Scheme), ("W4A8_MXFP", ModelSlimMXFP4W4A8Scheme), + ("W4A4_MXFP4", ModelSlimMXFP4Scheme), ] quant_schemes = [self.quant_description.get(prefix + ".weight", "")] diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py b/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py index 1b0233967..5d78a2d43 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py @@ -7,6 +7,7 @@ from .modelslim_scheme import ModelSlimLinearScheme, ModelSlimMoEScheme from .modelslim_mxfp8 import ModelSlimMXFP8Scheme from .modelslim_mxfp4_w4a8 import ModelSlimMXFP4W4A8Scheme +from .modelslim_mxfp4 import ModelSlimMXFP4Scheme # isort: on from .modelslim_w4a4_int4 import ModelSlimW4A4Int4 @@ -20,6 +21,7 @@ __all__ = [ "ModelSlimMoEScheme", "ModelSlimMXFP8Scheme", "ModelSlimMXFP4W4A8Scheme", + "ModelSlimMXFP4Scheme", "ModelSlimW8A8Int8", "ModelSlimW4A4Int4", "ModelSlimW4A4Int4MoE", diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp4.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp4.py new file mode 100644 index 000000000..1dcb4d5e5 --- /dev/null +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_mxfp4.py @@ -0,0 +1,96 @@ +"""ModelSlim W4A4_MXFP4 scheme for pre-quantized weight inference on Ascend NPU (SRT). + +The msmodelslim ``W4A4_MXFP4`` checkpoint stores weights in an **fp8 container**: + + weight: float8_e4m3fn shape [out, in] (one FP4 value per byte) + weight_scale: uint8 (UE8M0) shape [out, in//32] (block scales, group_size=32) + +This is a true W4(weight) A4(activation) scheme: both weights and activations are +single-level MXFP4. Weight post-processing and the matmul are delegated to +``NPUSingleLevelMXFP4OfflineLinearMethod`` (``self.kernel``): the fp8-container FP4 +is re-packed to ``float4_e2m1fn_x2`` (two FP4 per byte) + transposed and the scale +reshaped to 3D, then ``npu_quant_matmul`` runs with ``x1_dtype = x2_dtype = +float4_e2m1fn_x2`` and ``group_sizes=[1, 1, 32]`` — sharing the online +``NPUSingleLevelMXFP4LinearMethod`` matmul exactly (only the weight source differs). + +This differs from ``W4A8_MXFP`` (packed-uint8 FP4 weights + FP8 activations) and +from ``W8A8_MXFP8`` (float8_e4m3fn weights of shape [out, in]). +""" + +from typing import Dict, List, Optional + +import torch + +from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import ( + NPUSingleLevelMXFP4OfflineLinearMethod, +) +from sglang.srt.layers.parameter import GroupQuantScaleParameter, ModelWeightParameter +from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimLinearScheme + +# Fixed by the msmodelslim W4A4_MXFP4 export format (group_size=32). +MXFP4_BLOCK_SIZE = 32 + + +class ModelSlimMXFP4Scheme(ModelSlimLinearScheme): + """W4A4_MXFP4 offline scheme — fp8-container FP4 weights, MXFP4 activations.""" + + def __init__( + self, + quant_config: Optional[Dict[str, any]] = None, + prefix: Optional[str] = None, + ): + # quant_config / prefix accepted to match ModelSlimConfig.get_linear_scheme's + # dispatch signature; W4A4_MXFP4 needs no per-layer config beyond create_weights. + del quant_config, prefix + self.kernel = NPUSingleLevelMXFP4OfflineLinearMethod() + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + weight_loader = extra_weight_attrs.get("weight_loader") + output_size_per_partition = sum(output_partition_sizes) + + # msmodelslim exports weight as float8_e4m3fn, shape [out, in] — one FP4 + # value per byte (fp8 container). The kernel re-packs it to + # float4_e2m1fn_x2 (2 FP4 per byte) in process_weights_after_loading. + 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) + + # UE8M0 block scales: uint8, shape [out, in//32]. + scale_dim = input_size_per_partition // MXFP4_BLOCK_SIZE + weight_scale = GroupQuantScaleParameter( + data=torch.empty( + (output_size_per_partition, scale_dim), + dtype=torch.uint8, + ), + input_dim=1, + output_dim=0, + weight_loader=weight_loader, + ) + layer.register_parameter("weight_scale", weight_scale) + + def process_weights_after_loading(self, layer: torch.nn.Module): + self.kernel.process_weights_after_loading(layer) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return self.kernel.apply(layer, x, bias) diff --git a/python/sglang/srt/layers/quantization/npu_mxfp4_w4a4.py b/python/sglang/srt/layers/quantization/npu_mxfp4_w4a4.py new file mode 100644 index 000000000..ec8f4160b --- /dev/null +++ b/python/sglang/srt/layers/quantization/npu_mxfp4_w4a4.py @@ -0,0 +1,140 @@ +"""MXFP4 W4A4 online quantization config (dual-level MXFP4 weights + activations). + +Triggered by ``--quantization mxfp4`` on the Ascend NPU backend. On CUDA / AMD / +CPU the ``mxfp4`` key resolves to the upstream :class:`Mxfp4Config` (OCP MXFP4 +MoE) instead; the per-device split is done at registration time in +``sglang.srt.layers.quantization.__init__`` (this config is only registered +inside the ``is_npu()`` block, mirroring ``GPTQAscendConfig``). + +Online mode: FP16/BF16 weights are quantised to **dual-level** MXFP4 in +``process_weights_after_loading`` (a finer FP8 E4M3 L0 block scale plus a coarser +L1 scale); activations are dynamically quantised the same way and the matmul runs +via ``npu_dual_level_quant_matmul`` (see :class:`NPUDualLevelMXFP4LinearMethod`). +Dual-level is the sole online path — it captures per-block dynamic range far more +accurately than a single-level UE8M0 scale, avoiding the RTN degradation that made +single-level online decoding loop under greedy sampling. Requires Ascend 950 (A5). + +Offline (msmodelslim ``W4A4_MXFP4``) checkpoints are single-level (the checkpoint +stores UE8M0 scales) and are handled separately by the ``modelslim`` config +(``ModelSlimMXFP4Scheme`` → ``NPUSingleLevelMXFP4OfflineLinearMethod``), not this +class. +""" + +from __future__ import annotations + +import logging +from typing import Dict, List, Optional + +import torch + +from sglang.srt.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.unquant import ( + UnquantizedFusedMoEMethod, + UnquantizedLinearMethod, +) +from sglang.srt.layers.quantization.utils import is_layer_skipped +from sglang.srt.utils import is_npu + +logger = logging.getLogger(__name__) + + +class Mxfp4W4A4Config(QuantizationConfig): + """Single-level MXFP4 W4A4 online quantization config for Ascend NPU. + + True W4(weight) A4(activation): both weights and activations are quantised + to single-level MXFP4 (``float4_e2m1fn_x2``). The device-specific linear + method is selected in ``get_quant_method``; only Ascend NPU is wired up + today (on other devices ``mxfp4`` maps to the upstream ``Mxfp4Config``). + """ + + def __init__( + self, + ignored_layers: Optional[List[str]] = None, + packed_modules_mapping: Optional[Dict[str, str]] = None, + ): + super().__init__() + self.ignored_layers = ignored_layers or [] + self.packed_modules_mapping = packed_modules_mapping or {} + + @classmethod + def get_name(cls) -> str: + return "mxfp4" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + return 0 # NPU bypasses CUDA capability checks + + @classmethod + def get_config_filenames(cls) -> List[str]: + return [] + + @classmethod + def from_config(cls, config: Dict) -> Mxfp4W4A4Config: + ignored_layers = cls.get_from_keys_or( + config, ["ignored_layers", "modules_to_not_convert"], None + ) + if ignored_layers: + normalized: List[str] = [] + for layer in ignored_layers: + base = layer.removeprefix("model.") + normalized.append(base) + normalized.append(f"model.{base}") + ignored_layers = normalized + packed_modules_mapping = ( + cls.get_from_keys_or(config, ["packed_modules_mapping"], {}) or {} + ) + return cls( + ignored_layers=ignored_layers, + packed_modules_mapping=packed_modules_mapping, + ) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + if isinstance(layer, LinearBase): + if is_layer_skipped( + prefix, + self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedLinearMethod() + if is_npu(): + from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import ( + NPUDualLevelMXFP4LinearMethod, + ) + + # Online W4A4 always uses dual-level MXFP4 (finer FP8 L0 scales): + # single-level RTN was too lossy and degenerated under greedy + # decoding. Requires Ascend 950 (A5). The single-level kernel is + # retained only for the offline msmodelslim path. + return NPUDualLevelMXFP4LinearMethod(self) + raise NotImplementedError( + "mxfp4 W4A4 (single-level MXFP4 weights + activations) is currently " + "only implemented for the Ascend NPU backend; no CUDA/other-device " + "kernel exists in this config. Add a device branch here when one lands." + ) + elif isinstance(layer, FusedMoE): + # MoE single-level MXFP4 W4A4 not yet implemented; fall back to unquantised + logger.warning( + "MXFP4 W4A4 quantization is not yet supported for FusedMoE layers " + "(prefix=%s). Falling back to unquantized MoE — MoE weights will " + "run in full precision (BF16/FP16).", + prefix, + ) + return UnquantizedFusedMoEMethod( + layer.use_triton_kernels, layer.use_flashinfer_trtllm_moe + ) + return None + + def get_scaled_act_names(self) -> List[str]: + return []