diff --git a/python/sglang/srt/hardware_backend/gpu/quantization/awq_kernels.py b/python/sglang/srt/hardware_backend/gpu/quantization/awq_kernels.py new file mode 100644 index 000000000..59ef742e8 --- /dev/null +++ b/python/sglang/srt/hardware_backend/gpu/quantization/awq_kernels.py @@ -0,0 +1,255 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.layers.moe import MoeRunner +from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo +from sglang.srt.layers.quantization.marlin_utils import ( + apply_awq_marlin_linear, + awq_to_marlin_zero_points, + marlin_make_empty_g_idx, + marlin_make_workspace, + marlin_moe_permute_scales, + marlin_permute_scales, + moe_awq_to_marlin_zero_points, +) +from sglang.srt.layers.quantization.utils import get_scalar_types, replace_parameter +from sglang.srt.utils import is_hip, is_xpu + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) + from sglang.srt.layers.quantization.base_config import QuantizationConfig + +awq_marlin_moe_repack = None +awq_marlin_repack = None + + +def _unsupported_awq_dequantize(*args, **kwargs): + raise RuntimeError("AWQ GPU kernels are unavailable on the current platform.") + + +awq_dequantize = _unsupported_awq_dequantize + +if is_xpu(): + try: + from sgl_kernel import awq_dequantize + except ImportError: + pass +elif is_hip(): + try: + from sglang.srt.layers.quantization.awq.awq_triton import ( + awq_dequantize_triton as awq_dequantize, + ) + except ImportError: + pass +else: + try: + from sglang.jit_kernel.awq_dequantize import awq_dequantize + from sglang.jit_kernel.awq_marlin_repack import ( + awq_marlin_moe_repack, + awq_marlin_repack, + ) + from sglang.srt.utils.custom_op import register_custom_op_from_extern + + awq_dequantize = register_custom_op_from_extern( + awq_dequantize, + fake_impl=lambda qweight, scales, qzeros: qweight.new_empty( + qweight.shape[:-1] + (qweight.shape[-1] * 8,), dtype=scales.dtype + ), + ) + except ImportError: + try: + from sglang.srt.layers.quantization.awq.awq_triton import ( + awq_dequantize_triton as awq_dequantize, + ) + except ImportError: + try: + from sgl_kernel import awq_dequantize + except ImportError: + pass + +_, scalar_types = get_scalar_types() + + +class AWQLinearKernel: + def __init__(self, quant_config: Optional["QuantizationConfig"] = None): + self.quant_config = quant_config + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False) + layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False) + layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + qweight = layer.qweight + scales = layer.scales + qzeros = layer.qzeros + pack_factor = self.quant_config.pack_factor + out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,) + reshaped_x = x.reshape(-1, x.shape[-1]) + out = awq_dequantize(qweight, scales, qzeros) + out = torch.matmul(reshaped_x, out) + + if bias is not None: + out.add_(bias) + return out.reshape(out_shape) + + +class AWQMarlinLinearKernel: + def __init__(self, quant_config: Optional["QuantizationConfig"] = None): + self.quant_config = quant_config + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + device = layer.qweight.device + layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False) + layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False) + layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False) + + layer.workspace = marlin_make_workspace(device) + + marlin_qweight = awq_marlin_repack( + layer.qweight, + size_k=layer.input_size_per_partition, + size_n=layer.output_size_per_partition, + num_bits=self.quant_config.quant_type.size_bits, + ) + replace_parameter(layer, "qweight", marlin_qweight) + + marlin_scales = marlin_permute_scales( + layer.scales, + size_k=layer.input_size_per_partition, + size_n=layer.output_size_per_partition, + group_size=self.quant_config.group_size, + ) + replace_parameter(layer, "scales", marlin_scales) + + marlin_zp = awq_to_marlin_zero_points( + layer.qzeros, + size_k=layer.num_groups, + size_n=layer.output_size_per_partition, + num_bits=self.quant_config.quant_type.size_bits, + ) + replace_parameter(layer, "qzeros", marlin_zp) + + layer.g_idx = marlin_make_empty_g_idx(device) + layer.g_idx_sort_indices = marlin_make_empty_g_idx(device) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return apply_awq_marlin_linear( + input=x, + weight=layer.qweight, + weight_scale=layer.scales, + weight_zp=layer.qzeros, + g_idx=layer.g_idx, + g_idx_sort_indices=layer.g_idx_sort_indices, + workspace=layer.workspace, + quant_type=self.quant_config.quant_type, + output_size_per_partition=layer.output_size_per_partition, + input_size_per_partition=layer.input_size_per_partition, + bias=bias, + ) + + +class AWQMoEKernel: + def __init__(self, quant_config: Optional["QuantizationConfig"] = None): + self.quant_config = quant_config + self.runner: Optional[MoeRunner] = None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + num_experts = layer.w13_qweight.shape[0] + device = layer.w13_qweight.device + + layer.w13_g_idx_sort_indices = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + layer.w2_g_idx_sort_indices = torch.nn.Parameter( + torch.empty((num_experts, 0), dtype=torch.int32, device=device), + requires_grad=False, + ) + + marlin_w13_qweight = awq_marlin_moe_repack( + layer.w13_qweight, + layer.w13_g_idx_sort_indices, + size_k=layer.w13_qweight.shape[1], + size_n=layer.w13_qweight.shape[2] * self.quant_config.pack_factor, + num_bits=self.quant_config.weight_bits, + ) + replace_parameter(layer, "w13_qweight", marlin_w13_qweight) + + marlin_w2_qweight = awq_marlin_moe_repack( + layer.w2_qweight, + layer.w2_g_idx_sort_indices, + size_k=layer.w2_qweight.shape[1], + size_n=layer.w2_qweight.shape[2] * self.quant_config.pack_factor, + num_bits=self.quant_config.weight_bits, + ) + replace_parameter(layer, "w2_qweight", marlin_w2_qweight) + + marlin_w13_scales = marlin_moe_permute_scales( + s=layer.w13_scales, + size_k=layer.intermediate_size_per_partition, + size_n=layer.w13_scales.shape[2], + group_size=self.quant_config.group_size, + ) + replace_parameter(layer, "w13_scales", marlin_w13_scales) + + marlin_w2_scales = marlin_moe_permute_scales( + s=layer.w2_scales, + size_k=layer.intermediate_size_per_partition, + size_n=layer.w2_scales.shape[2], + group_size=self.quant_config.group_size, + ) + replace_parameter(layer, "w2_scales", marlin_w2_scales) + + marlin_w13_zp = moe_awq_to_marlin_zero_points( + layer.w13_qzeros, + size_k=layer.w13_qzeros.shape[1], + size_n=layer.w13_qzeros.shape[2] * self.quant_config.pack_factor, + num_bits=self.quant_config.weight_bits, + ) + replace_parameter(layer, "w13_qzeros", marlin_w13_zp) + + marlin_w2_zp = moe_awq_to_marlin_zero_points( + layer.w2_qzeros, + size_k=layer.w2_qzeros.shape[1], + size_n=layer.w2_qzeros.shape[2] * self.quant_config.pack_factor, + num_bits=self.quant_config.weight_bits, + ) + replace_parameter(layer, "w2_qzeros", marlin_w2_zp) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: "StandardDispatchOutput", + ) -> "CombineInput": + if self.runner is None: + raise RuntimeError("moe runner is not initialized") + + quant_info = MarlinMoeQuantInfo( + w13_qweight=layer.w13_qweight, + w2_qweight=layer.w2_qweight, + w13_scales=layer.w13_scales, + w2_scales=layer.w2_scales, + w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices, + w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices, + w13_qzeros=layer.w13_qzeros, + w2_qzeros=layer.w2_qzeros, + weight_bits=self.quant_config.weight_bits, + ) + return self.runner.run(dispatch_output, quant_info) diff --git a/python/sglang/srt/hardware_backend/npu/quantization/awq_kernels.py b/python/sglang/srt/hardware_backend/npu/quantization/awq_kernels.py new file mode 100644 index 000000000..166a5441a --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/quantization/awq_kernels.py @@ -0,0 +1,156 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( + NPUW4A16Int4DynamicMoEMethod, +) +from sglang.srt.layers.quantization.utils import replace_parameter + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput + from sglang.srt.layers.quantization.base_config import QuantizationConfig + +import torch_npu + + +class AWQAscendLinearKernel: + def __init__(self, quant_config: Optional["QuantizationConfig"] = None): + self.quant_config = quant_config + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False) + qweight_tmp = torch.zeros_like(layer.qweight.data) + qzeros_tmp = layer.qzeros.data + qzeros_list = [] + shifts = [0, 4, 1, 5, 2, 6, 3, 7] + + for i in range(0, self.quant_config.pack_factor): + shift_num = shifts[i] * 4 + qzeros_list.append((qzeros_tmp.reshape(-1, 1) >> shift_num) & 0xF) + qweight_tmp.bitwise_or_( + ((layer.qweight.data >> shift_num) & 0xF) << (4 * i) + ) + + qweight_tmp.bitwise_xor_(0x88888888) + + qzeros_tmp = torch.cat(qzeros_list, dim=-1).reshape(qzeros_tmp.shape[0], -1) + qzeros_tmp = -(qzeros_tmp - 8) + qzeros_tmp = qzeros_tmp.to(layer.scales.data.dtype) + + layer.zeros = torch.nn.Parameter(qzeros_tmp, requires_grad=False) + layer.weight = torch.nn.Parameter(qweight_tmp, requires_grad=False) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + qweight = layer.weight + scales = layer.scales + qzeros = layer.zeros + pack_factor = self.quant_config.pack_factor + out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,) + reshaped_x = x.reshape(-1, x.shape[-1]) + + if bias is not None and bias.dtype == torch.bfloat16: + bias = bias.float() + + out = torch_npu.npu_weight_quant_batchmatmul( + reshaped_x, + qweight, + antiquant_scale=scales, + antiquant_offset=qzeros, + antiquant_group_size=self.quant_config.group_size, + bias=bias, + ) + + return out.reshape(out_shape) + + +class AWQAscendMoEKernel: + def __init__(self, quant_config: Optional["QuantizationConfig"] = None): + self.quant_config = quant_config + self.kernel = NPUW4A16Int4DynamicMoEMethod() + + @staticmethod + def _register_or_replace_parameter( + layer: torch.nn.Module, name: str, tensor: torch.Tensor + ) -> None: + if hasattr(layer, name): + replace_parameter(layer, name, tensor) + else: + layer.register_parameter( + name, torch.nn.Parameter(tensor, requires_grad=False) + ) + + def _convert_awq_weight_to_npu_layout(self, qweight: torch.Tensor) -> torch.Tensor: + num_experts, input_size, _ = qweight.shape + unpacked_weight = ( + self.kernel._unpack_from_int32(qweight.flatten(0, 1), 4) + .view(num_experts, input_size, -1) + .transpose(1, 2) + .contiguous() + .int() + ) + return self.kernel._pack_to_int32(unpacked_weight) + + def _convert_awq_qzeros_to_npu_offset( + self, qzeros: torch.Tensor, dtype: torch.dtype + ) -> torch.Tensor: + num_experts, num_groups, _ = qzeros.shape + offset = ( + -self.kernel._unpack_from_int32(qzeros.flatten(0, 1), 4) + .view(num_experts, num_groups, -1) + .transpose(1, 2) + .contiguous() + ) + return offset.to(dtype) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + self._register_or_replace_parameter( + layer, + "w13_weight", + self._convert_awq_weight_to_npu_layout(layer.w13_qweight.data), + ) + self._register_or_replace_parameter( + layer, + "w2_weight", + self._convert_awq_weight_to_npu_layout(layer.w2_qweight.data), + ) + self._register_or_replace_parameter( + layer, + "w13_weight_scale", + layer.w13_scales.data.transpose(1, 2).contiguous(), + ) + self._register_or_replace_parameter( + layer, + "w2_weight_scale", + layer.w2_scales.data.transpose(1, 2).contiguous(), + ) + self._register_or_replace_parameter( + layer, + "w13_weight_offset", + self._convert_awq_qzeros_to_npu_offset( + layer.w13_qzeros.data, layer.w13_scales.data.dtype + ), + ) + self._register_or_replace_parameter( + layer, + "w2_weight_offset", + self._convert_awq_qzeros_to_npu_offset( + layer.w2_qzeros.data, layer.w2_scales.data.dtype + ), + ) + + self.kernel.process_weights_after_loading(layer) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: "StandardDispatchOutput", + ) -> torch.Tensor: + return self.kernel.apply(layer, dispatch_output) diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index a8e7bafda..477a99994 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -54,10 +54,7 @@ logger = logging.getLogger(__name__) WEIGHT_LOADER_V2_SUPPORTED = [ "CompressedTensorsLinearMethod", - "AWQMarlinLinearMethod", "AWQLinearMethod", - "AWQLinearAscendMethod", - "AWQLinearIntelAMXMethod", "GPTQMarlinLinearMethod", "Fp8LinearMethod", "BlockInt8LinearMethod", diff --git a/python/sglang/srt/layers/quantization/__init__.py b/python/sglang/srt/layers/quantization/__init__.py index 3538fdc82..b63db577f 100644 --- a/python/sglang/srt/layers/quantization/__init__.py +++ b/python/sglang/srt/layers/quantization/__init__.py @@ -17,8 +17,7 @@ class DummyConfig: CompressedTensorsConfig = DummyConfig from sglang.srt.layers.quantization.auto_round import AutoRoundConfig -from sglang.srt.layers.quantization.awq import AWQConfig, AWQMarlinConfig -from sglang.srt.layers.quantization.awq_cpu import CPUAWQConfig +from sglang.srt.layers.quantization.awq import AWQConfig, AWQCPUConfig, AWQMarlinConfig from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.bitsandbytes import BitsAndBytesConfig from sglang.srt.layers.quantization.blockwise_int8 import BlockInt8Config @@ -100,7 +99,7 @@ CPU_QUANTIZATION_METHODS = { "fp8": Fp8Config, "w8a8_int8": W8A8Int8Config, "compressed-tensors": CompressedTensorsConfig, - "awq": CPUAWQConfig, + "awq": AWQCPUConfig, "gptq": CPUGPTQConfig, } diff --git a/python/sglang/srt/layers/quantization/auto_round.py b/python/sglang/srt/layers/quantization/auto_round.py index 4690bbb9e..893a0045d 100644 --- a/python/sglang/srt/layers/quantization/auto_round.py +++ b/python/sglang/srt/layers/quantization/auto_round.py @@ -258,8 +258,8 @@ class AutoRoundConfig(QuantizationConfig): use_marlin = False if use_marlin: from sglang.srt.layers.quantization.awq import ( + AWQLinearMethod, AWQMarlinConfig, - AWQMarlinLinearMethod, AWQMoEMethod, ) @@ -282,6 +282,7 @@ class AutoRoundConfig(QuantizationConfig): if isinstance(layer, FusedMoE): if use_marlin: + layer.scheme = quant_args_marlin.get_moe_scheme(layer) return AWQMoEMethod(quant_args_marlin) from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config @@ -296,8 +297,10 @@ class AutoRoundConfig(QuantizationConfig): if isinstance(layer, (LinearBase, ParallelLMHead)): if use_marlin: - return AWQMarlinLinearMethod(quant_args_marlin) + layer.scheme = quant_args_marlin.get_linear_scheme(layer) + return AWQLinearMethod(quant_args_marlin) else: + layer.scheme = quant_args.get_linear_scheme(layer) return AWQLinearMethod(quant_args) return None diff --git a/python/sglang/srt/layers/quantization/awq.py b/python/sglang/srt/layers/quantization/awq.py deleted file mode 100755 index c05aa091a..000000000 --- a/python/sglang/srt/layers/quantization/awq.py +++ /dev/null @@ -1,966 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -import logging -import warnings -from typing import TYPE_CHECKING, Any, Dict, List, Optional - -import torch - -from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( - npu_fused_experts, -) -from sglang.srt.layers.linear import LinearBase, set_weight_attrs -from sglang.srt.layers.moe import ( - MoeRunner, - MoeRunnerBackend, - MoeRunnerConfig, - get_moe_runner_backend, -) -from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo -from sglang.srt.layers.parameter import GroupQuantScaleParameter, PackedvLLMParameter -from sglang.srt.layers.quantization.base_config import ( - FusedMoEMethodBase, - LinearMethodBase, - QuantizationConfig, - QuantizeMethodBase, -) -from sglang.srt.layers.quantization.marlin_utils import ( - apply_awq_marlin_linear, - awq_to_marlin_zero_points, - check_marlin_supported, - check_marlin_supports_layer, - check_moe_marlin_supports_layer, - marlin_make_empty_g_idx, - marlin_make_workspace, - marlin_moe_permute_scales, - marlin_permute_scales, - moe_awq_to_marlin_zero_points, - verify_marlin_supported, - verify_marlin_supports_shape, -) -from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod -from sglang.srt.layers.quantization.utils import get_scalar_types, replace_parameter -from sglang.srt.utils.patch_torch import register_fake_if_exists - -if TYPE_CHECKING: - from sglang.srt.layers.moe.token_dispatcher import ( - CombineInput, - StandardDispatchOutput, - ) - -from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu - -_is_cuda = is_cuda() -_is_hip = is_hip() -_is_xpu = is_xpu() -_is_npu = is_npu() - -if _is_npu: - import torch_npu - -if _is_cuda: - from sglang.jit_kernel.awq_dequantize import awq_dequantize - from sglang.jit_kernel.awq_marlin_repack import ( - awq_marlin_moe_repack, - awq_marlin_repack, - ) - from sglang.srt.utils.custom_op import register_custom_op_from_extern - - awq_dequantize = register_custom_op_from_extern( - awq_dequantize, - fake_impl=lambda qweight, scales, qzeros: qweight.new_empty( - qweight.shape[:-1] + (qweight.shape[-1] * 8,), dtype=scales.dtype - ), - ) - -elif _is_hip: - from sglang.srt.layers.quantization.awq_triton import ( - awq_dequantize_triton as awq_dequantize, - ) - -elif _is_xpu: - from sgl_kernel import awq_dequantize - - warnings.warn(f"XPU does not support fused_marlin_moe currently.") -else: - warnings.warn(f"Only CUDA, HIP and XPU support AWQ currently.") - -logger = logging.getLogger(__name__) - - -ScalarType, scalar_types = get_scalar_types() - - -def is_layer_skipped_awq(prefix: str, modules_to_not_convert: List[str]): - return any(module_name in prefix for module_name in modules_to_not_convert) - - -class AWQConfig(QuantizationConfig): - """Config class for AWQ. - - Reference: https://arxiv.org/abs/2306.00978 - """ - - def __init__( - self, - weight_bits: int, - group_size: int, - zero_point: bool, - modules_to_not_convert: Optional[List[str]] = None, - ) -> None: - super().__init__() - self.weight_bits = weight_bits - self.group_size = group_size - self.zero_point = zero_point - self.modules_to_not_convert = modules_to_not_convert or [] - - if self.weight_bits != 4: - raise ValueError( - "Currently, only 4-bit weight quantization is supported for " - f"AWQ, but got {self.weight_bits} bits." - ) - self.pack_factor = 32 // self.weight_bits - - def __repr__(self) -> str: - return ( - f"AWQConfig(weight_bits={self.weight_bits}, " - f"group_size={self.group_size}, " - f"zero_point={self.zero_point}, " - f"modules_to_not_convert={self.modules_to_not_convert})" - ) - - def get_scaled_act_names(self) -> List[str]: - return [] - - def get_name(self) -> str: - return "awq" - - def get_supported_act_dtypes(self) -> List[torch.dtype]: - return [torch.float16] if not _is_npu else [torch.float16, torch.bfloat16] - - @classmethod - def get_min_capability(cls) -> int: - # The AWQ kernel only supports Turing or newer GPUs. - if _is_npu: - raise NotImplementedError( - 'NPU hardware does not support "get_min_capability" feature.' - ) - else: - return 75 - - @staticmethod - def get_config_filenames() -> List[str]: - return [ - "quant_config.json", # E.g., casperhansen/vicuna-7b-v1.5-awq - # E.g., abhinavkulkarni/mosaicml-mpt-7b-instruct-w4-g128-awq - "quantize_config.json", - ] - - @classmethod - def from_config(cls, config: Dict[str, Any]) -> AWQConfig: - weight_bits = cls.get_from_keys(config, ["w_bit", "bits"]) - group_size = cls.get_from_keys(config, ["q_group_size", "group_size"]) - zero_point = cls.get_from_keys(config, ["zero_point"]) - modules_to_not_convert = cls.get_from_keys_or( - config, ["modules_to_not_convert"], None - ) - return cls(weight_bits, group_size, zero_point, modules_to_not_convert) - - def get_quant_method( - self, layer: torch.nn.Module, prefix: str - ) -> Optional[LinearMethodBase]: - from sglang.srt.layers.linear import LinearBase - from sglang.srt.layers.moe.fused_moe_triton import FusedMoE - - if _is_npu: - if isinstance(layer, LinearBase): - if is_layer_skipped_awq(prefix, self.modules_to_not_convert): - return UnquantizedLinearMethod() - return AWQLinearAscendMethod(self) - elif isinstance(layer, FusedMoE): - return AWQMoEAscendMethod(self) - return None - - if isinstance(layer, LinearBase): - if is_layer_skipped_awq(prefix, self.modules_to_not_convert): - return UnquantizedLinearMethod() - return AWQLinearMethod(self) - return None - - -class AWQMarlinConfig(QuantizationConfig): - """Config class for AWQ Marlin""" - - # num_bits -> type - TYPE_MAP = { - 4: scalar_types.uint4, - 8: scalar_types.uint8, - } - - def __init__( - self, - weight_bits: int, - group_size: int, - zero_point: bool, - lm_head_quantized: bool, - modules_to_not_convert: Optional[list[str]], - full_config: dict[str, Any], - ) -> None: - super().__init__() - if _is_hip: - warnings.warn(f"HIP does not support fused_marlin_moe currently.") - self.pack_factor = 32 // weight_bits # packed into int32 - self.group_size = group_size - self.zero_point = zero_point - self.lm_head_quantized = lm_head_quantized - self.weight_bits = weight_bits - self.modules_to_not_convert = modules_to_not_convert or [] - self.full_config = full_config - - if self.weight_bits not in self.TYPE_MAP: - raise ValueError( - f"Unsupported num_bits = {self.weight_bits}. " - f"Supported num_bits = {self.TYPE_MAP.keys()}" - ) - - self.quant_type = self.TYPE_MAP[self.weight_bits] - - verify_marlin_supported( - self.quant_type, group_size=self.group_size, has_zp=self.zero_point - ) - - def __repr__(self) -> str: - return ( - f"AWQMarlinConfig(quant_type={self.quant_type}, " - f"group_size={self.group_size}, " - f"zero_point={self.zero_point}, " - f"lm_head_quantized={self.lm_head_quantized}, " - f"modules_to_not_convert={self.modules_to_not_convert})" - ) - - def get_scaled_act_names(self) -> List[str]: - return [] - - @classmethod - def get_name(cls) -> str: - return "awq_marlin" - - @classmethod - def get_supported_act_dtypes(cls) -> list[torch.dtype]: - return [torch.half, torch.bfloat16] - - @classmethod - def get_min_capability(cls) -> int: - return 80 - - @classmethod - def get_config_filenames(cls) -> list[str]: - return ["quantize_config.json"] - - @classmethod - def from_config(cls, config: dict[str, Any]) -> AWQMarlinConfig: - weight_bits = cls.get_from_keys(config, ["bits"]) - group_size = cls.get_from_keys(config, ["group_size"]) - zero_point = cls.get_from_keys(config, ["zero_point"]) - lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False) - modules_to_not_convert = cls.get_from_keys_or( - config, ["modules_to_not_convert"], None - ) - return cls( - weight_bits, - group_size, - zero_point, - lm_head_quantized, - modules_to_not_convert, - config, - ) - - @classmethod - def override_quantization_method(cls, hf_quant_cfg, user_quant) -> Optional[str]: - can_convert = cls.is_awq_marlin_compatible(hf_quant_cfg) - is_valid_user_quant = ( - user_quant is None or user_quant == "marlin" or user_quant == "awq_marlin" - ) - - if can_convert and is_valid_user_quant: - msg = ( - "The model is convertible to {} during runtime." - " Using {} kernel.".format(cls.get_name(), cls.get_name()) - ) - logger.info(msg) - return cls.get_name() - - if can_convert and user_quant == "awq": - logger.info( - "Detected that the model can run with awq_marlin" - ", however you specified quantization=awq explicitly," - " so forcing awq. Use quantization=awq_marlin for" - " faster inference" - ) - return None - - def get_quant_method( - self, layer: torch.nn.Module, prefix: str - ) -> Optional[QuantizeMethodBase]: - from sglang.srt.layers.moe.fused_moe_triton import FusedMoE - from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead - - if isinstance(layer, LinearBase) or ( - isinstance(layer, ParallelLMHead) and self.lm_head_quantized - ): - if is_layer_skipped_awq(prefix, self.modules_to_not_convert): - return UnquantizedLinearMethod() - # Check if the layer is supported by AWQMarlin. - if not check_marlin_supports_layer(layer, self.group_size): - logger.warning_once( - "Layer '%s' is not supported by AWQMarlin. Falling back to unoptimized AWQ kernels.", # noqa: E501 - prefix, - ) - return AWQConfig.from_config(self.full_config).get_quant_method( - layer, prefix - ) - return AWQMarlinLinearMethod(self) - elif isinstance(layer, FusedMoE): - from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config - - if not check_moe_marlin_supports_layer(layer, self.group_size): - logger.warning_once( - f"Layer '{prefix}' is not supported by AWQMoeMarlin. " - "Falling back to Moe WNA16 kernels." - ) - return MoeWNA16Config.from_config(self.full_config).get_quant_method( - layer, prefix - ) - return AWQMoEMethod(self) - return None - - @classmethod - def is_awq_marlin_compatible(cls, quant_config: dict[str, Any]): - # Extract data from quant config. - quant_method = quant_config.get("quant_method", "").lower() - num_bits = quant_config.get("bits") - group_size = quant_config.get("group_size") - zero_point = quant_config.get("zero_point") - - if not _is_cuda: - return False - - if quant_method != "awq": - return False - - # If we cannot find the info needed in the config, cannot convert. - if num_bits is None or group_size is None or zero_point is None: - return False - - if num_bits not in cls.TYPE_MAP: - return False - - return check_marlin_supported( - quant_type=cls.TYPE_MAP[num_bits], group_size=group_size, has_zp=zero_point - ) - - -class AWQLinearMethod(LinearMethodBase): - """Linear method for AWQ. - - Args: - quant_config: The AWQ quantization config. - """ - - def __init__(self, quant_config: AWQConfig): - self.quant_config = quant_config - - def create_weights( - self, - layer: torch.nn.Module, - input_size_per_partition: int, - output_partition_sizes: List[int], - input_size: int, - output_size: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - if input_size_per_partition % self.quant_config.group_size != 0: - raise ValueError( - "The input size is not aligned with the quantized " - "weight shape. This can be caused by too large " - "tensor parallel size." - ) - - output_size_per_partition = sum(output_partition_sizes) - if output_size_per_partition % self.quant_config.pack_factor != 0: - raise ValueError( - "The output size is not aligned with the quantized " - "weight shape. This can be caused by too large " - "tensor parallel size." - ) - - weight_loader = extra_weight_attrs.get("weight_loader") - qweight = PackedvLLMParameter( - data=torch.empty( - input_size_per_partition, - output_size_per_partition // self.quant_config.pack_factor, - dtype=torch.int32, - ), - input_dim=0, - output_dim=1, - packed_dim=1, - packed_factor=self.quant_config.pack_factor, - weight_loader=weight_loader, - ) - - qzeros = PackedvLLMParameter( - data=torch.empty( - input_size_per_partition // self.quant_config.group_size, - output_size_per_partition // self.quant_config.pack_factor, - dtype=torch.int32, - ), - input_dim=0, - output_dim=1, - packed_dim=1, - packed_factor=self.quant_config.pack_factor, - weight_loader=weight_loader, - ) - - scales = GroupQuantScaleParameter( - data=torch.empty( - input_size_per_partition // self.quant_config.group_size, - output_size_per_partition, - dtype=params_dtype, - ), - input_dim=0, - output_dim=1, - weight_loader=weight_loader, - ) - - layer.register_parameter("qweight", qweight) - layer.register_parameter("qzeros", qzeros) - layer.register_parameter("scales", scales) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False) - layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False) - layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False) - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - qweight = layer.qweight - scales = layer.scales - qzeros = layer.qzeros - pack_factor = self.quant_config.pack_factor - out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,) - reshaped_x = x.reshape(-1, x.shape[-1]) - out = awq_dequantize(qweight, scales, qzeros) - out = torch.matmul(reshaped_x, out) - - if bias is not None: - out.add_(bias) - return out.reshape(out_shape) - - -class AWQMarlinLinearMethod(LinearMethodBase): - """Linear method for AWQ Marlin. - - Args: - quant_config: The AWQ Marlin quantization config. - """ - - def __init__(self, quant_config: AWQMarlinConfig) -> None: - self.quant_config = quant_config - - def create_weights( - self, - layer: torch.nn.Module, - input_size_per_partition: int, - output_partition_sizes: list[int], - input_size: int, - output_size: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ) -> None: - del output_size - output_size_per_partition = sum(output_partition_sizes) - weight_loader = extra_weight_attrs.get("weight_loader") - - # Normalize group_size - if self.quant_config.group_size != -1: - group_size = self.quant_config.group_size - else: - group_size = input_size - - verify_marlin_supports_shape( - output_size_per_partition=output_size_per_partition, - input_size_per_partition=input_size_per_partition, - input_size=input_size, - group_size=group_size, - ) - - qweight = PackedvLLMParameter( - data=torch.empty( - input_size_per_partition, - output_size_per_partition // self.quant_config.pack_factor, - dtype=torch.int32, - ), - input_dim=0, - output_dim=1, - packed_dim=1, - packed_factor=self.quant_config.pack_factor, - weight_loader=weight_loader, - ) - - num_groups = input_size_per_partition // group_size - - qzeros = PackedvLLMParameter( - data=torch.empty( - num_groups, - output_size_per_partition // self.quant_config.pack_factor, - dtype=torch.int32, - ), - input_dim=0, - output_dim=1, - packed_dim=1, - packed_factor=self.quant_config.pack_factor, - weight_loader=weight_loader, - ) - - scales = GroupQuantScaleParameter( - data=torch.empty( - num_groups, - output_size_per_partition, - dtype=params_dtype, - ), - input_dim=0, - output_dim=1, - weight_loader=weight_loader, - ) - - layer.register_parameter("qweight", qweight) - layer.register_parameter("qzeros", qzeros) - layer.register_parameter("scales", scales) - - layer.input_size_per_partition = input_size_per_partition - layer.output_size_per_partition = output_size_per_partition - layer.num_groups = num_groups - - # TODO: Update this docs - # Checkpoints are serialized in AutoAWQ format, which is different from the - # marlin format. This function is called after the weights are loaded. - # Here, we handle the repacking - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - device = layer.qweight.device - layer.qweight = torch.nn.Parameter(layer.qweight.data, requires_grad=False) - layer.qzeros = torch.nn.Parameter(layer.qzeros.data, requires_grad=False) - layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False) - - # Allocate marlin workspace - layer.workspace = marlin_make_workspace(device) - - # Repack weights from AWQ format to marlin format. - marlin_qweight = awq_marlin_repack( - layer.qweight, - size_k=layer.input_size_per_partition, - size_n=layer.output_size_per_partition, - num_bits=self.quant_config.quant_type.size_bits, - ) - replace_parameter(layer, "qweight", marlin_qweight) - - # Permute scales from AWQ format to marlin format. - marlin_scales = marlin_permute_scales( - layer.scales, - size_k=layer.input_size_per_partition, - size_n=layer.output_size_per_partition, - group_size=self.quant_config.group_size, - ) - replace_parameter(layer, "scales", marlin_scales) - - # Permute zero-points from AWQ format to marlin format. - marlin_zp = awq_to_marlin_zero_points( - layer.qzeros, - size_k=layer.num_groups, - size_n=layer.output_size_per_partition, - num_bits=self.quant_config.quant_type.size_bits, - ) - replace_parameter(layer, "qzeros", marlin_zp) - - # Not-used - layer.g_idx = marlin_make_empty_g_idx(device) - layer.g_idx_sort_indices = marlin_make_empty_g_idx(device) - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - return apply_awq_marlin_linear( - input=x, - weight=layer.qweight, - weight_scale=layer.scales, - weight_zp=layer.qzeros, - g_idx=layer.g_idx, - g_idx_sort_indices=layer.g_idx_sort_indices, - workspace=layer.workspace, - quant_type=self.quant_config.quant_type, - output_size_per_partition=layer.output_size_per_partition, - input_size_per_partition=layer.input_size_per_partition, - bias=bias, - ) - - -class AWQLinearAscendMethod(AWQLinearMethod): - """Linear method for AWQ on Ascend. - - Args: - quant_config: The AWQ quantization config. - """ - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False) - qweight_tmp = torch.zeros_like(layer.qweight.data) - qzeros_tmp = layer.qzeros.data - qzeros_list = [] - shifts = [0, 4, 1, 5, 2, 6, 3, 7] - - for i in range(0, self.quant_config.pack_factor): - shift_num = shifts[i] * 4 - qzeros_list.append((qzeros_tmp.reshape(-1, 1) >> shift_num) & 0xF) - qweight_tmp.bitwise_or_( - ((layer.qweight.data >> shift_num) * (2 ** (4 * i))) & (0xF << (4 * i)) - ) - - qweight_tmp.bitwise_xor_(0x88888888) - - qzeros_tmp = torch.cat(qzeros_list, dim=-1).reshape(qzeros_tmp.shape[0], -1) - qzeros_tmp = -(qzeros_tmp - 8) - qzeros_tmp = qzeros_tmp.to(layer.scales.data.dtype) - - layer.zeros = torch.nn.Parameter(qzeros_tmp, requires_grad=False) - layer.weight = torch.nn.Parameter(qweight_tmp, requires_grad=False) - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - qweight = layer.weight - scales = layer.scales - qzeros = layer.zeros - pack_factor = self.quant_config.pack_factor - out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,) - reshaped_x = x.reshape(-1, x.shape[-1]) - - if bias is not None and bias.dtype == torch.bfloat16: - bias = bias.float() - - out = torch_npu.npu_weight_quant_batchmatmul( - reshaped_x, - qweight, - antiquant_scale=scales, - antiquant_offset=qzeros, - antiquant_group_size=self.quant_config.group_size, - bias=bias, - ) - - return out.reshape(out_shape) - - -class AWQMoEMethod(FusedMoEMethodBase): - - def __init__(self, quant_config: AWQMarlinConfig): - self.quant_config = quant_config - if self.quant_config.weight_bits != 4: - raise ValueError("AWQMoEMethod only supports 4bit now.") - self.quant_type = scalar_types.uint4 - - 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, - ): - # Delay the import to avoid circular dependency - from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported - - extra_weight_attrs.update( - { - "is_transposed": True, - "quant_method": FusedMoeWeightScaleSupported.GROUP.value, - } - ) - - w13_qweight = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - 2 * intermediate_size_per_partition // self.quant_config.pack_factor, - 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, - hidden_size // self.quant_config.pack_factor, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_qweight", w2_qweight) - set_weight_attrs(w2_qweight, extra_weight_attrs) - - num_groups_w13 = hidden_size // self.quant_config.group_size - num_groups_w2 = intermediate_size_per_partition // self.quant_config.group_size - - # WEIGHT_SCALES - # Allocate 2 scales for w1 and w3 respectively. - w13_scales = torch.nn.Parameter( - torch.empty( - num_experts, - num_groups_w13, - intermediate_size_per_partition * 2, - 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) - - # WEIGHT_ZERO_POINT - # Allocate 2 zero points for w1 and w3 respectively. - w13_qzeros = torch.nn.Parameter( - torch.empty( - num_experts, - num_groups_w13, - 2 * intermediate_size_per_partition // self.quant_config.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 // self.quant_config.pack_factor, - dtype=torch.int32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_qzeros", w2_qzeros) - set_weight_attrs(w2_qzeros, extra_weight_attrs) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - num_experts = layer.w13_qweight.shape[0] - device = layer.w13_qweight.device - - layer.w13_g_idx_sort_indices = torch.nn.Parameter( - torch.empty((num_experts, 0), dtype=torch.int32, device=device), - requires_grad=False, - ) - layer.w2_g_idx_sort_indices = torch.nn.Parameter( - torch.empty((num_experts, 0), dtype=torch.int32, device=device), - requires_grad=False, - ) - - marlin_w13_qweight = awq_marlin_moe_repack( - layer.w13_qweight, - layer.w13_g_idx_sort_indices, - size_k=layer.w13_qweight.shape[1], - size_n=layer.w13_qweight.shape[2] * self.quant_config.pack_factor, - num_bits=self.quant_config.weight_bits, - ) - replace_parameter(layer, "w13_qweight", marlin_w13_qweight) - - marlin_w2_qweight = awq_marlin_moe_repack( - layer.w2_qweight, - layer.w2_g_idx_sort_indices, - size_k=layer.w2_qweight.shape[1], - size_n=layer.w2_qweight.shape[2] * self.quant_config.pack_factor, - num_bits=self.quant_config.weight_bits, - ) - replace_parameter(layer, "w2_qweight", marlin_w2_qweight) - - # hidden_size->intermediate_size - marlin_w13_scales = marlin_moe_permute_scales( - s=layer.w13_scales, - size_k=layer.intermediate_size_per_partition, - size_n=layer.w13_scales.shape[2], - group_size=self.quant_config.group_size, - ) - - replace_parameter(layer, "w13_scales", marlin_w13_scales) - - marlin_w2_scales = marlin_moe_permute_scales( - s=layer.w2_scales, - size_k=layer.intermediate_size_per_partition, - size_n=layer.w2_scales.shape[2], - group_size=self.quant_config.group_size, - ) - replace_parameter(layer, "w2_scales", marlin_w2_scales) - - marlin_w13_zp = moe_awq_to_marlin_zero_points( - layer.w13_qzeros, - size_k=layer.w13_qzeros.shape[1], - size_n=layer.w13_qzeros.shape[2] * self.quant_config.pack_factor, - num_bits=self.quant_config.weight_bits, - ) - replace_parameter(layer, "w13_qzeros", marlin_w13_zp) - - marlin_w2_zp = moe_awq_to_marlin_zero_points( - layer.w2_qzeros, - size_k=layer.w2_qzeros.shape[1], - size_n=layer.w2_qzeros.shape[2] * self.quant_config.pack_factor, - num_bits=self.quant_config.weight_bits, - ) - replace_parameter(layer, "w2_qzeros", marlin_w2_zp) - - def create_moe_runner( - self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig - ): - assert get_moe_runner_backend().is_auto() - self.moe_runner_config = moe_runner_config - self.runner = MoeRunner(MoeRunnerBackend.MARLIN, moe_runner_config) - - def apply( - self, - layer: torch.nn.Module, - dispatch_output: StandardDispatchOutput, - ) -> CombineInput: - - quant_info = MarlinMoeQuantInfo( - w13_qweight=layer.w13_qweight, - w2_qweight=layer.w2_qweight, - w13_scales=layer.w13_scales, - w2_scales=layer.w2_scales, - w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices, - w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices, - w13_qzeros=layer.w13_qzeros, - w2_qzeros=layer.w2_qzeros, - weight_bits=self.quant_config.weight_bits, - ) - - return self.runner.run(dispatch_output, quant_info) - - -class AWQMoEAscendMethod(AWQMoEMethod): - def __init__(self, quant_config: AWQConfig): - self.quant_config = quant_config - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - w13_qweight_tmp = torch.zeros_like(layer.w13_qweight.data) - w2_qweight_tmp = torch.zeros_like(layer.w2_qweight.data) - w13_qzeros_list = [] - w2_qzeros_list = [] - shifts = [0, 4, 1, 5, 2, 6, 3, 7] - for i in range(0, self.quant_config.pack_factor): - shift_num = shifts[i] * 4 - w13_qzeros_list.append( - (layer.w13_qzeros.data.reshape(-1, 1) >> shift_num) & 0xF - ) - w2_qzeros_list.append( - (layer.w2_qzeros.data.reshape(-1, 1) >> shift_num) & 0xF - ) - w13_qweight_tmp.bitwise_or_( - ((layer.w13_qweight.data >> shift_num) * (2 ** (4 * i))) - & (0xF << (4 * i)) - ) - w2_qweight_tmp.bitwise_or_( - ((layer.w2_qweight.data >> shift_num) * (2 ** (4 * i))) - & (0xF << (4 * i)) - ) - - w13_qweight_tmp.bitwise_xor_(0x88888888) - w2_qweight_tmp.bitwise_xor_(0x88888888) - - w13_qzeros_tmp = torch.cat(w13_qzeros_list, dim=-1).reshape( - layer.w13_qzeros.shape[0], layer.w13_qzeros.shape[1], -1 - ) - w13_qzeros_tmp = -(w13_qzeros_tmp - 8) - w13_qzeros_tmp = w13_qzeros_tmp.to(layer.w13_scales.data.dtype) - w2_qzeros_tmp = torch.cat(w2_qzeros_list, dim=-1).reshape( - layer.w2_qzeros.shape[0], layer.w2_qzeros.shape[1], -1 - ) - w2_qzeros_tmp = -(w2_qzeros_tmp - 8) - w2_qzeros_tmp = w2_qzeros_tmp.to(layer.w2_scales.data.dtype) - - layer.register_parameter( - "w13_qzeros", torch.nn.Parameter(w13_qzeros_tmp, requires_grad=False) - ) - layer.register_parameter( - "w13_qweight", torch.nn.Parameter(w13_qweight_tmp, requires_grad=False) - ) - layer.register_parameter( - "w2_qzeros", torch.nn.Parameter(w2_qzeros_tmp, requires_grad=False) - ) - layer.register_parameter( - "w2_qweight", torch.nn.Parameter(w2_qweight_tmp, requires_grad=False) - ) - - def create_moe_runner( - self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig - ): - self.moe_runner_config = moe_runner_config - - 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.activation == "silu" - ), "Only SiLU activation is supported." - - 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) - - -# Register fake implementations for torch.compile support -if _is_cuda: - - @register_fake_if_exists("sgl_kernel::awq_marlin_repack") - def _(b_q_weight, size_k, size_n, num_bits): - return b_q_weight.new_empty( - (size_k // 16, size_n * (num_bits // 2)), dtype=b_q_weight.dtype - ) diff --git a/python/sglang/srt/layers/quantization/awq/__init__.py b/python/sglang/srt/layers/quantization/awq/__init__.py new file mode 100644 index 000000000..f54a5eab8 --- /dev/null +++ b/python/sglang/srt/layers/quantization/awq/__init__.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 + +from .awq import ( + AWQConfig, + AWQCPUConfig, + AWQLinearMethod, + AWQMarlinConfig, + AWQMoEMethod, +) +from .awq_triton import awq_dequantize_decomposition, awq_dequantize_triton +from .schemes import ( + AWQAscendLinearScheme, + AWQAscendMoEScheme, + AWQLinearScheme, + AWQMarlinLinearScheme, + AWQMoEScheme, +) + +__all__ = [ + "AWQConfig", + "AWQCPUConfig", + "AWQMarlinConfig", + "AWQLinearMethod", + "AWQMoEMethod", + "AWQLinearScheme", + "AWQMarlinLinearScheme", + "AWQAscendLinearScheme", + "AWQMoEScheme", + "AWQAscendMoEScheme", + "awq_dequantize_triton", + "awq_dequantize_decomposition", +] diff --git a/python/sglang/srt/layers/quantization/awq/awq.py b/python/sglang/srt/layers/quantization/awq/awq.py new file mode 100644 index 000000000..4d238d9a3 --- /dev/null +++ b/python/sglang/srt/layers/quantization/awq/awq.py @@ -0,0 +1,484 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import logging +import warnings +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +import torch + +from sglang.srt.layers.linear import LinearBase +from sglang.srt.layers.moe import MoeRunnerConfig +from sglang.srt.layers.quantization.base_config import ( + FusedMoEMethodBase, + LinearMethodBase, + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.srt.layers.quantization.marlin_utils import ( + check_marlin_supported, + check_marlin_supports_layer, + check_moe_marlin_supports_layer, + verify_marlin_supported, +) +from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod +from sglang.srt.layers.quantization.utils import get_scalar_types +from sglang.srt.utils.patch_torch import register_fake_if_exists + +from .schemes import ( + AWQAscendLinearScheme, + AWQAscendMoEScheme, + AWQIntelAMXLinearScheme, + AWQIntelAMXMoEScheme, + AWQLinearScheme, + AWQMarlinLinearScheme, + AWQMoEScheme, +) + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) + +from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu + +_is_cuda = is_cuda() +_is_hip = is_hip() +_is_xpu = is_xpu() +_is_npu = is_npu() + +if not (_is_cuda or _is_hip or _is_xpu or _is_npu): + warnings.warn(f"Only CUDA, HIP and XPU support AWQ currently.") + +logger = logging.getLogger(__name__) + + +ScalarType, scalar_types = get_scalar_types() + + +def is_layer_skipped_awq(prefix: str, modules_to_not_convert: List[str]): + return any(module_name in prefix for module_name in modules_to_not_convert) + + +class AWQConfig(QuantizationConfig): + """Config class for AWQ. + + Reference: https://arxiv.org/abs/2306.00978 + """ + + def __init__( + self, + weight_bits: int, + group_size: int, + zero_point: bool, + modules_to_not_convert: Optional[List[str]] = None, + ) -> None: + super().__init__() + self.weight_bits = weight_bits + self.group_size = group_size + self.zero_point = zero_point + self.modules_to_not_convert = modules_to_not_convert or [] + + if self.weight_bits != 4: + raise ValueError( + "Currently, only 4-bit weight quantization is supported for " + f"AWQ, but got {self.weight_bits} bits." + ) + self.pack_factor = 32 // self.weight_bits + + def __repr__(self) -> str: + return ( + f"AWQConfig(weight_bits={self.weight_bits}, " + f"group_size={self.group_size}, " + f"zero_point={self.zero_point}, " + f"modules_to_not_convert={self.modules_to_not_convert})" + ) + + def get_scaled_act_names(self) -> List[str]: + return [] + + def get_name(self) -> str: + return "awq" + + def get_supported_act_dtypes(self) -> List[torch.dtype]: + return [torch.float16] if not _is_npu else [torch.float16, torch.bfloat16] + + @classmethod + def get_min_capability(cls) -> int: + # The AWQ kernel only supports Turing or newer GPUs. + if _is_npu: + raise NotImplementedError( + 'NPU hardware does not support "get_min_capability" feature.' + ) + else: + return 75 + + @staticmethod + def get_config_filenames() -> List[str]: + return [ + "quant_config.json", # E.g., casperhansen/vicuna-7b-v1.5-awq + # E.g., abhinavkulkarni/mosaicml-mpt-7b-instruct-w4-g128-awq + "quantize_config.json", + ] + + @classmethod + def from_config(cls, config: Dict[str, Any]) -> AWQConfig: + weight_bits = cls.get_from_keys(config, ["w_bit", "bits"]) + group_size = cls.get_from_keys(config, ["q_group_size", "group_size"]) + zero_point = cls.get_from_keys(config, ["zero_point"]) + modules_to_not_convert = cls.get_from_keys_or( + config, ["modules_to_not_convert"], None + ) + return cls(weight_bits, group_size, zero_point, modules_to_not_convert) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[LinearMethodBase]: + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + if _is_npu: + if isinstance(layer, LinearBase): + if is_layer_skipped_awq(prefix, self.modules_to_not_convert): + return UnquantizedLinearMethod() + layer.scheme = self.get_linear_scheme(layer) + return AWQLinearMethod(self) + elif isinstance(layer, FusedMoE): + layer.scheme = self.get_moe_scheme(layer) + return AWQMoEMethod(self) + return None + + if isinstance(layer, LinearBase): + if is_layer_skipped_awq(prefix, self.modules_to_not_convert): + return UnquantizedLinearMethod() + layer.scheme = self.get_linear_scheme(layer) + return AWQLinearMethod(self) + return None + + def get_linear_scheme(self, layer: torch.nn.Module): + assert isinstance(layer, LinearBase) + # TODO: move platform-specific AWQ scheme selection into the platform + # plugin factory once quantization hooks are available there. + if _is_npu: + return AWQAscendLinearScheme(self) + return AWQLinearScheme(self) + + def get_moe_scheme(self, layer: torch.nn.Module): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + assert isinstance(layer, FusedMoE) + # This is currently only reached by the NPU path in get_quant_method. + if _is_npu: + return AWQAscendMoEScheme(self) + raise NotImplementedError("AWQConfig only supports MoE scheme on NPU.") + + +class AWQCPUConfig(AWQConfig): + """CPU Config class for AWQ, inherit from AWQConfig""" + + def get_supported_act_dtypes(self) -> List[torch.dtype]: + return [torch.float16, torch.bfloat16] + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[LinearMethodBase]: + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + if isinstance(layer, LinearBase): + if is_layer_skipped_awq(prefix, self.modules_to_not_convert): + return UnquantizedLinearMethod() + layer.scheme = self.get_linear_scheme(layer) + return AWQLinearMethod(self) + elif isinstance(layer, FusedMoE): + layer.scheme = self.get_moe_scheme(layer) + return AWQMoEMethod(self) + return None + + def get_linear_scheme(self, layer: torch.nn.Module): + from sglang.srt.layers.linear import LinearBase + + assert isinstance(layer, LinearBase) + return AWQIntelAMXLinearScheme(self) + + def get_moe_scheme(self, layer: torch.nn.Module): + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + + assert isinstance(layer, FusedMoE) + return AWQIntelAMXMoEScheme(self) + + +class AWQMarlinConfig(QuantizationConfig): + """Config class for AWQ Marlin""" + + # num_bits -> type + TYPE_MAP = { + 4: scalar_types.uint4, + 8: scalar_types.uint8, + } + + def __init__( + self, + weight_bits: int, + group_size: int, + zero_point: bool, + lm_head_quantized: bool, + modules_to_not_convert: Optional[list[str]], + full_config: dict[str, Any], + ) -> None: + super().__init__() + if _is_hip: + warnings.warn(f"HIP does not support fused_marlin_moe currently.") + self.pack_factor = 32 // weight_bits # packed into int32 + self.group_size = group_size + self.zero_point = zero_point + self.lm_head_quantized = lm_head_quantized + self.weight_bits = weight_bits + self.modules_to_not_convert = modules_to_not_convert or [] + self.full_config = full_config + + if self.weight_bits not in self.TYPE_MAP: + raise ValueError( + f"Unsupported num_bits = {self.weight_bits}. " + f"Supported num_bits = {self.TYPE_MAP.keys()}" + ) + + self.quant_type = self.TYPE_MAP[self.weight_bits] + + verify_marlin_supported( + self.quant_type, group_size=self.group_size, has_zp=self.zero_point + ) + + def __repr__(self) -> str: + return ( + f"AWQMarlinConfig(quant_type={self.quant_type}, " + f"group_size={self.group_size}, " + f"zero_point={self.zero_point}, " + f"lm_head_quantized={self.lm_head_quantized}, " + f"modules_to_not_convert={self.modules_to_not_convert})" + ) + + def get_scaled_act_names(self) -> List[str]: + return [] + + @classmethod + def get_name(cls) -> str: + return "awq_marlin" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.half, torch.bfloat16] + + @classmethod + def get_min_capability(cls) -> int: + return 80 + + @classmethod + def get_config_filenames(cls) -> list[str]: + return ["quantize_config.json"] + + @classmethod + def from_config(cls, config: dict[str, Any]) -> AWQMarlinConfig: + weight_bits = cls.get_from_keys(config, ["bits"]) + group_size = cls.get_from_keys(config, ["group_size"]) + zero_point = cls.get_from_keys(config, ["zero_point"]) + lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False) + modules_to_not_convert = cls.get_from_keys_or( + config, ["modules_to_not_convert"], None + ) + return cls( + weight_bits, + group_size, + zero_point, + lm_head_quantized, + modules_to_not_convert, + config, + ) + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant) -> Optional[str]: + can_convert = cls.is_awq_marlin_compatible(hf_quant_cfg) + is_valid_user_quant = ( + user_quant is None or user_quant == "marlin" or user_quant == "awq_marlin" + ) + + if can_convert and is_valid_user_quant: + msg = ( + "The model is convertible to {} during runtime." + " Using {} kernel.".format(cls.get_name(), cls.get_name()) + ) + logger.info(msg) + return cls.get_name() + + if can_convert and user_quant == "awq": + logger.info( + "Detected that the model can run with awq_marlin" + ", however you specified quantization=awq explicitly," + " so forcing awq. Use quantization=awq_marlin for" + " faster inference" + ) + return None + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> Optional[QuantizeMethodBase]: + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead + + if isinstance(layer, LinearBase) or ( + isinstance(layer, ParallelLMHead) and self.lm_head_quantized + ): + if is_layer_skipped_awq(prefix, self.modules_to_not_convert): + return UnquantizedLinearMethod() + # Check if the layer is supported by AWQMarlin. + if not check_marlin_supports_layer(layer, self.group_size): + logger.warning_once( + "Layer '%s' is not supported by AWQMarlin. Falling back to unoptimized AWQ kernels.", # noqa: E501 + prefix, + ) + return AWQConfig.from_config(self.full_config).get_quant_method( + layer, prefix + ) + layer.scheme = self.get_linear_scheme(layer) + return AWQLinearMethod(self) + elif isinstance(layer, FusedMoE): + from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config + + if not check_moe_marlin_supports_layer(layer, self.group_size): + logger.warning_once( + f"Layer '{prefix}' is not supported by AWQMoeMarlin. " + "Falling back to Moe WNA16 kernels." + ) + return MoeWNA16Config.from_config(self.full_config).get_quant_method( + layer, prefix + ) + layer.scheme = self.get_moe_scheme(layer) + return AWQMoEMethod(self) + return None + + def get_linear_scheme(self, layer: torch.nn.Module): + return AWQMarlinLinearScheme(self) + + def get_moe_scheme(self, layer: torch.nn.Module): + return AWQMoEScheme(self) + + @classmethod + def is_awq_marlin_compatible(cls, quant_config: dict[str, Any]): + # Extract data from quant config. + quant_method = quant_config.get("quant_method", "").lower() + num_bits = quant_config.get("bits") + group_size = quant_config.get("group_size") + zero_point = quant_config.get("zero_point") + + if not _is_cuda: + return False + + if quant_method != "awq": + return False + + # If we cannot find the info needed in the config, cannot convert. + if num_bits is None or group_size is None or zero_point is None: + return False + + if num_bits not in cls.TYPE_MAP: + return False + + return check_marlin_supported( + quant_type=cls.TYPE_MAP[num_bits], group_size=group_size, has_zp=zero_point + ) + + +class AWQLinearMethod(LinearMethodBase): + """Linear method for AWQ. + + Args: + quant_config: The AWQ quantization config. + """ + + def __init__(self, quant_config: AWQConfig): + self.quant_config = quant_config + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + weight_loader = extra_weight_attrs.get("weight_loader") + layer.scheme.create_weights( + layer=layer, + input_size_per_partition=input_size_per_partition, + output_partition_sizes=output_partition_sizes, + input_size=input_size, + output_size=output_size, + params_dtype=params_dtype, + weight_loader=weight_loader, + ) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.scheme.process_weights_after_loading(layer) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return layer.scheme.apply_weights(layer, x, bias) + + +class AWQMoEMethod(FusedMoEMethodBase): + + def __init__(self, quant_config: AWQMarlinConfig): + self.quant_config = quant_config + self.quant_type = scalar_types.uint4 + if self.quant_config.weight_bits != 4: + raise ValueError("AWQMoEMethod only supports 4bit now.") + + 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, + ): + layer.scheme.create_weights( + layer=layer, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size_per_partition=intermediate_size_per_partition, + params_dtype=params_dtype, + **extra_weight_attrs, + ) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.scheme.process_weights_after_loading(layer) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + layer.scheme.create_moe_runner(layer, moe_runner_config) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + return layer.scheme.apply_weights(layer, dispatch_output) + + +# Register fake implementations for torch.compile support +if _is_cuda: + + @register_fake_if_exists("sgl_kernel::awq_marlin_repack") + def _(b_q_weight, size_k, size_n, num_bits): + return b_q_weight.new_empty( + (size_k // 16, size_n * (num_bits // 2)), dtype=b_q_weight.dtype + ) diff --git a/python/sglang/srt/layers/quantization/awq_triton.py b/python/sglang/srt/layers/quantization/awq/awq_triton.py similarity index 100% rename from python/sglang/srt/layers/quantization/awq_triton.py rename to python/sglang/srt/layers/quantization/awq/awq_triton.py diff --git a/python/sglang/srt/layers/quantization/awq/schemes/__init__.py b/python/sglang/srt/layers/quantization/awq/schemes/__init__.py new file mode 100644 index 000000000..cbf38c3a8 --- /dev/null +++ b/python/sglang/srt/layers/quantization/awq/schemes/__init__.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 + +from .awq_cpu import AWQIntelAMXLinearScheme, AWQIntelAMXMoEScheme +from .awq_linear import AWQAscendLinearScheme, AWQLinearScheme +from .awq_marlin import AWQMarlinLinearScheme +from .awq_moe import AWQAscendMoEScheme, AWQMoEScheme +from .awq_scheme import AWQLinearSchemeBase, AWQMoESchemeBase + +__all__ = [ + "AWQLinearSchemeBase", + "AWQMoESchemeBase", + "AWQLinearScheme", + "AWQAscendLinearScheme", + "AWQIntelAMXLinearScheme", + "AWQMarlinLinearScheme", + "AWQMoEScheme", + "AWQAscendMoEScheme", + "AWQIntelAMXMoEScheme", +] diff --git a/python/sglang/srt/layers/quantization/awq_cpu.py b/python/sglang/srt/layers/quantization/awq/schemes/awq_cpu.py old mode 100755 new mode 100644 similarity index 58% rename from python/sglang/srt/layers/quantization/awq_cpu.py rename to python/sglang/srt/layers/quantization/awq/schemes/awq_cpu.py index 20d64bf07..7f258685d --- a/python/sglang/srt/layers/quantization/awq_cpu.py +++ b/python/sglang/srt/layers/quantization/awq/schemes/awq_cpu.py @@ -1,65 +1,29 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -import logging -from typing import TYPE_CHECKING, List, Optional +from typing import TYPE_CHECKING, Optional import torch -from sglang.srt.layers.moe import ( - MoeRunnerConfig, -) -from sglang.srt.layers.quantization.base_config import ( - LinearMethodBase, -) -from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod -from sglang.srt.layers.quantization.utils import get_scalar_types - -from .awq import AWQConfig, AWQLinearMethod, AWQMoEMethod - -if TYPE_CHECKING: - from sglang.srt.layers.moe.token_dispatcher import ( - StandardDispatchOutput, - ) - from sglang.srt.layers.amx_utils import ( CPUQuantMethod, _amx_process_weight_after_loading, ) +from sglang.srt.layers.moe import MoeRunnerConfig -logger = logging.getLogger(__name__) +from .awq_linear import AWQLinearScheme +from .awq_moe import AWQMoEScheme + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput + from sglang.srt.layers.quantization.awq.awq import AWQConfig + +__all__ = ["AWQIntelAMXLinearScheme", "AWQIntelAMXMoEScheme"] -ScalarType, scalar_types = get_scalar_types() - - -def is_layer_skipped_awq(prefix: str, modules_to_not_convert: List[str]): - return any(module_name in prefix for module_name in modules_to_not_convert) - - -class CPUAWQConfig(AWQConfig): - """CPU Config class for AWQ, inherit from AWQConfig""" - - def get_supported_act_dtypes(self) -> List[torch.dtype]: - return [torch.float16, torch.bfloat16] - - def get_quant_method( - self, layer: torch.nn.Module, prefix: str - ) -> Optional[LinearMethodBase]: - from sglang.srt.layers.linear import LinearBase - from sglang.srt.layers.moe.fused_moe_triton import FusedMoE - - if isinstance(layer, LinearBase): - if is_layer_skipped_awq(prefix, self.modules_to_not_convert): - return UnquantizedLinearMethod() - return AWQLinearIntelAMXMethod(self) - elif isinstance(layer, FusedMoE): - return AWQMoEIntelAMXMethod(self) - return None - - -class AWQLinearIntelAMXMethod(AWQLinearMethod): - """Linear method for AWQ on Intel CPU with AMX.""" +class AWQIntelAMXLinearKernel: + def __init__(self, quant_config: "AWQConfig"): + self.quant_config = quant_config def process_weights_after_loading(self, layer: torch.nn.Module) -> None: _amx_process_weight_after_loading( @@ -75,7 +39,6 @@ class AWQLinearIntelAMXMethod(AWQLinearMethod): x: torch.Tensor, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: - return torch.ops.sgl_kernel.int4_scaled_mm_cpu( x, layer.qweight, @@ -85,8 +48,16 @@ class AWQLinearIntelAMXMethod(AWQLinearMethod): ) -class AWQMoEIntelAMXMethod(AWQMoEMethod): - """MoE method for AWQ on Intel CPU with AMX.""" +class AWQIntelAMXLinearScheme(AWQLinearScheme): + """Linear scheme for AWQ on Intel CPU with AMX.""" + + def _init_kernel(self, quant_config: "AWQConfig"): + return AWQIntelAMXLinearKernel(quant_config) + + +class AWQIntelAMXMoEKernel: + def __init__(self, quant_config: "AWQConfig"): + self.quant_config = quant_config def process_weights_after_loading(self, layer: torch.nn.Module) -> None: _amx_process_weight_after_loading( @@ -104,7 +75,7 @@ class AWQMoEIntelAMXMethod(AWQMoEMethod): def apply( self, layer: torch.nn.Module, - dispatch_output: StandardDispatchOutput, + dispatch_output: "StandardDispatchOutput", ) -> torch.Tensor: from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput @@ -131,3 +102,16 @@ class AWQMoEIntelAMXMethod(AWQMoEMethod): True, # is_vnni ) return StandardCombineInput(hidden_states=output) + + +class AWQIntelAMXMoEScheme(AWQMoEScheme): + """MoE scheme for AWQ on Intel CPU with AMX.""" + + def _init_kernel(self, quant_config: "AWQConfig"): + return AWQIntelAMXMoEKernel(quant_config) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + self.kernel.create_moe_runner(layer, moe_runner_config) diff --git a/python/sglang/srt/layers/quantization/awq/schemes/awq_linear.py b/python/sglang/srt/layers/quantization/awq/schemes/awq_linear.py new file mode 100644 index 000000000..535776477 --- /dev/null +++ b/python/sglang/srt/layers/quantization/awq/schemes/awq_linear.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from typing import TYPE_CHECKING, List, Optional + +import torch + +from sglang.srt.layers.parameter import GroupQuantScaleParameter, PackedvLLMParameter + +from .awq_scheme import AWQLinearSchemeBase + +if TYPE_CHECKING: + from sglang.srt.layers.quantization.awq.awq import AWQConfig + +__all__ = ["AWQLinearScheme", "AWQAscendLinearScheme"] + + +class AWQLinearScheme(AWQLinearSchemeBase): + def __init__(self, quant_config: "AWQConfig"): + self.quant_config = quant_config + self.kernel = self._init_kernel(quant_config) + + def _init_kernel(self, quant_config: "AWQConfig"): + from sglang.srt.hardware_backend.gpu.quantization.awq_kernels import ( + AWQLinearKernel, + ) + + return AWQLinearKernel(quant_config) + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: List[int], + params_dtype: torch.dtype, + weight_loader, + **kwargs, + ): + if input_size_per_partition % self.quant_config.group_size != 0: + raise ValueError( + "The input size is not aligned with the quantized " + "weight shape. This can be caused by too large " + "tensor parallel size." + ) + + output_size_per_partition = sum(output_partition_sizes) + if output_size_per_partition % self.quant_config.pack_factor != 0: + raise ValueError( + "The output size is not aligned with the quantized " + "weight shape. This can be caused by too large " + "tensor parallel size." + ) + + qweight = PackedvLLMParameter( + data=torch.empty( + input_size_per_partition, + output_size_per_partition // self.quant_config.pack_factor, + dtype=torch.int32, + ), + input_dim=0, + output_dim=1, + packed_dim=1, + packed_factor=self.quant_config.pack_factor, + weight_loader=weight_loader, + ) + + qzeros = PackedvLLMParameter( + data=torch.empty( + input_size_per_partition // self.quant_config.group_size, + output_size_per_partition // self.quant_config.pack_factor, + dtype=torch.int32, + ), + input_dim=0, + output_dim=1, + packed_dim=1, + packed_factor=self.quant_config.pack_factor, + weight_loader=weight_loader, + ) + + scales = GroupQuantScaleParameter( + data=torch.empty( + input_size_per_partition // self.quant_config.group_size, + output_size_per_partition, + dtype=params_dtype, + ), + input_dim=0, + output_dim=1, + weight_loader=weight_loader, + ) + + layer.register_parameter("qweight", qweight) + layer.register_parameter("qzeros", qzeros) + layer.register_parameter("scales", scales) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + self.kernel.process_weights_after_loading(layer) + + def apply_weights( + self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] + ): + return self.kernel.apply(layer, x, bias) + + +class AWQAscendLinearScheme(AWQLinearScheme): + def _init_kernel(self, quant_config: "AWQConfig"): + from sglang.srt.hardware_backend.npu.quantization.awq_kernels import ( + AWQAscendLinearKernel, + ) + + return AWQAscendLinearKernel(quant_config) diff --git a/python/sglang/srt/layers/quantization/awq/schemes/awq_marlin.py b/python/sglang/srt/layers/quantization/awq/schemes/awq_marlin.py new file mode 100644 index 000000000..b92a7cba9 --- /dev/null +++ b/python/sglang/srt/layers/quantization/awq/schemes/awq_marlin.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.hardware_backend.gpu.quantization.awq_kernels import ( + AWQMarlinLinearKernel, +) +from sglang.srt.layers.parameter import GroupQuantScaleParameter, PackedvLLMParameter +from sglang.srt.layers.quantization.marlin_utils import verify_marlin_supports_shape + +from .awq_scheme import AWQLinearSchemeBase + +if TYPE_CHECKING: + from sglang.srt.layers.quantization.awq.awq import AWQMarlinConfig + +__all__ = ["AWQMarlinLinearScheme"] + + +class AWQMarlinLinearScheme(AWQLinearSchemeBase): + def __init__(self, quant_config: "AWQMarlinConfig"): + self.quant_config = quant_config + self.kernel = AWQMarlinLinearKernel(quant_config) + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + params_dtype: torch.dtype, + weight_loader, + **kwargs, + ) -> None: + output_size_per_partition = sum(output_partition_sizes) + + group_size = ( + self.quant_config.group_size + if self.quant_config.group_size != -1 + else input_size + ) + + verify_marlin_supports_shape( + output_size_per_partition=output_size_per_partition, + input_size_per_partition=input_size_per_partition, + input_size=input_size, + group_size=group_size, + ) + + qweight = PackedvLLMParameter( + data=torch.empty( + input_size_per_partition, + output_size_per_partition // self.quant_config.pack_factor, + dtype=torch.int32, + ), + input_dim=0, + output_dim=1, + packed_dim=1, + packed_factor=self.quant_config.pack_factor, + weight_loader=weight_loader, + ) + + num_groups = input_size_per_partition // group_size + + qzeros = PackedvLLMParameter( + data=torch.empty( + num_groups, + output_size_per_partition // self.quant_config.pack_factor, + dtype=torch.int32, + ), + input_dim=0, + output_dim=1, + packed_dim=1, + packed_factor=self.quant_config.pack_factor, + weight_loader=weight_loader, + ) + + scales = GroupQuantScaleParameter( + data=torch.empty( + num_groups, + output_size_per_partition, + dtype=params_dtype, + ), + input_dim=0, + output_dim=1, + weight_loader=weight_loader, + ) + + layer.register_parameter("qweight", qweight) + layer.register_parameter("qzeros", qzeros) + layer.register_parameter("scales", scales) + + layer.input_size_per_partition = input_size_per_partition + layer.output_size_per_partition = output_size_per_partition + layer.num_groups = num_groups + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + self.kernel.process_weights_after_loading(layer) + + def apply_weights( + self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] + ): + return self.kernel.apply(layer, x, bias) diff --git a/python/sglang/srt/layers/quantization/awq/schemes/awq_moe.py b/python/sglang/srt/layers/quantization/awq/schemes/awq_moe.py new file mode 100644 index 000000000..2233aa99a --- /dev/null +++ b/python/sglang/srt/layers/quantization/awq/schemes/awq_moe.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.srt.layers.linear import set_weight_attrs +from sglang.srt.layers.moe import ( + MoeRunner, + MoeRunnerBackend, + MoeRunnerConfig, + get_moe_runner_backend, +) + +from .awq_scheme import AWQMoESchemeBase + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput + from sglang.srt.layers.quantization.awq.awq import AWQConfig, AWQMarlinConfig + +__all__ = ["AWQMoEScheme", "AWQAscendMoEScheme"] + + +class AWQMoEScheme(AWQMoESchemeBase): + def __init__(self, quant_config: "AWQMarlinConfig"): + self.quant_config = quant_config + if self.quant_config.weight_bits != 4: + raise ValueError("AWQMoEScheme only supports 4bit now.") + self.kernel = self._init_kernel(quant_config) + + def _init_kernel(self, quant_config: "AWQMarlinConfig"): + from sglang.srt.hardware_backend.gpu.quantization.awq_kernels import ( + AWQMoEKernel, + ) + + return AWQMoEKernel(quant_config) + + 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 + + extra_weight_attrs.update( + { + "is_transposed": True, + "quant_method": FusedMoeWeightScaleSupported.GROUP.value, + } + ) + + w13_qweight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + 2 * intermediate_size_per_partition // self.quant_config.pack_factor, + 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, + hidden_size // self.quant_config.pack_factor, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_qweight", w2_qweight) + set_weight_attrs(w2_qweight, extra_weight_attrs) + + num_groups_w13 = hidden_size // self.quant_config.group_size + num_groups_w2 = intermediate_size_per_partition // self.quant_config.group_size + + w13_scales = torch.nn.Parameter( + torch.empty( + num_experts, + num_groups_w13, + intermediate_size_per_partition * 2, + 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 // self.quant_config.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 // self.quant_config.pack_factor, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_qzeros", w2_qzeros) + set_weight_attrs(w2_qzeros, 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 + ): + assert get_moe_runner_backend().is_auto() + self.moe_runner_config = moe_runner_config + self.kernel.runner = MoeRunner(MoeRunnerBackend.MARLIN, moe_runner_config) + + def apply_weights( + self, + layer: torch.nn.Module, + dispatch_output: "StandardDispatchOutput", + ): + return self.kernel.apply(layer, dispatch_output) + + +class AWQAscendMoEScheme(AWQMoEScheme): + def _init_kernel(self, quant_config: "AWQConfig"): + from sglang.srt.hardware_backend.npu.quantization.awq_kernels import ( + AWQAscendMoEKernel, + ) + + return AWQAscendMoEKernel(quant_config) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config diff --git a/python/sglang/srt/layers/quantization/awq/schemes/awq_scheme.py b/python/sglang/srt/layers/quantization/awq/schemes/awq_scheme.py new file mode 100644 index 000000000..8094d45f3 --- /dev/null +++ b/python/sglang/srt/layers/quantization/awq/schemes/awq_scheme.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 + +from abc import abstractmethod +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.layers.moe import MoeRunnerConfig +from sglang.srt.layers.quantization.base_scheme import BaseLinearScheme, BaseMoEScheme + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput + +__all__ = ["AWQLinearSchemeBase", "AWQMoESchemeBase"] + + +class AWQLinearSchemeBase(BaseLinearScheme): + @abstractmethod + def create_weights(self, *args, **kwargs): + raise NotImplementedError + + @abstractmethod + def process_weights_after_loading(self, layer: torch.nn.Module): + raise NotImplementedError + + @abstractmethod + def apply_weights( + self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] + ): + raise NotImplementedError + + +class AWQMoESchemeBase(BaseMoEScheme): + @abstractmethod + def create_weights(self, *args, **kwargs): + raise NotImplementedError + + @abstractmethod + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + raise NotImplementedError + + @abstractmethod + def process_weights_after_loading(self, layer: torch.nn.Module): + raise NotImplementedError + + @abstractmethod + def apply_weights( + self, + layer: torch.nn.Module, + dispatch_output: "StandardDispatchOutput", + ): + raise NotImplementedError diff --git a/python/sglang/srt/models/bailing_moe_linear.py b/python/sglang/srt/models/bailing_moe_linear.py index 3f47c038c..1cdaba88b 100644 --- a/python/sglang/srt/models/bailing_moe_linear.py +++ b/python/sglang/srt/models/bailing_moe_linear.py @@ -100,7 +100,7 @@ if _is_cuda: elif _is_cpu and _is_cpu_amx_available: pass elif _is_hip: - from sglang.srt.layers.quantization.awq_triton import ( + from sglang.srt.layers.quantization.awq.awq_triton import ( awq_dequantize_triton as awq_dequantize, ) else: diff --git a/python/sglang/srt/models/deepseek_common/utils.py b/python/sglang/srt/models/deepseek_common/utils.py index fab9594ef..59f0d8970 100644 --- a/python/sglang/srt/models/deepseek_common/utils.py +++ b/python/sglang/srt/models/deepseek_common/utils.py @@ -79,14 +79,14 @@ def awq_dequantize_func(): return awq_dequantize elif _is_hip: from sglang.kernel_api_logging import debug_kernel_api - from sglang.srt.layers.quantization.awq_triton import ( + from sglang.srt.layers.quantization.awq.awq_triton import ( awq_dequantize_triton as awq_dequantize, ) return debug_kernel_api(awq_dequantize, op_name="DeepseekCommon.awq_dequantize") elif _is_npu: from sglang.kernel_api_logging import debug_kernel_api - from sglang.srt.layers.quantization.awq_triton import ( + from sglang.srt.layers.quantization.awq.awq_triton import ( awq_dequantize_decomposition as awq_dequantize, ) diff --git a/python/sglang/srt/models/longcat_flash.py b/python/sglang/srt/models/longcat_flash.py index 028148a00..6536c46f0 100644 --- a/python/sglang/srt/models/longcat_flash.py +++ b/python/sglang/srt/models/longcat_flash.py @@ -117,7 +117,7 @@ if _is_cuda: elif _is_cpu and _is_cpu_amx_available: pass elif _is_hip: - from sglang.srt.layers.quantization.awq_triton import ( + from sglang.srt.layers.quantization.awq.awq_triton import ( awq_dequantize_triton as awq_dequantize, ) else: diff --git a/python/sglang/srt/models/longcat_flash_nextn.py b/python/sglang/srt/models/longcat_flash_nextn.py index d3dc20f36..c5a630cf3 100644 --- a/python/sglang/srt/models/longcat_flash_nextn.py +++ b/python/sglang/srt/models/longcat_flash_nextn.py @@ -97,7 +97,7 @@ if _is_cuda: elif _is_cpu and _is_cpu_amx_available: pass elif _is_hip: - from sglang.srt.layers.quantization.awq_triton import ( + from sglang.srt.layers.quantization.awq.awq_triton import ( awq_dequantize_triton as awq_dequantize, ) else: diff --git a/test/registered/quant/test_awq_dequant.py b/test/registered/quant/test_awq_dequant.py index ee0e80f1f..41f60ac74 100644 --- a/test/registered/quant/test_awq_dequant.py +++ b/test/registered/quant/test_awq_dequant.py @@ -12,7 +12,7 @@ import unittest import torch -from sglang.srt.layers.quantization.awq_triton import ( +from sglang.srt.layers.quantization.awq.awq_triton import ( AWQ_TRITON_SUPPORTED_GROUP_SIZES, awq_dequantize_triton, awq_gemm_triton,