From 6910c1b281fc696a2912864513c7dd2ebc8576f5 Mon Sep 17 00:00:00 2001 From: YeChang Guo <52730608+YChange01@users.noreply.github.com> Date: Wed, 4 Mar 2026 21:02:19 +0800 Subject: [PATCH] [Feature][NPU]: add runtime support for GPTQ-quantized MoE models (#16364) Co-authored-by: GuoYechang <52730608+GuoYechang@users.noreply.github.com> Co-authored-by: root --- python/sglang/srt/layers/linear.py | 2 + python/sglang/srt/layers/quantization/gptq.py | 296 +++++++++++++++++- test/srt/ascend/test_ascend_gptq_moe.py | 76 +++++ test/srt/run_suite.py | 1 + 4 files changed, 370 insertions(+), 5 deletions(-) create mode 100644 test/srt/ascend/test_ascend_gptq_moe.py diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index 27a8e8695..02d11cc45 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -64,6 +64,8 @@ WEIGHT_LOADER_V2_SUPPORTED = [ "GPTQMarlin24LinearMethod", "TPUInt8LinearMethod", "GPTQLinearMethod", + "GPTQLinearAscendMethod", + "GPTQMoEAscendMethod", "FBGEMMFp8LinearMethod", "GPTQLinearAscendMethod", "ModelOptFp8LinearMethod", diff --git a/python/sglang/srt/layers/quantization/gptq.py b/python/sglang/srt/layers/quantization/gptq.py index baf4c3d94..9e5c2945f 100644 --- a/python/sglang/srt/layers/quantization/gptq.py +++ b/python/sglang/srt/layers/quantization/gptq.py @@ -7,6 +7,9 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union import torch +from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( + npu_fused_experts, +) from sglang.srt.layers.moe import ( MoeRunner, MoeRunnerBackend, @@ -174,8 +177,8 @@ class GPTQConfig(QuantizationConfig): f"group_size={self.group_size}, " f"desc_act={self.desc_act})," f"lm_head_quantized={self.lm_head_quantized}), " - f"dynamic={self.dynamic}", - f"checkpoint_format={self.checkpoint_format})", + f"dynamic={self.dynamic}," + f"checkpoint_format={self.checkpoint_format})" ) def get_scaled_act_names(self) -> List[str]: @@ -236,11 +239,10 @@ class GPTQConfig(QuantizationConfig): from sglang.srt.layers.moe.fused_moe_triton import FusedMoE if _is_npu: + if isinstance(layer, FusedMoE): + return GPTQMoEAscendMethod(self) if isinstance(layer, LinearBase): return GPTQLinearAscendMethod(self) - elif isinstance(layer, FusedMoE): - # TODO: support GPTQ quantization MoE on npu. - raise NotImplementedError("GPTQ Method does not support MoE yet.") return None if isinstance(layer, FusedMoE): @@ -597,6 +599,290 @@ class GPTQLinearMethod(LinearMethodBase): return output.reshape(out_shape) +class GPTQMoEAscendMethod(FusedMoEMethodBase): + + def __init__(self, quant_config: GPTQConfig): + super().__init__() + self.quant_config = quant_config + self.use_v2_format = quant_config.checkpoint_format == "gptq_v2" + self.moe_runner_config: Optional[MoeRunnerConfig] = None + + 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, + ): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + pack_factor = self.quant_config.pack_factor + + num_groups_w13 = hidden_size // self.quant_config.group_size + num_groups_w2 = intermediate_size_per_partition // self.quant_config.group_size + + extra_weight_attrs.update( + { + "is_transposed": True, + "quant_method": FusedMoeWeightScaleSupported.GROUP.value, + } + ) + + w13_qweight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size // pack_factor, + 2 * intermediate_size_per_partition, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_qweight", w13_qweight) + set_weight_attrs(w13_qweight, extra_weight_attrs) + + w2_qweight = torch.nn.Parameter( + torch.empty( + num_experts, + intermediate_size_per_partition // pack_factor, + hidden_size, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_qweight", w2_qweight) + set_weight_attrs(w2_qweight, extra_weight_attrs) + + w13_scales = torch.nn.Parameter( + torch.empty( + num_experts, + num_groups_w13, + 2 * intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_scales", w13_scales) + set_weight_attrs(w13_scales, extra_weight_attrs) + + w2_scales = torch.nn.Parameter( + torch.empty( + num_experts, + num_groups_w2, + hidden_size, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_scales", w2_scales) + set_weight_attrs(w2_scales, extra_weight_attrs) + + w13_qzeros = torch.nn.Parameter( + torch.empty( + num_experts, + num_groups_w13, + 2 * intermediate_size_per_partition // pack_factor, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_qzeros", w13_qzeros) + set_weight_attrs(w13_qzeros, extra_weight_attrs) + + w2_qzeros = torch.nn.Parameter( + torch.empty( + num_experts, + num_groups_w2, + hidden_size // pack_factor, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_qzeros", w2_qzeros) + set_weight_attrs(w2_qzeros, extra_weight_attrs) + + def create_moe_runner( + self, + layer: torch.nn.Module, + moe_runner_config: MoeRunnerConfig, + **extra_weight_attrs, + ): + self.moe_runner_config = moe_runner_config + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + w13_qzeros_2d = layer.w13_qzeros.data.contiguous().reshape( + -1, layer.w13_qzeros.shape[-1] + ) + layer.w13_qzeros = torch.nn.Parameter( + unpack_from_int32( + w13_qzeros_2d, + self.quant_config.weight_bits, + packed_dim=1, + ) + .reshape(layer.w13_qzeros.shape[0], layer.w13_qzeros.shape[1], -1) + .to(layer.w13_scales.dtype), + requires_grad=False, + ) + if not self.use_v2_format: + layer.w13_qzeros += 1 + + w2_qzeros_2d = layer.w2_qzeros.data.contiguous().reshape( + -1, layer.w2_qzeros.shape[-1] + ) + layer.w2_qzeros = torch.nn.Parameter( + unpack_from_int32( + w2_qzeros_2d, + self.quant_config.weight_bits, + packed_dim=1, + ) + .reshape(layer.w2_qzeros.shape[0], layer.w2_qzeros.shape[1], -1) + .to(layer.w2_scales.dtype), + requires_grad=False, + ) + if not self.use_v2_format: + layer.w2_qzeros += 1 + + w13_qweight_2d = ( + layer.w13_qweight.data.transpose(-1, -2) + .contiguous() + .reshape(-1, layer.w13_qweight.shape[-2]) + ) + w13_qweight_tmp = unpack_from_int32( + w13_qweight_2d, self.quant_config.weight_bits, packed_dim=1 + ) + + if self.quant_config.weight_bits == 4: + group_size = self.quant_config.group_size + scale_expanded = layer.w13_scales.data.repeat_interleave(group_size, dim=1) + + neg_mask = scale_expanded < 0 + + if neg_mask.any(): + neg_mask = neg_mask.transpose(-1, -2) + neg_mask = neg_mask.contiguous().reshape(w13_qweight_tmp.shape) + w13_qweight_tmp[neg_mask] = -w13_qweight_tmp[neg_mask] + + if w13_qweight_tmp.max() > 7: + w13_qweight_tmp.clamp_(max=7) + + layer.w13_scales.data.abs_() + + layer.w13_qweight = torch.nn.Parameter( + torch_npu.npu_convert_weight_to_int4pack( + w13_qweight_tmp.reshape( + layer.w13_qweight.shape[0], layer.w13_qweight.shape[2], -1 + ) + .transpose(-1, -2) + .contiguous() + .reshape(-1, layer.w13_qweight.shape[2]) + .to(torch.int32) + ) + .reshape(layer.w13_qweight.shape[0], layer.w13_qweight.shape[1] * 8, -1) + .contiguous(), + requires_grad=False, + ) + # use int8 to store weight by default + else: + layer.w13_qweight = torch.nn.Parameter( + w13_qweight_tmp.reshape( + layer.w13_qweight.shape[0], layer.w13_qweight.shape[2], -1 + ) + .transpose(-1, -2) + .contiguous(), + requires_grad=False, + ) + + w2_qweight_2d = ( + layer.w2_qweight.data.transpose(-1, -2) + .contiguous() + .reshape(-1, layer.w2_qweight.shape[-2]) + ) + w2_qweight_tmp = unpack_from_int32( + w2_qweight_2d, self.quant_config.weight_bits, packed_dim=1 + ) + + if self.quant_config.weight_bits == 4: + group_size = self.quant_config.group_size + scale_expanded = layer.w2_scales.data.repeat_interleave(group_size, dim=1) + + neg_mask = scale_expanded < 0 + + if neg_mask.any(): + neg_mask = neg_mask.transpose(-1, -2) + neg_mask = neg_mask.contiguous().reshape(w2_qweight_tmp.shape) + w2_qweight_tmp[neg_mask] = -w2_qweight_tmp[neg_mask] + + if w2_qweight_tmp.max() > 7: + w2_qweight_tmp.clamp_(max=7) + + layer.w2_scales.data.abs_() + + layer.w2_qweight = torch.nn.Parameter( + torch_npu.npu_convert_weight_to_int4pack( + w2_qweight_tmp.reshape( + layer.w2_qweight.shape[0], layer.w2_qweight.shape[2], -1 + ) + .transpose(-1, -2) + .contiguous() + .reshape(-1, layer.w2_qweight.shape[2]) + .to(torch.int32) + ) + .reshape(layer.w2_qweight.shape[0], layer.w2_qweight.shape[1] * 8, -1) + .contiguous(), + requires_grad=False, + ) + # use int8 to store weight by default + else: + layer.w2_qweight = torch.nn.Parameter( + w2_qweight_tmp.reshape( + layer.w2_qweight.shape[0], layer.w2_qweight.shape[2], -1 + ) + .transpose(-1, -2) + .contiguous(), + requires_grad=False, + ) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> torch.Tensor: + from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + + assert ( + self.moe_runner_config is not None + ), "moe_runner_config is not set. Did you forget to call create_weights/create_moe_runner?" + + assert self.moe_runner_config.activation in ("silu", "swiglu"), ( + f"Only SiLU/Swiglu activation is supported, " + f"got {self.moe_runner_config.activation!r}." + ) + + 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( + hidden_states=x, + w13=layer.w13_qweight, + w13_scale=layer.w13_scales, + w13_offset=layer.w13_qzeros, + w2=layer.w2_qweight, + w2_scale=layer.w2_scales, + w2_offset=layer.w2_qzeros, + topk_weights=topk_weights, + topk_ids=topk_ids, + top_k=topk_ids.shape[1], + use_wna16=True, + ) + + return StandardCombineInput(hidden_states=output) + + class GPTQMarlinLinearMethod(LinearMethodBase): """Linear method for GPTQ Marlin. diff --git a/test/srt/ascend/test_ascend_gptq_moe.py b/test/srt/ascend/test_ascend_gptq_moe.py new file mode 100644 index 000000000..22b954379 --- /dev/null +++ b/test/srt/ascend/test_ascend_gptq_moe.py @@ -0,0 +1,76 @@ +import logging +import unittest +from types import SimpleNamespace +from urllib.parse import urlparse + +from sglang.srt.utils import kill_process_tree +from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +logger = logging.getLogger(__name__) + +TEST_MODEL_MATRIX = { + "/root/.cache/modelscope/hub/models/Qwen/Qwen3-30B-A3B-GPTQ-Int4": { + "accuracy": 0.85, + }, +} + + +class TestAscendGPTQMoEInt4(CustomTestCase): + + @classmethod + def setUpClass(cls): + cls.models = TEST_MODEL_MATRIX.keys() + cls.base_url = DEFAULT_URL_FOR_TEST + cls.url = urlparse(DEFAULT_URL_FOR_TEST) + cls.common_args = [ + "--trust-remote-code", + "--mem-fraction-static", + 0.8, + "--attention-backend", + "ascend", + "--quantization", + "gptq", + ] + + def test_a_gsm8k(self): + for model in self.models: + with self.subTest(model=model): + logger.info(f"##=== Testing accuracy: {model} ===##") + + process = popen_launch_server( + model, + self.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + *self.common_args, + ], + ) + + try: + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=1319, + max_new_tokens=512, + parallel=128, + host=f"http://{self.url.hostname}", + port=int(self.url.port), + ) + + metrics = run_eval_few_shot_gsm8k(args) + self.assertGreaterEqual( + metrics["accuracy"], + TEST_MODEL_MATRIX[model]["accuracy"], + ) + finally: + kill_process_tree(process.pid) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index a5dc7a145..8f8de3885 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -84,6 +84,7 @@ suite_xpu = { suite_ascend = { "per-commit-1-npu-a2": [ TestFile("ascend/test_ascend_gptq.py", 400), + TestFile("ascend/test_ascend_gptq_moe.py", 400), TestFile("ascend/test_ascend_graph_tp1_bf16.py", 400), TestFile("ascend/test_ascend_piecewise_graph_prefill.py", 400), TestFile("ascend/test_ascend_hicache_mha.py", 400),