chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[NPU] [Quantization] w4a4 MoE layer support (#18924)

Browse Source
This commit is contained in:
Артем Савкин
2026-03-11 16:52:35 +03:00
committed by GitHub
parent 9991debde3
commit ed42af99a9
7 changed files with 342 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/sglang/srt/layers/quantization/modelslim/modelslim.py
View File

@@ -16,6 +16,7 @@ from sglang.srt.layers.quantization.base_config import (
from sglang.srt.layers.quantization.compressed_tensors.utils import should_ignore_layer
from sglang.srt.layers.quantization.modelslim.schemes import (
ModelSlimW4A4Int4,
ModelSlimW4A4Int4MoE,
ModelSlimW4A8Int8MoE,
ModelSlimW8A8Int8,
ModelSlimW8A8Int8MoE,
@@ -214,7 +215,11 @@ class ModelSlimConfig(QuantizationConfig):
# TODO: @dsikka: refactor this to use schemes as other kernels
# are supported + check if the layer is being ignored.
prefix_in_quant_config = prefix + ".0.down_proj.weight"
prefix_in_quant_config = prefix + ".0.gate_proj.weight"
is_moe_w4a4_dynamic = (
self.quant_description.get(prefix_in_quant_config, "STATIC")
== "W4A4_DYNAMIC"
)
is_moe_w4a8_dynamic = (
self.quant_description.get(prefix_in_quant_config, "STATIC")
== "W4A8_DYNAMIC"
@@ -223,7 +228,10 @@ class ModelSlimConfig(QuantizationConfig):
self.quant_description.get(prefix_in_quant_config, "STATIC")
== "W8A8_DYNAMIC"
)
if is_moe_w4a8_dynamic:
if is_moe_w4a4_dynamic:
logger.info_once("Using ModelSlimW4A4Int4MoE")
return ModelSlimW4A4Int4MoE(self)
elif is_moe_w4a8_dynamic:
logger.info_once("Using ModelSlimW4A8Int8MoE")
return ModelSlimW4A8Int8MoE(self)
elif is_moe_w8a8_dynamic:

2
python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py
View File

@@ -2,6 +2,7 @@
from .modelslim_scheme import ModelSlimLinearScheme, ModelSlimMoEScheme
from .modelslim_w4a4_int4 import ModelSlimW4A4Int4
from .modelslim_w4a4_int4_moe import ModelSlimW4A4Int4MoE
from .modelslim_w4a8_int8_moe import ModelSlimW4A8Int8MoE
from .modelslim_w8a8_int8 import ModelSlimW8A8Int8
from .modelslim_w8a8_int8_moe import ModelSlimW8A8Int8MoE
@@ -11,6 +12,7 @@ __all__ = [
"ModelSlimMoEScheme",
"ModelSlimW8A8Int8",
"ModelSlimW4A4Int4",
"ModelSlimW4A4Int4MoE",
"ModelSlimW4A8Int8MoE",
"ModelSlimW8A8Int8MoE",
]

135
python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4_moe.py Normal file
View File

@@ -0,0 +1,135 @@
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, Any, Dict
import torch
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
NPUW4A4Int4DynamicMoEMethod,
)
from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimMoEScheme
from sglang.srt.utils import set_weight_attrs
if TYPE_CHECKING:
from sglang.srt.layers.moe import MoeRunnerConfig
from sglang.srt.layers.moe.token_dispatcher import (
CombineInput,
StandardDispatchOutput,
)
logger = logging.getLogger(__name__)
__all__ = [
"ModelSlimW4A4Int4MoE",
]
class ModelSlimW4A4Int4MoE(ModelSlimMoEScheme):
def __init__(
self,
quant_config: Dict[str, Any],
prefix: str = None,
):
self.quant_config = quant_config
self.kernel = NPUW4A4Int4DynamicMoEMethod()
def create_weights(
self,
layer: torch.nn.Module,
num_experts: int,
hidden_size: int,
intermediate_size_per_partition: int,
params_dtype: torch.dtype,
**extra_weight_attrs,
) -> None:
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
self.num_experts = num_experts
extra_weight_attrs.update(
{"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
)
# weight
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size,
dtype=torch.int8,
),
requires_grad=False,
)
layer.register_parameter("w13_weight", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
w2_weight = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition,
dtype=torch.int8,
),
requires_grad=False,
)
layer.register_parameter("w2_weight", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
# scale
w13_weight_scale = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale", w13_weight_scale)
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
w2_weight_scale = torch.nn.Parameter(
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
requires_grad=False,
)
layer.register_parameter("w2_weight_scale", w2_weight_scale)
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
# offset
w13_weight_offset = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_weight_offset", w13_weight_offset)
set_weight_attrs(w13_weight_offset, extra_weight_attrs)
w2_weight_offset = torch.nn.Parameter(
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
requires_grad=False,
)
layer.register_parameter("w2_weight_offset", w2_weight_offset)
set_weight_attrs(w2_weight_offset, extra_weight_attrs)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
self.kernel.process_weights_after_loading(layer)
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig"
):
self.moe_runner_config = moe_runner_config
def apply_weights(
self,
layer,
dispatch_output: "StandardDispatchOutput",
) -> "CombineInput":
return self.kernel.apply(layer, dispatch_output)
def apply_without_routing_weights(
self,
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
):
# FIXME W4A4 MoE does not work with DeepEP
raise NotImplementedError(
f"DeepEP currently does not support quantization in int4, please disable --moe-a2a-backend deepep"
)