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
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)