From ed42af99a92fa6f69c5de7f05d023b5a673ddf6b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=90=D1=80=D1=82=D0=B5=D0=BC=20=D0=A1=D0=B0=D0=B2=D0=BA?= =?UTF-8?q?=D0=B8=D0=BD?= <58187114+OrangeRedeng@users.noreply.github.com> Date: Wed, 11 Mar 2026 16:52:35 +0300 Subject: [PATCH] [NPU] [Quantization] w4a4 MoE layer support (#18924) --- docs/platforms/ascend_npu_quantization.md | 1 + .../npu/quantization/fused_moe_method_npu.py | 154 ++++++++++++++++++ .../quantization/modelslim/modelslim.py | 12 +- .../modelslim/schemes/__init__.py | 2 + .../schemes/modelslim_w4a4_int4_moe.py | 135 +++++++++++++++ .../sglang/test/ascend/test_ascend_utils.py | 3 + .../llm_models/test_ascend_qwen3_30b_w4a4.py | 37 +++++ 7 files changed, 342 insertions(+), 2 deletions(-) create mode 100644 python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4_moe.py create mode 100644 test/registered/ascend/llm_models/test_ascend_qwen3_30b_w4a4.py diff --git a/docs/platforms/ascend_npu_quantization.md b/docs/platforms/ascend_npu_quantization.md index 234eb6c34..fb4adb54f 100644 --- a/docs/platforms/ascend_npu_quantization.md +++ b/docs/platforms/ascend_npu_quantization.md @@ -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 diff --git a/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py b/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py index f9f3ce4a7..7214c13d5 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py @@ -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: diff --git a/python/sglang/srt/layers/quantization/modelslim/modelslim.py b/python/sglang/srt/layers/quantization/modelslim/modelslim.py index 65433c57c..e282eb936 100644 --- a/python/sglang/srt/layers/quantization/modelslim/modelslim.py +++ b/python/sglang/srt/layers/quantization/modelslim/modelslim.py @@ -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: diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py b/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py index 14005eb3d..c349fd3c4 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py @@ -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", ] diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4_moe.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4_moe.py new file mode 100644 index 000000000..9faf1b7a4 --- /dev/null +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4_moe.py @@ -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" + ) diff --git a/python/sglang/test/ascend/test_ascend_utils.py b/python/sglang/test/ascend/test_ascend_utils.py index 183f96760..b47c5dc51 100644 --- a/python/sglang/test/ascend/test_ascend_utils.py +++ b/python/sglang/test/ascend/test_ascend_utils.py @@ -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") diff --git a/test/registered/ascend/llm_models/test_ascend_qwen3_30b_w4a4.py b/test/registered/ascend/llm_models/test_ascend_qwen3_30b_w4a4.py new file mode 100644 index 000000000..b31f3c7e2 --- /dev/null +++ b/test/registered/ascend/llm_models/test_ascend_qwen3_30b_w4a4.py @@ -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()