From 1f5dc2cdca0a2f2c19076c2ad3539a128aaf2fb3 Mon Sep 17 00:00:00 2001 From: Yaochen Han <48639761+Alisehen@users.noreply.github.com> Date: Mon, 8 Jun 2026 18:14:40 +0800 Subject: [PATCH] [GPTQ] Refactor CPU quantization schemes (#26786) Co-authored-by: ronnie_zheng --- .../cpu/quantization/awq_kernels.py | 99 ++++++++ .../cpu/quantization/gptq_kernels.py | 99 ++++++++ .../srt/layers/quantization/__init__.py | 6 +- .../srt/layers/quantization/auto_round.py | 10 +- .../quantization/awq/schemes/awq_cpu.py | 89 +------- .../quantization/awq/schemes/awq_marlin.py | 12 +- .../schemes/compressed_tensors_wNa16_moe.py | 4 +- .../srt/layers/quantization/gptq/__init__.py | 17 +- .../srt/layers/quantization/gptq/gptq.py | 46 +++- .../quantization/gptq/schemes/__init__.py | 3 + .../{ => gptq/schemes}/gptq_cpu.py | 214 +++++------------- .../quantization/gptq/schemes/gptq_linear.py | 13 +- .../quantization/gptq/schemes/gptq_marlin.py | 14 +- 13 files changed, 348 insertions(+), 278 deletions(-) create mode 100644 python/sglang/srt/hardware_backend/cpu/quantization/awq_kernels.py create mode 100644 python/sglang/srt/hardware_backend/cpu/quantization/gptq_kernels.py rename python/sglang/srt/layers/quantization/{ => gptq/schemes}/gptq_cpu.py (55%) diff --git a/python/sglang/srt/hardware_backend/cpu/quantization/awq_kernels.py b/python/sglang/srt/hardware_backend/cpu/quantization/awq_kernels.py new file mode 100644 index 000000000..1ef41d153 --- /dev/null +++ b/python/sglang/srt/hardware_backend/cpu/quantization/awq_kernels.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.layers.amx_utils import ( + CPUQuantMethod, + _amx_process_weight_after_loading, +) +from sglang.srt.layers.moe import MoeRunnerConfig + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput + from sglang.srt.layers.quantization.awq.awq import AWQConfig + +__all__ = ["AWQIntelAMXLinearKernel", "AWQIntelAMXMoEKernel"] + + +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( + layer, ["qweight", "qzeros", "scales"], None, "awq" + ) + 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: + return torch.ops.sgl_kernel.int4_scaled_mm_cpu( + x, + layer.qweight, + layer.qzeros, + layer.scales, + bias, + ) + + +class AWQIntelAMXMoEKernel: + def __init__(self, quant_config: "AWQConfig"): + self.quant_config = quant_config + self.moe_runner_config: Optional[MoeRunnerConfig] = None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + _amx_process_weight_after_loading( + layer, ["w13_qweight", "w13_qzeros", "w13_scales"], None, "awq" + ) + _amx_process_weight_after_loading( + layer, ["w2_qweight", "w2_qzeros", "w2_scales"], None, "awq" + ) + + 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 + output = torch.ops.sgl_kernel.fused_experts_cpu( + x, + layer.w13_qweight, + layer.w2_qweight, + topk_weights, + topk_ids, + False, # inplace See [Note] inplace should be False in fused_experts. + CPUQuantMethod.INT4_W4A8, + layer.w13_scales, # w1_scale + layer.w2_scales, # w2_scale + layer.w13_qzeros, + layer.w2_qzeros, + None, # block_size + None, # w1 bias + None, # w3 bias + None, # alpha + None, # limit + True, # is_vnni + ) + return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/hardware_backend/cpu/quantization/gptq_kernels.py b/python/sglang/srt/hardware_backend/cpu/quantization/gptq_kernels.py new file mode 100644 index 000000000..8923c4b94 --- /dev/null +++ b/python/sglang/srt/hardware_backend/cpu/quantization/gptq_kernels.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.layers.amx_utils import ( + CPUQuantMethod, + _amx_process_weight_after_loading, +) +from sglang.srt.layers.moe import MoeRunnerConfig + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput + from sglang.srt.layers.quantization.gptq.gptq import GPTQConfig + +__all__ = ["GPTQIntelAMXLinearKernel", "GPTQIntelAMXMoEKernel"] + + +class GPTQIntelAMXLinearKernel: + def __init__(self, quant_config: "GPTQConfig"): + self.quant_config = quant_config + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + _amx_process_weight_after_loading( + layer, ["qweight", "qzeros", "scales"], None, "gptq" + ) + 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: + return torch.ops.sgl_kernel.int4_scaled_mm_cpu( + x, + layer.qweight, + layer.qzeros, + layer.scales, + bias, + ) + + +class GPTQIntelAMXMoEKernel: + def __init__(self, quant_config: "GPTQConfig"): + self.quant_config = quant_config + self.moe_runner_config: Optional[MoeRunnerConfig] = None + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + _amx_process_weight_after_loading( + layer, ["w13_qweight", "w13_qzeros", "w13_scales"], None, "gptq" + ) + _amx_process_weight_after_loading( + layer, ["w2_qweight", "w2_qzeros", "w2_scales"], None, "gptq" + ) + + 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 + output = torch.ops.sgl_kernel.fused_experts_cpu( + x, + layer.w13_qweight, + layer.w2_qweight, + topk_weights, + topk_ids, + False, # inplace See [Note] inplace should be False in fused_experts. + CPUQuantMethod.INT4_W4A8, + layer.w13_scales, # w1_scale + layer.w2_scales, # w2_scale + layer.w13_qzeros, + layer.w2_qzeros, + None, # block_size + None, # w1 bias + None, # w3 bias + None, # alpha + None, # limit + True, # is_vnni + ) + return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/quantization/__init__.py b/python/sglang/srt/layers/quantization/__init__.py index fb2c47513..82668bded 100644 --- a/python/sglang/srt/layers/quantization/__init__.py +++ b/python/sglang/srt/layers/quantization/__init__.py @@ -30,11 +30,11 @@ from sglang.srt.layers.quantization.fp8 import Fp8Config from sglang.srt.layers.quantization.fpgemm_fp8 import FBGEMMFp8Config from sglang.srt.layers.quantization.gguf import GGUFConfig from sglang.srt.layers.quantization.gptq import ( + CPUGPTQConfig, GPTQAscendConfig, GPTQConfig, GPTQMarlinConfig, ) -from sglang.srt.layers.quantization.gptq_cpu import CPUGPTQConfig from sglang.srt.layers.quantization.mlx import MlxQuantizationConfig from sglang.srt.layers.quantization.modelopt_quant import ( ModelOptFp4Config, @@ -54,7 +54,6 @@ from sglang.srt.layers.quantization.w8a8_int8 import W8A8Int8Config from sglang.srt.platforms import current_platform from sglang.srt.utils import ( cpu_has_amx_support, - is_cpu, is_cuda, is_hip, is_mps, @@ -98,7 +97,7 @@ BASE_QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = { } -if is_cpu() or is_cuda() or (_is_mxfp_supported and is_hip()): +if is_cuda() or (_is_mxfp_supported and is_hip()): BASE_QUANTIZATION_METHODS.update( { "mxfp4": Mxfp4Config, @@ -129,7 +128,6 @@ CPU_QUANTIZATION_METHODS = { "compressed-tensors": CompressedTensorsConfig, "awq": AWQCPUConfig, "gptq": CPUGPTQConfig, - "mxfp4": Mxfp4Config, } QUANTIZATION_METHODS = {**BASE_QUANTIZATION_METHODS} diff --git a/python/sglang/srt/layers/quantization/auto_round.py b/python/sglang/srt/layers/quantization/auto_round.py index e82456dc3..fb3451490 100644 --- a/python/sglang/srt/layers/quantization/auto_round.py +++ b/python/sglang/srt/layers/quantization/auto_round.py @@ -309,8 +309,8 @@ class AutoRoundConfig(QuantizationConfig): from sglang.srt.layers.moe.fused_moe_triton import FusedMoE from sglang.srt.layers.quantization.gptq import ( GPTQAscendConfig, - GPTQLinearAscendMethod, - GPTQMoEAscendMethod, + GPTQLinearMethod, + GPTQMoEMethod, ) from sglang.srt.layers.quantization.marlin_utils import ( check_marlin_supported, @@ -345,10 +345,12 @@ class AutoRoundConfig(QuantizationConfig): quant_args.sym = sym if isinstance(layer, FusedMoE): - return GPTQMoEAscendMethod(quant_args) + layer.scheme = quant_args.get_moe_scheme(layer) + return GPTQMoEMethod(quant_args) if isinstance(layer, (LinearBase, ParallelLMHead)): - return GPTQLinearAscendMethod(quant_args) + layer.scheme = quant_args.get_linear_scheme(layer) + return GPTQLinearMethod(quant_args) return None diff --git a/python/sglang/srt/layers/quantization/awq/schemes/awq_cpu.py b/python/sglang/srt/layers/quantization/awq/schemes/awq_cpu.py index f1dc785ae..3560a88d1 100644 --- a/python/sglang/srt/layers/quantization/awq/schemes/awq_cpu.py +++ b/python/sglang/srt/layers/quantization/awq/schemes/awq_cpu.py @@ -1,13 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING import torch -from sglang.srt.layers.amx_utils import ( - CPUQuantMethod, - _amx_process_weight_after_loading, +from sglang.srt.hardware_backend.cpu.quantization.awq_kernels import ( + AWQIntelAMXLinearKernel, + AWQIntelAMXMoEKernel, ) from sglang.srt.layers.moe import MoeRunnerConfig @@ -15,39 +15,11 @@ 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"] -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( - layer, ["qweight", "qzeros", "scales"], None, "awq" - ) - 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: - return torch.ops.sgl_kernel.int4_scaled_mm_cpu( - x, - layer.qweight, - layer.qzeros, - layer.scales, - bias, - ) - - class AWQIntelAMXLinearScheme(AWQLinearScheme): """Linear scheme for AWQ on Intel CPU with AMX.""" @@ -55,59 +27,6 @@ class AWQIntelAMXLinearScheme(AWQLinearScheme): 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( - layer, ["w13_qweight", "w13_qzeros", "w13_scales"], None, "awq" - ) - _amx_process_weight_after_loading( - layer, ["w2_qweight", "w2_qzeros", "w2_scales"], None, "awq" - ) - - 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 - output = torch.ops.sgl_kernel.fused_experts_cpu( - x, - layer.w13_qweight, - layer.w2_qweight, - topk_weights, - topk_ids, - False, # inplace See [Note] inplace should be False in fused_experts. - CPUQuantMethod.INT4_W4A8, - layer.w13_scales, # w1_scale - layer.w2_scales, # w2_scale - layer.w13_qzeros, - layer.w2_qzeros, - None, # block_size - None, # w1 bias - None, # w3 bias - None, # alpha - None, # limit - True, # is_vnni - ) - return StandardCombineInput(hidden_states=output) - - class AWQIntelAMXMoEScheme(AWQMoEScheme): """MoE scheme for AWQ on Intel CPU with AMX.""" diff --git a/python/sglang/srt/layers/quantization/awq/schemes/awq_marlin.py b/python/sglang/srt/layers/quantization/awq/schemes/awq_marlin.py index b92a7cba9..10f115dd5 100644 --- a/python/sglang/srt/layers/quantization/awq/schemes/awq_marlin.py +++ b/python/sglang/srt/layers/quantization/awq/schemes/awq_marlin.py @@ -5,9 +5,6 @@ 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 @@ -22,7 +19,14 @@ __all__ = ["AWQMarlinLinearScheme"] class AWQMarlinLinearScheme(AWQLinearSchemeBase): def __init__(self, quant_config: "AWQMarlinConfig"): self.quant_config = quant_config - self.kernel = AWQMarlinLinearKernel(quant_config) + self.kernel = self._init_kernel(quant_config) + + def _init_kernel(self, quant_config: "AWQMarlinConfig"): + from sglang.srt.hardware_backend.gpu.quantization.awq_kernels import ( + AWQMarlinLinearKernel, + ) + + return AWQMarlinLinearKernel(quant_config) def create_weights( self, diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py index 60e4ef5d6..58562bb23 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py @@ -8,6 +8,9 @@ from typing import TYPE_CHECKING import torch from compressed_tensors import CompressionFormat +from sglang.srt.hardware_backend.gpu.quantization.gptq_kernels import ( + gptq_marlin_moe_repack, +) from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( NPUW4A16Int4DynamicMoEMethod, ) @@ -16,7 +19,6 @@ from sglang.srt.layers.quantization.compressed_tensors.schemes import ( WNA16_SUPPORTED_BITS, CompressedTensorsMoEScheme, ) -from sglang.srt.layers.quantization.gptq import gptq_marlin_moe_repack from sglang.srt.layers.quantization.marlin_utils import ( marlin_make_workspace, marlin_moe_permute_scales, diff --git a/python/sglang/srt/layers/quantization/gptq/__init__.py b/python/sglang/srt/layers/quantization/gptq/__init__.py index c754cc081..f706e77ec 100644 --- a/python/sglang/srt/layers/quantization/gptq/__init__.py +++ b/python/sglang/srt/layers/quantization/gptq/__init__.py @@ -1,22 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 -from sglang.srt.hardware_backend.gpu.quantization.gptq_kernels import ( - gptq_marlin_moe_repack, -) - from .gptq import ( + CPUGPTQConfig, GPTQAscendConfig, GPTQConfig, - GPTQLinearAscendMethod, GPTQLinearMethod, GPTQMarlinConfig, GPTQMarlinLinearMethod, GPTQMarlinMoEMethod, - GPTQMoEAscendMethod, + GPTQMoEMethod, check_marlin_format, ) from .schemes import ( GPTQAscendLinearScheme, + GPTQIntelAMXLinearScheme, + GPTQIntelAMXMoEScheme, GPTQLinearScheme, GPTQMarlinLinearScheme, GPTQMarlinMoEScheme, @@ -26,17 +24,18 @@ from .schemes import ( __all__ = [ "GPTQConfig", "GPTQAscendConfig", + "CPUGPTQConfig", "GPTQMarlinConfig", "GPTQLinearMethod", - "GPTQMoEAscendMethod", + "GPTQMoEMethod", "GPTQMarlinLinearMethod", - "GPTQLinearAscendMethod", "GPTQMarlinMoEMethod", "GPTQLinearScheme", "GPTQAscendLinearScheme", + "GPTQIntelAMXLinearScheme", + "GPTQIntelAMXMoEScheme", "GPTQMarlinLinearScheme", "GPTQMoEAscendScheme", "GPTQMarlinMoEScheme", "check_marlin_format", - "gptq_marlin_moe_repack", ] diff --git a/python/sglang/srt/layers/quantization/gptq/gptq.py b/python/sglang/srt/layers/quantization/gptq/gptq.py index 9112ba8ff..8cc9f81d3 100644 --- a/python/sglang/srt/layers/quantization/gptq/gptq.py +++ b/python/sglang/srt/layers/quantization/gptq/gptq.py @@ -22,6 +22,8 @@ from sglang.srt.utils.patch_torch import register_fake_if_exists from .schemes import ( GPTQAscendLinearScheme, + GPTQIntelAMXLinearScheme, + GPTQIntelAMXMoEScheme, GPTQLinearScheme, GPTQMarlinLinearScheme, GPTQMarlinMoEScheme, @@ -209,10 +211,10 @@ class GPTQAscendConfig(GPTQConfig): if isinstance(layer, FusedMoE): layer.scheme = self.get_moe_scheme(layer) - return GPTQMoEAscendMethod(self) + return GPTQMoEMethod(self) if isinstance(layer, LinearBase): layer.scheme = self.get_linear_scheme(layer) - return GPTQLinearAscendMethod(self) + return GPTQLinearMethod(self) return None def get_linear_scheme(self, layer: torch.nn.Module): @@ -225,6 +227,40 @@ class GPTQAscendConfig(GPTQConfig): return GPTQMoEAscendScheme(self) +class CPUGPTQConfig(GPTQConfig): + """CPU Config class for GPTQ on Intel CPU with AMX.""" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return [torch.half, 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): + layer.scheme = self.get_linear_scheme(layer) + return GPTQLinearMethod(self) + if isinstance(layer, FusedMoE): + layer.scheme = self.get_moe_scheme(layer) + return GPTQMoEMethod(self) + return None + + def get_linear_scheme(self, layer: torch.nn.Module): + from sglang.srt.layers.linear import LinearBase + + assert isinstance(layer, LinearBase) + return GPTQIntelAMXLinearScheme(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 GPTQIntelAMXMoEScheme(self) + + class GPTQMarlinConfig(QuantizationConfig): """Config class for GPTQ Marlin""" @@ -460,7 +496,7 @@ class GPTQLinearMethod(LinearMethodBase): return layer.scheme.apply_weights(layer, x, bias) -class GPTQMoEAscendMethod(FusedMoEMethodBase): +class GPTQMoEMethod(FusedMoEMethodBase): def __init__(self, quant_config: GPTQConfig): super().__init__() @@ -552,10 +588,6 @@ class GPTQMarlinLinearMethod(LinearMethodBase): return layer.scheme.apply_weights(layer, x, bias) -class GPTQLinearAscendMethod(GPTQLinearMethod): - """Linear method for GPTQ on Ascend NPU.""" - - class GPTQMarlinMoEMethod(FusedMoEMethodBase): """MoE Marlin method with quantization.""" diff --git a/python/sglang/srt/layers/quantization/gptq/schemes/__init__.py b/python/sglang/srt/layers/quantization/gptq/schemes/__init__.py index 326157ca7..72682db95 100644 --- a/python/sglang/srt/layers/quantization/gptq/schemes/__init__.py +++ b/python/sglang/srt/layers/quantization/gptq/schemes/__init__.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 +from .gptq_cpu import GPTQIntelAMXLinearScheme, GPTQIntelAMXMoEScheme from .gptq_linear import GPTQAscendLinearScheme, GPTQLinearScheme from .gptq_marlin import GPTQMarlinLinearScheme from .gptq_moe import GPTQMarlinMoEScheme, GPTQMoEAscendScheme @@ -10,7 +11,9 @@ __all__ = [ "GPTQMoESchemeBase", "GPTQLinearScheme", "GPTQAscendLinearScheme", + "GPTQIntelAMXLinearScheme", "GPTQMarlinLinearScheme", "GPTQMoEAscendScheme", + "GPTQIntelAMXMoEScheme", "GPTQMarlinMoEScheme", ] diff --git a/python/sglang/srt/layers/quantization/gptq_cpu.py b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_cpu.py similarity index 55% rename from python/sglang/srt/layers/quantization/gptq_cpu.py rename to python/sglang/srt/layers/quantization/gptq/schemes/gptq_cpu.py index e50f14fb6..ebf463659 100644 --- a/python/sglang/srt/layers/quantization/gptq_cpu.py +++ b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_cpu.py @@ -1,12 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -from typing import TYPE_CHECKING, List, Optional +from typing import TYPE_CHECKING import torch -from sglang.srt.layers.moe import ( - MoeRunnerConfig, +from sglang.srt.hardware_backend.cpu.quantization.gptq_kernels import ( + GPTQIntelAMXLinearKernel, + GPTQIntelAMXMoEKernel, ) +from sglang.srt.layers.linear import set_weight_attrs +from sglang.srt.layers.moe import MoeRunnerConfig from sglang.srt.layers.parameter import ( ChannelQuantScaleParameter, GroupQuantScaleParameter, @@ -14,52 +18,36 @@ from sglang.srt.layers.parameter import ( PackedvLLMParameter, RowvLLMParameter, ) -from sglang.srt.layers.quantization.base_config import ( - FusedMoEMethodBase, - LinearMethodBase, -) + +from .gptq_linear import GPTQLinearScheme +from .gptq_scheme import GPTQMoESchemeBase if TYPE_CHECKING: - from sglang.srt.layers.moe.token_dispatcher import ( - StandardDispatchOutput, - ) + from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput + from sglang.srt.layers.quantization.gptq.gptq import GPTQConfig -from sglang.srt.layers.amx_utils import ( - CPUQuantMethod, - _amx_process_weight_after_loading, -) - -from .gptq import GPTQConfig +__all__ = ["GPTQIntelAMXLinearScheme", "GPTQIntelAMXMoEScheme"] -class CPUGPTQConfig(GPTQConfig): - """CPU Config class for AWQ, inherit from AWQConfig""" - - @classmethod - def get_supported_act_dtypes(cls) -> List[torch.dtype]: - return [torch.half, torch.bfloat16] - - def get_quant_method( - self, layer: torch.nn.Module, prefix: str - ) -> Optional[LinearMethodBase]: - # Delay the import to avoid circular dependency - from sglang.srt.layers.linear import LinearBase - from sglang.srt.layers.moe.fused_moe_triton import FusedMoE - - if isinstance(layer, FusedMoE): - return GPTQMoEIntelAMXMethod(self) - - if isinstance(layer, LinearBase): - return GPTQLinearIntelAMXMethod(self) +def _check_cpu_amx_support(quant_config: "GPTQConfig") -> None: + if quant_config.desc_act and not ( + quant_config.true_sequential and quant_config.static_groups + ): + raise ValueError( + "Currently, desc_act (True) is only supported with sequential " + "and static group on CPU with AMX." + ) + if quant_config.weight_bits != 4: + raise ValueError("Currently, only 4bits is supported on CPU with AMX.") + if quant_config.checkpoint_format == "gptq_v2": + raise ValueError("Currently, gptq_v2 is not supported on CPU with AMX.") -class GPTQLinearIntelAMXMethod(LinearMethodBase): - """Linear method for GPTQ on Intel CPU with AMX.""" +class GPTQIntelAMXLinearScheme(GPTQLinearScheme): + """Linear scheme for GPTQ on Intel CPU with AMX.""" - def __init__(self, quant_config: GPTQConfig): - self.quant_config = quant_config - # GPTQ v1 and v2 format deals with zero points differently - self.use_v2_format = quant_config.checkpoint_format == "gptq_v2" + def _init_kernel(self, quant_config: "GPTQConfig"): + return GPTQIntelAMXLinearKernel(quant_config) def create_weights( self, @@ -67,12 +55,12 @@ class GPTQLinearIntelAMXMethod(LinearMethodBase): 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, + **kwargs, ): - del output_size # Unused. - weight_loader = extra_weight_attrs.get("weight_loader") + _check_cpu_amx_support(self.quant_config) + if input_size_per_partition % self.quant_config.group_size != 0: raise ValueError( "The input size is not aligned with the quantized " @@ -87,17 +75,6 @@ class GPTQLinearIntelAMXMethod(LinearMethodBase): "tensor parallel size." ) - if self.quant_config.desc_act and not ( - self.quant_config.true_sequential and self.quant_config.static_groups - ): - raise ValueError( - "Currently, desc_act (True) is only supported with sequential and static group on CPU with AMX." - ) - if self.quant_config.weight_bits != 4: - raise ValueError("Currently, only 4bits is supported on CPU with AMX.") - if self.use_v2_format: - raise ValueError("Currently, gptq_v2 is not supported on CPU with AMX.") - if self.quant_config.group_size != -1: group_size = self.quant_config.group_size else: @@ -154,7 +131,6 @@ class GPTQLinearIntelAMXMethod(LinearMethodBase): packed_factor=self.quant_config.pack_factor, **qzeros_args, ) - else: scales = GroupQuantScaleParameter( output_dim=1, input_dim=0, **weight_scale_args @@ -172,34 +148,13 @@ class GPTQLinearIntelAMXMethod(LinearMethodBase): layer.register_parameter("qzeros", qzeros) layer.register_parameter("scales", scales) - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - _amx_process_weight_after_loading( - layer, ["qweight", "qzeros", "scales"], None, "gptq" - ) - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - return torch.ops.sgl_kernel.int4_scaled_mm_cpu( - x, - layer.qweight, - layer.qzeros, - layer.scales, - bias, - ) +class GPTQIntelAMXMoEScheme(GPTQMoESchemeBase): + """MoE scheme for GPTQ on Intel CPU with AMX.""" - -class GPTQMoEIntelAMXMethod(FusedMoEMethodBase): - """MoE method for GPTQ on Intel CPU with AMX.""" - - def __init__(self, quant_config: GPTQConfig): - super().__init__() + def __init__(self, quant_config: "GPTQConfig"): self.quant_config = quant_config - self.use_v2_format = quant_config.checkpoint_format == "gptq_v2" - self.moe_runner_config: Optional[MoeRunnerConfig] = None + self.kernel = GPTQIntelAMXMoEKernel(quant_config) def create_weights( self, @@ -210,20 +165,11 @@ class GPTQMoEIntelAMXMethod(FusedMoEMethodBase): params_dtype: torch.dtype, **extra_weight_attrs, ): - if self.quant_config.desc_act and not ( - self.quant_config.true_sequential and self.quant_config.static_groups - ): - raise ValueError( - "Currently, desc_act (True) is only supported with sequential and static group on CPU with AMX." - ) - if self.quant_config.weight_bits != 4: - raise ValueError("Currently, only 4bits is supported on CPU with AMX.") - if self.use_v2_format: - raise ValueError("Currently, gptq_v2 is not supported on CPU with AMX.") - # Delay the import to avoid circular dependency - from sglang.srt.layers.linear import set_weight_attrs from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + _check_cpu_amx_support(self.quant_config) + pack_factor = self.quant_config.pack_factor + if self.quant_config.group_size != -1: scales_size13 = hidden_size // self.quant_config.group_size w2_scales_size = intermediate_size_per_partition @@ -235,11 +181,11 @@ class GPTQMoEIntelAMXMethod(FusedMoEMethodBase): strategy = FusedMoeWeightScaleSupported.CHANNEL.value extra_weight_attrs.update({"quant_method": strategy, "is_transposed": True}) - # Fused gate_up_proj (column parallel) + w13_qweight = torch.nn.Parameter( torch.empty( num_experts, - hidden_size // self.quant_config.pack_factor, + hidden_size // pack_factor, 2 * intermediate_size_per_partition, dtype=torch.int32, ), @@ -247,11 +193,11 @@ class GPTQMoEIntelAMXMethod(FusedMoEMethodBase): ) layer.register_parameter("w13_qweight", w13_qweight) set_weight_attrs(w13_qweight, extra_weight_attrs) - # down_proj (row parallel) + w2_qweight = torch.nn.Parameter( torch.empty( num_experts, - intermediate_size_per_partition // self.quant_config.pack_factor, + intermediate_size_per_partition // pack_factor, hidden_size, dtype=torch.int32, ), @@ -259,7 +205,7 @@ class GPTQMoEIntelAMXMethod(FusedMoEMethodBase): ) layer.register_parameter("w2_qweight", w2_qweight) set_weight_attrs(w2_qweight, extra_weight_attrs) - # up_proj scales + w13_scales = torch.nn.Parameter( torch.empty( num_experts, @@ -271,51 +217,47 @@ class GPTQMoEIntelAMXMethod(FusedMoEMethodBase): ) layer.register_parameter("w13_scales", w13_scales) set_weight_attrs(w13_scales, extra_weight_attrs) - # down_proj scales + w2_scales = torch.nn.Parameter( torch.empty(num_experts, scales_size2, hidden_size, dtype=params_dtype), requires_grad=False, ) layer.register_parameter("w2_scales", w2_scales) set_weight_attrs(w2_scales, extra_weight_attrs) - # dont shard the w2 scales when running act order set_weight_attrs(w2_scales, {"load_full_w2": self.quant_config.desc_act}) - # up_proj scales + w13_qzeros = torch.nn.Parameter( torch.empty( num_experts, scales_size13, - 2 * intermediate_size_per_partition // self.quant_config.pack_factor, + 2 * intermediate_size_per_partition // pack_factor, dtype=torch.int32, ), requires_grad=False, ) layer.register_parameter("w13_qzeros", w13_qzeros) set_weight_attrs(w13_qzeros, extra_weight_attrs) - # down_proj scales + w2_qzeros = torch.nn.Parameter( torch.empty( num_experts, scales_size2, - hidden_size // self.quant_config.pack_factor, + hidden_size // pack_factor, dtype=torch.int32, ), requires_grad=False, ) layer.register_parameter("w2_qzeros", w2_qzeros) set_weight_attrs(w2_qzeros, extra_weight_attrs) - # dont shard the w2 scales when running act order set_weight_attrs(w2_qzeros, {"load_full_w2": self.quant_config.desc_act}) + w13_g_idx = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - dtype=torch.int32, - ), + torch.empty(num_experts, hidden_size, dtype=torch.int32), requires_grad=False, ) layer.register_parameter("w13_g_idx", w13_g_idx) set_weight_attrs(w13_g_idx, extra_weight_attrs) + w2_g_idx = torch.nn.Parameter( torch.empty( num_experts, @@ -328,52 +270,16 @@ class GPTQMoEIntelAMXMethod(FusedMoEMethodBase): set_weight_attrs(w2_g_idx, extra_weight_attrs) def create_moe_runner( - self, - layer: torch.nn.Module, - moe_runner_config: MoeRunnerConfig, - **extra_weight_attrs, + 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) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - _amx_process_weight_after_loading( - layer, ["w13_qweight", "w13_qzeros", "w13_scales"], None, "gptq" - ) - _amx_process_weight_after_loading( - layer, ["w2_qweight", "w2_qzeros", "w2_scales"], None, "gptq" - ) + self.kernel.process_weights_after_loading(layer) - def apply( + def apply_weights( 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 - output = torch.ops.sgl_kernel.fused_experts_cpu( - x, - layer.w13_qweight, - layer.w2_qweight, - topk_weights, - topk_ids, - False, # inplace See [Note] inplace should be False in fused_experts. - CPUQuantMethod.INT4_W4A8, - layer.w13_scales, # w1_scale - layer.w2_scales, # w2_scale - layer.w13_qzeros, - layer.w2_qzeros, - None, # block_size - None, # w1 bias - None, # w3 bias - None, # alpha - None, # limit - True, # is_vnni - ) - return StandardCombineInput(hidden_states=output) + dispatch_output: "StandardDispatchOutput", + ): + return self.kernel.apply(layer, dispatch_output) diff --git a/python/sglang/srt/layers/quantization/gptq/schemes/gptq_linear.py b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_linear.py index a02dd31b2..547d5e5e6 100644 --- a/python/sglang/srt/layers/quantization/gptq/schemes/gptq_linear.py +++ b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_linear.py @@ -5,7 +5,6 @@ from typing import TYPE_CHECKING, Optional import torch -from sglang.srt.hardware_backend.gpu.quantization.gptq_kernels import GPTQLinearKernel from sglang.srt.layers.parameter import ( ChannelQuantScaleParameter, GroupQuantScaleParameter, @@ -30,6 +29,10 @@ class GPTQLinearScheme(GPTQLinearSchemeBase): self.kernel = self._init_kernel(quant_config) def _init_kernel(self, quant_config: "GPTQConfig"): + from sglang.srt.hardware_backend.gpu.quantization.gptq_kernels import ( + GPTQLinearKernel, + ) + return GPTQLinearKernel(quant_config) def create_weights( @@ -157,12 +160,12 @@ class GPTQAscendLinearScheme(GPTQLinearScheme): return GPTQLinearAscendKernel(quant_config) def create_weights(self, layer: torch.nn.Module, **kwargs): - super().create_weights(layer=layer, **kwargs) - set_weight_attrs(layer.qzeros, {"pack_factor": self.quant_config.pack_factor}) - set_weight_attrs(layer.qweight, {"pack_factor": self.quant_config.pack_factor}) - if self.quant_config.desc_act: raise ValueError( "Currently, desc_act (True) is not supported by GPTQ " "quantization on npu." ) + + super().create_weights(layer=layer, **kwargs) + set_weight_attrs(layer.qzeros, {"pack_factor": self.quant_config.pack_factor}) + set_weight_attrs(layer.qweight, {"pack_factor": self.quant_config.pack_factor}) diff --git a/python/sglang/srt/layers/quantization/gptq/schemes/gptq_marlin.py b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_marlin.py index c630da829..d54c800df 100644 --- a/python/sglang/srt/layers/quantization/gptq/schemes/gptq_marlin.py +++ b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_marlin.py @@ -5,10 +5,6 @@ from typing import TYPE_CHECKING, Optional import torch -from sglang.srt.hardware_backend.gpu.quantization.gptq_kernels import ( - GPTQMarlinLinearKernel, - MarlinLinearLayerConfig, -) from sglang.srt.layers.parameter import ( ChannelQuantScaleParameter, GroupQuantScaleParameter, @@ -17,6 +13,7 @@ from sglang.srt.layers.parameter import ( RowvLLMParameter, ) from sglang.srt.layers.quantization.marlin_utils import ( + MarlinLinearLayerConfig, marlin_repeat_scales_on_all_ranks, verify_marlin_supported, ) @@ -32,13 +29,20 @@ __all__ = ["GPTQMarlinLinearScheme"] class GPTQMarlinLinearScheme(GPTQLinearSchemeBase): def __init__(self, quant_config: "GPTQMarlinConfig"): self.quant_config = quant_config - self.kernel = GPTQMarlinLinearKernel(quant_config) + self.kernel = self._init_kernel(quant_config) verify_marlin_supported( quant_type=self.quant_config.quant_type, group_size=self.quant_config.group_size, ) + def _init_kernel(self, quant_config: "GPTQMarlinConfig"): + from sglang.srt.hardware_backend.gpu.quantization.gptq_kernels import ( + GPTQMarlinLinearKernel, + ) + + return GPTQMarlinLinearKernel(quant_config) + def create_weights( self, layer: torch.nn.Module,