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

1
docs/platforms/ascend_npu_quantization.md
View File

@@ -6,6 +6,7 @@ To load already quantized models, simply load the model weights and config. Agai
- [x] W4A4 dynamic linear
- [x] W8A8 static linear
- [x] W8A8 dynamic linear
- [x] W4A4 dynamic MOE
- [x] W4A8 dynamic MOE
- [x] W8A8 dynamic MOE

154
python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py
View File

@@ -14,6 +14,92 @@ if TYPE_CHECKING:
from sglang.srt.layers.quantization.base_config import QuantizationConfig
def npu_fused_experts_w4a4(
hidden_states: torch.Tensor,
w13: torch.Tensor,
w13_scale: torch.Tensor,
w2: torch.Tensor,
w2_scale: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
top_k: int,
):
original_shape = hidden_states.shape
original_dtype = hidden_states.dtype
scale_dtype = original_dtype if original_dtype == torch.bfloat16 else torch.float32
if len(original_shape) == 3:
hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
num_tokens = hidden_states.shape[0]
num_experts = w13.shape[0]
hidden_states, expanded_row_idx, expert_tokens, _ = (
torch.ops.npu.npu_moe_init_routing_v2(
hidden_states,
topk_ids,
active_num=num_tokens * top_k,
expert_num=num_experts,
expert_tokens_num_type=1,
expert_tokens_num_flag=True,
active_expert_range=[0, num_experts],
quant_mode=-1,
)
)
expert_tokens = expert_tokens.to(torch.int64)
# gmm1: gate_up_proj
hidden_states, pertoken_scale = torch.ops.npu.npu_dynamic_quant(
hidden_states, dst_type=torch.quint4x2
)
scale_args13 = {
"scale": [w13_scale],
"per_token_scale": [pertoken_scale],
}
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[w13],
**scale_args13,
split_item=2,
group_list_type=1,
group_type=0,
group_list=expert_tokens,
output_dtype=original_dtype,
)[0]
# act_fn: swiglu
hidden_states = torch.ops.npu.npu_swiglu(hidden_states)
hidden_states, pertoken_scale = torch.ops.npu.npu_dynamic_quant(hidden_states)
scale_args2 = {
"scale": [w2_scale.to(scale_dtype)],
"per_token_scale": [pertoken_scale],
}
# gmm2: down_proj
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[w2],
**scale_args2,
split_item=2,
group_list_type=1,
group_type=0,
group_list=expert_tokens,
output_dtype=original_dtype,
)[0]
final_hidden_states = torch.ops.npu.npu_moe_finalize_routing(
hidden_states,
skip1=None,
skip2=None,
bias=None,
scales=topk_weights,
expanded_src_to_dst_row=expanded_row_idx,
export_for_source_row=topk_ids,
drop_pad_mode=2,
)
if len(original_shape) == 3:
final_hidden_states = final_hidden_states.view(original_shape)
return final_hidden_states
def npu_fused_experts(
hidden_states: torch.Tensor,
w13: torch.Tensor,
@@ -231,6 +317,74 @@ class _NPUFusedMoEMethodBase(FusedMoEMethodBase):
self.quant_config = quant_config
class NPUW4A4Int4DynamicMoEMethod(_NPUFusedMoEMethodBase):
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
layer.w13_weight.data = npu_format_cast(layer.w13_weight.data.transpose(1, 2))
layer.w13_weight.data = self._pack_to_int32(
layer.w13_weight.data.to(torch.int32)
)
layer.w2_weight.data = npu_format_cast(layer.w2_weight.data.transpose(1, 2))
scale_np = layer.w13_weight_scale.data.cpu().numpy()
scale_np.dtype = np.uint32
scale_uint64_tensor = torch.from_numpy(scale_np.astype(np.int64)).npu()
layer.w13_weight_scale = torch.nn.Parameter(
scale_uint64_tensor.squeeze(-1), requires_grad=False
)
layer.w2_weight_scale = torch.nn.Parameter(
layer.w2_weight_scale.data.squeeze(-1), requires_grad=False
)
# Compressed-tensors format doesn't have this field
if hasattr(layer, "w13_weight_offset"):
layer.w13_weight_offset = torch.nn.Parameter(
layer.w13_weight_offset.data.squeeze(-1),
requires_grad=False,
)
if hasattr(layer, "w2_weight_offset"):
layer.w2_weight_offset = torch.nn.Parameter(
layer.w2_weight_offset.data.squeeze(-1),
requires_grad=False,
)
def _pack_to_int32(self, weight: torch.Tensor):
# pack 8 int4 to int32, we use a int32 to represent a int4
assert (
weight.shape[-1] % 8 == 0
), "the last dim of weight needs to be divided by 8"
new_weight = torch.ops.npu.npu_convert_weight_to_int4pack(weight.flatten(0, 1))
new_weight = new_weight.view(weight.shape[0], weight.shape[1], -1)
return new_weight
def apply(
self,
layer,
dispatch_output: "StandardDispatchOutput",
) -> "CombineInput":
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
x = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output
topk_weights, topk_ids, _ = topk_output
topk_ids = topk_ids.to(torch.int32)
topk_weights = topk_weights.to(x.dtype)
output = npu_fused_experts_w4a4(
hidden_states=x,
w13=layer.w13_weight,
w13_scale=layer.w13_weight_scale,
w2=layer.w2_weight,
w2_scale=layer.w2_weight_scale,
topk_weights=topk_weights,
topk_ids=topk_ids,
top_k=topk_ids.shape[1],
)
return StandardCombineInput(hidden_states=output)
class NPUW8A8Int8DynamicMoEMethod(_NPUFusedMoEMethodBase):
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

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

3
python/sglang/test/ascend/test_ascend_utils.py
View File

@@ -155,6 +155,9 @@ QWEN3_VL_235B_A22B_INSTRUCT_WEIGHTS_PATH = os.path.join(
)
QWEN3_30B_A3B_WEIGHTS_PATH = os.path.join(MODEL_WEIGHTS_DIR, "Qwen/Qwen3-30B-A3B")
QWEN3_30B_MODELSLIM_INT4_WEIGHTS_PATH = os.path.join(
MODEL_WEIGHTS_DIR, "Eco-Tech/Qwen3-30B-A3B-w4a4-LAOS"
)
# Embedding model weights path
BGE_LARGE_EN_V1_5_WEIGHTS_PATH = os.path.join(MODEL_WEIGHTS_DIR, "bge-large-en-v1.5")

37
test/registered/ascend/llm_models/test_ascend_qwen3_30b_w4a4.py Normal file
View File

@@ -0,0 +1,37 @@
import unittest
from sglang.test.ascend.gsm8k_ascend_mixin import GSM8KAscendMixin
from sglang.test.ascend.test_ascend_utils import QWEN3_30B_MODELSLIM_INT4_WEIGHTS_PATH
from sglang.test.ci.ci_register import register_npu_ci
from sglang.test.test_utils import CustomTestCase
register_npu_ci(est_time=400, suite="per-commit-2-npu-a2")
class TestQwen317BGPTQInt8(GSM8KAscendMixin, CustomTestCase):
"""Testcase: Verify that the inference accuracy of the Eco-Tech/Qwen3-30B-A3B-w4a4-LAOS model on the GSM8K dataset is no less than 0.85.
[Test Category] Model
[Test Target] Qwen/Qwen3-1.7B-GPTQ-Int8
"""
model = QWEN3_30B_MODELSLIM_INT4_WEIGHTS_PATH
accuracy = 0.85
other_args = [
"--trust-remote-code",
"--mem-fraction-static",
0.8,
"--max-running-requests",
32,
"--attention-backend",
"ascend",
"--disable-cuda-graph",
"--cuda-graph-max-bs",
32,
"--tp-size",
2,
]
if __name__ == "__main__":
unittest.main()