minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
LinyuanLi
2026-08-18 19:06:05 +03:00
committed by GitHub
parent 97dedd1ce9
commit 9485c083bb
7 changed files with 199 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/docs/advanced_features/quantization.mdx
+2 -1
View File
@@ -47,7 +47,7 @@ The following table summarizes quantization method support across NVIDIA and AMD
<td>Yes</td>
<td>Yes</td>
<td>Yes (A5)</td>
<td>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 <code>W4A4_MXFP4</code> checkpoints (single-level) are auto-detected via <code>modelslim</code></td>
<td>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 <code>W4A4_MXFP4</code> dense and MoE checkpoints (single-level) are auto-detected via <code>modelslim</code></td>
</tr>
<tr>
<td><code>mxfp8</code></td>
@@ -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
docs/docs/hardware-platforms/ascend-npus/optimization/quantization.mdx
+25
View File
@@ -84,6 +84,14 @@ SGLang supports **mix-bits** quantization (independently defines and loads each
<td><strong style={{color: 'blue'}}>WIP</strong></td>
<td><strong style={{color: 'red'}}>x</strong></td>
</tr>
<tr>
<td><a href="https://github.com/sgl-project/sglang/pull/30319">MXFP4 W4A4 (ModelSlim)</a></td>
<td>MoE</td>
<td><strong style={{color: 'red'}}>x</strong></td>
<td><strong style={{color: 'red'}}>x</strong></td>
<td><strong style={{color: 'green'}}>√</strong></td>
<td><strong style={{color: 'red'}}>x</strong></td>
</tr>
<tr>
<td>W4A4 dynamic</td>
<td>MoE</td>
@@ -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.
python/sglang/srt/hardware_backend/npu/moe/quant.py
+6 -5
View File
@@ -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
+80
View File
@@ -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
+2
View File
@@ -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
+2
View File
@@ -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
+82
View File
@@ -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)