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

[NPU] Add mxfp4-w4a8 MOE Quantization Support for NPU (#30318)

Browse Source
This commit is contained in:
LinyuanLi
2026-08-16 14:03:17 +08:00
committed by GitHub
parent 24ab8f9ed9
commit 0da87024d3
6 changed files with 202 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/docs/advanced_features/quantization.mdx
+2 -1
View File
@@ -61,7 +61,7 @@ The following table summarizes quantization method support across NVIDIA and AMD
<td>No</td>
<td>No</td>
<td>Yes (A5)</td>
<td>Ascend NPU only; online W4A8 for Qwen3 dense LLM (MXFP4 weights + MXFP8 activations) on A5 series; offline <code>W4A8_MXFP</code> checkpoints are auto-detected via <code>modelslim</code></td>
<td>Ascend NPU only; online W4A8 for Qwen3 dense LLM (MXFP4 weights + MXFP8 activations) on A5 series; offline <code>W4A8_MXFP</code> dense and MoE checkpoints are auto-detected via <code>modelslim</code></td>
</tr>
<tr>
<td><code>blockwise_int8</code></td>
@@ -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
docs/docs/hardware-platforms/ascend-npus/optimization/quantization.mdx
+25
View File
@@ -68,6 +68,14 @@ SGLang supports **mix-bits** quantization (independently defines and loads each
<td><strong style={{color: 'green'}}>√</strong></td>
<td><strong style={{color: 'red'}}>x</strong></td>
</tr>
<tr>
<td><a href="https://github.com/sgl-project/sglang/pull/30318">MXFP4 W4A8 (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><a href="https://github.com/sgl-project/sglang/pull/23795">MXFP4 W4A4</a></td>
<td>Linear</td>
@@ -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).
python/sglang/srt/hardware_backend/npu/quantization/moe_methods.py
+86
View File
@@ -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
# ---------------------------------------------------------------------------
python/sglang/srt/layers/quantization/modelslim/modelslim.py
+4 -2
View File
@@ -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),
python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py
+2
View File
@@ -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",
python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a8_mxfp4_moe.py
+83
View File
@@ -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)