[GPTQ] Refactor CPU quantization schemes (#26786)

Co-authored-by: ronnie_zheng <zl19940307@163.com>

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)

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)

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}

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,
+            GPTQLinearMethod,
-            GPTQMoEAscendMethod,
+            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
 

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 (
+from sglang.srt.hardware_backend.cpu.quantization.awq_kernels import (
-    CPUQuantMethod,
+    AWQIntelAMXLinearKernel,
-    _amx_process_weight_after_loading,
+    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."""
 

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,

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,

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",
 ]

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

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",
 ]

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 (
+from sglang.srt.hardware_backend.cpu.quantization.gptq_kernels import (
-    MoeRunnerConfig,
+    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,
+from .gptq_linear import GPTQLinearScheme
-    LinearMethodBase,
+from .gptq_scheme import GPTQMoESchemeBase
-)
 
 if TYPE_CHECKING:
-    from sglang.srt.layers.moe.token_dispatcher import (
+    from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
-        StandardDispatchOutput,
+    from sglang.srt.layers.quantization.gptq.gptq import GPTQConfig
-    )
 
-from sglang.srt.layers.amx_utils import (
+__all__ = ["GPTQIntelAMXLinearScheme", "GPTQIntelAMXMoEScheme"]
-    CPUQuantMethod,
-    _amx_process_weight_after_loading,
-)
-
-from .gptq import GPTQConfig
 
 
-class CPUGPTQConfig(GPTQConfig):
+def _check_cpu_amx_support(quant_config: "GPTQConfig") -> None:
-    """CPU Config class for AWQ, inherit from AWQConfig"""
+    if quant_config.desc_act and not (
-
+        quant_config.true_sequential and quant_config.static_groups
-    @classmethod
+    ):
-    def get_supported_act_dtypes(cls) -> List[torch.dtype]:
+        raise ValueError(
-        return [torch.half, torch.bfloat16]
+            "Currently, desc_act (True) is only supported with sequential "
-
+            "and static group on CPU with AMX."
-    def get_quant_method(
+        )
-        self, layer: torch.nn.Module, prefix: str
+    if quant_config.weight_bits != 4:
-    ) -> Optional[LinearMethodBase]:
+        raise ValueError("Currently, only 4bits is supported on CPU with AMX.")
-        # Delay the import to avoid circular dependency
+    if quant_config.checkpoint_format == "gptq_v2":
-        from sglang.srt.layers.linear import LinearBase
+        raise ValueError("Currently, gptq_v2 is not supported on CPU with AMX.")
-        from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
-
-        if isinstance(layer, FusedMoE):
-            return GPTQMoEIntelAMXMethod(self)
-
-        if isinstance(layer, LinearBase):
-            return GPTQLinearIntelAMXMethod(self)
 
 
-class GPTQLinearIntelAMXMethod(LinearMethodBase):
+class GPTQIntelAMXLinearScheme(GPTQLinearScheme):
-    """Linear method for GPTQ on Intel CPU with AMX."""
+    """Linear scheme for GPTQ on Intel CPU with AMX."""
 
-    def __init__(self, quant_config: GPTQConfig):
+    def _init_kernel(self, quant_config: "GPTQConfig"):
-        self.quant_config = quant_config
+        return GPTQIntelAMXLinearKernel(quant_config)
-        # GPTQ v1 and v2 format deals with zero points differently
-        self.use_v2_format = quant_config.checkpoint_format == "gptq_v2"
 
     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.
+        _check_cpu_amx_support(self.quant_config)
-        weight_loader = extra_weight_attrs.get("weight_loader")
+
         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(
+class GPTQIntelAMXMoEScheme(GPTQMoESchemeBase):
-        self,
+    """MoE scheme for GPTQ on Intel CPU with AMX."""
-        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,
-        )
 
-
+    def __init__(self, quant_config: "GPTQConfig"):
-class GPTQMoEIntelAMXMethod(FusedMoEMethodBase):
-    """MoE method for GPTQ on Intel CPU with AMX."""
-
-    def __init__(self, quant_config: GPTQConfig):
-        super().__init__()
         self.quant_config = quant_config
-        self.use_v2_format = quant_config.checkpoint_format == "gptq_v2"
+        self.kernel = GPTQIntelAMXMoEKernel(quant_config)
-        self.moe_runner_config: Optional[MoeRunnerConfig] = None
 
     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(
+            torch.empty(num_experts, hidden_size, dtype=torch.int32),
-                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,
+        self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
-        layer: torch.nn.Module,
-        moe_runner_config: MoeRunnerConfig,
-        **extra_weight_attrs,
     ):
-        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(
+        self.kernel.process_weights_after_loading(layer)
-            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(
+    def apply_weights(
         self,
         layer: torch.nn.Module,
-        dispatch_output: StandardDispatchOutput,
+        dispatch_output: "StandardDispatchOutput",
-    ) -> torch.Tensor:
+    ):
-        from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
+        return self.kernel.apply(layer, dispatch_output)
-
-        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)

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

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,