minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[GPTQ] Refactor CPU quantization schemes (#26786)

Browse Source 
Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
Yaochen Han
2026-06-08 13:14:40 +03:00
committed by GitHub
co-authored by ronnie_zheng
parent a26587dd4e
commit 1f5dc2cdca
13 changed files with 348 additions and 278 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/cpu/quantization/awq_kernels.py
+99
View File
@@ -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
+99
View File
@@ -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
+2 -4
View File
@@ -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
+6 -4
View File
@@ -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
python/sglang/srt/layers/quantization/awq/schemes/awq_cpu.py
+4 -85
View File
@@ -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."""
python/sglang/srt/layers/quantization/awq/schemes/awq_marlin.py
+8 -4
View File
@@ -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
+3 -1
View File
@@ -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
+8 -9
View File
@@ -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
+39 -7
View File
@@ -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
+3
View File
@@ -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_cpu.py → python/sglang/srt/layers/quantization/gptq/schemes/gptq_cpu.py
+60 -154
View File
@@ -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)
python/sglang/srt/layers/quantization/gptq/schemes/gptq_linear.py
+8 -5
View File
@@ -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
+9 -5
View File
@@ -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,