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

[CPU] [Quantization] Add GPTQ/AWQ 4bits quantization support for CPU (#22685)

Browse Source 
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
jianan-gu
2026-04-22 13:34:02 -07:00
committed by GitHub
co-authored by Ma Mingfei
parent 0b77284587
commit ad0fc88810
14 changed files with 835 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/configs/update_config.py
+12 -7
View File
@@ -17,19 +17,24 @@ def may_get_weight_block_size(model_config, load_config):
if quant_config is not None and hasattr(quant_config, "weight_block_size"):
return getattr(quant_config, "weight_block_size")
if quant_config is not None and hasattr(quant_config, "group_size"):
return [getattr(quant_config, "group_size")]
return None
def get_moe_padding_size(weight_block_size):
if weight_block_size is not None:
# See NOTE(HandH1998): To ensure proper alignment of the block-wise quantization scales, the output_size of the weights for both the gate and up layers must be divisible by block_n.
assert (
len(weight_block_size) == 2
), "Only len(weight_block_size) == 2 is supported"
assert (
weight_block_size[0] == weight_block_size[1]
), "Only weight_block_size[0] == weight_block_size[1] is supported"
assert len(weight_block_size) in [
1,
2,
], "Only len(weight_block_size) in [1, 2] is supported"
if len(weight_block_size) == 2:
assert (
weight_block_size[0] == weight_block_size[1]
), "Only weight_block_size[0] == weight_block_size[1] is supported"
return weight_block_size[0]
return DEFAULT_MOE_PADDING_SIZE
python/sglang/srt/layers/amx_utils.py
Regular → Executable
+65 -31
View File
@@ -16,6 +16,11 @@ class CPUQuantMethod(IntEnum):
INT4_W4A8 = 3
class CPUQuantAlgo(IntEnum):
AWQ = 0
GPTQ = 1
def amx_process_weight_after_loading(weight, is_conv=False):
if weight.device != torch.device("cpu"):
return weight
@@ -74,7 +79,7 @@ def _init_amx_conv_state(conv_state):
def _amx_process_weight_after_loading(
module, weight_names, transpose_dims=None
module, weight_names, transpose_dims=None, qweight_packed_method=None
) -> None:
# Pack weight for get better performance on CPU
devices = {getattr(module, weight_name).device for weight_name in weight_names}
@@ -86,40 +91,69 @@ def _amx_process_weight_after_loading(
transpose_dims
), "len(weight_names) should be equal to len(transpose_dims)"
for i, weight_name in enumerate(weight_names):
weight_tensor = getattr(module, weight_name)
if transpose_dims and transpose_dims[i]:
weight_tensor = weight_tensor.transpose(*transpose_dims[i])
is_conv_weight = is_dim_conv_weight(weight_tensor)
# We don't pack weight or use intel amx backend if any weight of this module has unsupported dim.
if (
(not dim_is_supported(weight_tensor))
or not dtype_is_supported(weight_tensor)
) and (not is_conv_weight):
logger.warning(
f"Unsupported dimension or dtype for prepacking for weight '{weight_name}' with shape {weight_tensor.shape} and dtype {weight_tensor.dtype} in {module}. "
f"The derived (OC, IC) dimensions must be divisible by (16, 32). "
)
module.use_intel_amx_backend = False
return
packed_weight = torch.nn.Parameter(
amx_process_weight_after_loading(weight_tensor, is_conv_weight),
requires_grad=False,
)
packed_weight.__dict__ = weight_tensor.__dict__
setattr(module, weight_name, packed_weight)
if is_conv_weight:
# need to use inplace copy for conv weight amx packing,
# as its usage in radix_linear_attention will use the original conv weight.
weight_tensor = weight_tensor.view(-1, weight_tensor.size(-1))
weight_tensor.copy_(packed_weight)
module.use_intel_amx_backend = (
device == torch.device("cpu") and cpu_has_amx_support()
)
if qweight_packed_method is None:
for i, weight_name in enumerate(weight_names):
weight_tensor = getattr(module, weight_name)
if transpose_dims and transpose_dims[i]:
weight_tensor = weight_tensor.transpose(*transpose_dims[i])
is_conv_weight = is_dim_conv_weight(weight_tensor)
# We don't pack weight or use intel amx backend if any weight of this module has unsupported dim.
if (
(not dim_is_supported(weight_tensor))
or not dtype_is_supported(weight_tensor)
) and (not is_conv_weight):
logger.warning(
f"Unsupported dimension or dtype for prepacking for weight '{weight_name}' with shape {weight_tensor.shape} and dtype {weight_tensor.dtype} in {module}. "
f"The derived (OC, IC) dimensions must be divisible by (16, 32). "
)
module.use_intel_amx_backend = False
return
packed_weight = torch.nn.Parameter(
amx_process_weight_after_loading(weight_tensor, is_conv_weight),
requires_grad=False,
)
packed_weight.__dict__ = weight_tensor.__dict__
setattr(module, weight_name, packed_weight)
if is_conv_weight:
# need to use inplace copy for conv weight amx packing,
# as its usage in radix_linear_attention will use the original conv weight.
weight_tensor = weight_tensor.view(-1, weight_tensor.size(-1))
weight_tensor.copy_(packed_weight)
else:
assert qweight_packed_method in ["awq", "gptq"]
qweight_tensor = getattr(module, weight_names[0])
qzeros_tensor = getattr(module, weight_names[1])
scales_tensor = getattr(module, weight_names[2])
qweight, qzeros, scales = torch.ops.sgl_kernel.convert_weight_packed_scale_zp(
qweight_tensor,
qzeros_tensor,
scales_tensor,
CPUQuantAlgo.AWQ if qweight_packed_method == "awq" else CPUQuantAlgo.GPTQ,
)
packed_qweight = torch.nn.Parameter(
qweight.detach(),
requires_grad=False,
)
packed_qzeros = torch.nn.Parameter(
qzeros.detach(),
requires_grad=False,
)
packed_scales = torch.nn.Parameter(
scales.detach(),
requires_grad=False,
)
packed_qweight.__dict__ = qweight_tensor.__dict__
packed_qzeros.__dict__ = qzeros_tensor.__dict__
packed_scales.__dict__ = scales_tensor.__dict__
setattr(module, weight_names[0], packed_qweight)
setattr(module, weight_names[1], packed_qzeros)
setattr(module, weight_names[2], packed_scales)
if (
module.use_intel_amx_backend
and hasattr(module, "bias")
python/sglang/srt/layers/linear.py
+3
View File
@@ -57,6 +57,7 @@ WEIGHT_LOADER_V2_SUPPORTED = [
"AWQMarlinLinearMethod",
"AWQLinearMethod",
"AWQLinearAscendMethod",
"AWQLinearIntelAMXMethod",
"GPTQMarlinLinearMethod",
"Fp8LinearMethod",
"BlockInt8LinearMethod",
@@ -67,7 +68,9 @@ WEIGHT_LOADER_V2_SUPPORTED = [
"GPTQLinearMethod",
"FBGEMMFp8LinearMethod",
"GPTQLinearAscendMethod",
"GPTQLinearIntelAMXMethod",
"GPTQMoEAscendMethod",
"GPTQMoEIntelAMXMethod",
"ModelOptFp8LinearMethod",
"ModelOptFp4LinearMethod",
"IPEXAWQLinearMethod",
python/sglang/srt/layers/quantization/__init__.py
+28 -1
View File
@@ -18,6 +18,7 @@ CompressedTensorsConfig = DummyConfig
from sglang.srt.layers.quantization.auto_round import AutoRoundConfig
from sglang.srt.layers.quantization.awq import AWQConfig, AWQMarlinConfig
from sglang.srt.layers.quantization.awq_cpu import CPUAWQConfig
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.quantization.bitsandbytes import BitsAndBytesConfig
from sglang.srt.layers.quantization.blockwise_int8 import BlockInt8Config
@@ -28,6 +29,7 @@ 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 GPTQConfig, GPTQMarlinConfig
from sglang.srt.layers.quantization.gptq_cpu import CPUGPTQConfig
from sglang.srt.layers.quantization.modelopt_quant import (
ModelOptFp4Config,
ModelOptFp8Config,
@@ -43,7 +45,13 @@ from sglang.srt.layers.quantization.quark_int4fp8_moe import QuarkInt4Fp8Config
from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config
from sglang.srt.layers.quantization.w8a8_fp8 import W8A8Fp8Config
from sglang.srt.layers.quantization.w8a8_int8 import W8A8Int8Config
from sglang.srt.utils import is_cuda, is_hip, is_npu, mxfp_supported
from sglang.srt.utils import (
cpu_has_amx_support,
is_cuda,
is_hip,
is_npu,
mxfp_supported,
)
_is_mxfp_supported = mxfp_supported()
@@ -87,6 +95,15 @@ if is_cuda() or (_is_mxfp_supported and is_hip()):
}
)
# subset of above quant methods, supported on CPU
CPU_QUANTIZATION_METHODS = {
"fp8": Fp8Config,
"w8a8_int8": W8A8Int8Config,
"compressed-tensors": CompressedTensorsConfig,
"awq": CPUAWQConfig,
"gptq": CPUGPTQConfig,
}
QUANTIZATION_METHODS = {**BASE_QUANTIZATION_METHODS}
@@ -96,6 +113,16 @@ def get_quantization_config(quantization: str) -> Type[QuantizationConfig]:
f"Invalid quantization method: {quantization}. "
f"Available methods: {list(QUANTIZATION_METHODS.keys())}"
)
from sglang.srt.utils import is_cpu
if is_cpu() and cpu_has_amx_support():
if quantization not in CPU_QUANTIZATION_METHODS:
raise ValueError(
f"Invalid quantization method on CPU: {quantization}. "
f"Available methods on CPU: {list(QUANTIZATION_METHODS.keys())}"
)
else:
return CPU_QUANTIZATION_METHODS[quantization]
return QUANTIZATION_METHODS[quantization]
python/sglang/srt/layers/quantization/awq.py
Regular → Executable
View File
python/sglang/srt/layers/quantization/awq_cpu.py
Executable
+133
View File
@@ -0,0 +1,133 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import logging
from typing import TYPE_CHECKING, List, Optional
import torch
from sglang.srt.layers.moe import (
MoeRunnerConfig,
)
from sglang.srt.layers.quantization.base_config import (
LinearMethodBase,
)
from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
from sglang.srt.layers.quantization.utils import get_scalar_types
from .awq import AWQConfig, AWQLinearMethod, AWQMoEMethod
if TYPE_CHECKING:
from sglang.srt.layers.moe.token_dispatcher import (
StandardDispatchOutput,
)
from sglang.srt.layers.amx_utils import (
CPUQuantMethod,
_amx_process_weight_after_loading,
)
logger = logging.getLogger(__name__)
ScalarType, scalar_types = get_scalar_types()
def is_layer_skipped_awq(prefix: str, modules_to_not_convert: List[str]):
return any(module_name in prefix for module_name in modules_to_not_convert)
class CPUAWQConfig(AWQConfig):
"""CPU Config class for AWQ, inherit from AWQConfig"""
def get_supported_act_dtypes(self) -> List[torch.dtype]:
return [torch.float16, torch.bfloat16]
def get_quant_method(
self, layer: torch.nn.Module, prefix: str
) -> Optional[LinearMethodBase]:
from sglang.srt.layers.linear import LinearBase
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
if isinstance(layer, LinearBase):
if is_layer_skipped_awq(prefix, self.modules_to_not_convert):
return UnquantizedLinearMethod()
return AWQLinearIntelAMXMethod(self)
elif isinstance(layer, FusedMoE):
return AWQMoEIntelAMXMethod(self)
return None
class AWQLinearIntelAMXMethod(AWQLinearMethod):
"""Linear method for AWQ on Intel CPU with AMX."""
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 AWQMoEIntelAMXMethod(AWQMoEMethod):
"""MoE method for AWQ on Intel CPU with AMX."""
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
True, # is_vnni
)
return StandardCombineInput(hidden_states=output)
python/sglang/srt/layers/quantization/gptq.py
+10 -1
View File
@@ -129,6 +129,8 @@ class GPTQConfig(QuantizationConfig):
lm_head_quantized: bool,
dynamic: Dict[str, Dict[str, Union[int, bool]]],
checkpoint_format: str = "",
true_sequential: bool = False,
static_groups: bool = False,
) -> None:
# GPTQModel use `dynamic` config property to allow per module
# quantization config so each module can be individually optimized.
@@ -165,6 +167,8 @@ class GPTQConfig(QuantizationConfig):
# Currently GPTQModel stores v1 format checkpoints by default,
# but provides the option to set `format="gptq_v2"` in `QuantizeConfig`.
self.checkpoint_format = checkpoint_format
self.true_sequential = true_sequential
self.static_groups = static_groups
if self.weight_bits not in [2, 3, 4, 8]:
raise ValueError(
"Currently, only 2/3/4/8-bit weight quantization is "
@@ -222,6 +226,10 @@ class GPTQConfig(QuantizationConfig):
checkpoint_format = cls.get_from_keys_or(
config, ["checkpoint_format"], default=""
)
true_sequential = cls.get_from_keys_or(
config, ["true_sequential"], default=False
)
static_groups = cls.get_from_keys_or(config, ["static_groups"], default=False)
return cls(
weight_bits,
group_size,
@@ -229,6 +237,8 @@ class GPTQConfig(QuantizationConfig):
lm_head_quantized,
dynamic,
checkpoint_format,
true_sequential,
static_groups,
)
def get_quant_method(
@@ -477,7 +487,6 @@ class GPTQLinearMethod(LinearMethodBase):
group_size = self.quant_config.group_size
else:
group_size = input_size
self.use_shuffle = True
scale_and_zero_size = input_size // group_size
scale_and_zero_input_dim = None
python/sglang/srt/layers/quantization/gptq_cpu.py
+375
View File
@@ -0,0 +1,375 @@
from __future__ import annotations
from typing import TYPE_CHECKING, List, Optional
import torch
from sglang.srt.layers.moe import (
MoeRunnerConfig,
)
from sglang.srt.layers.parameter import (
ChannelQuantScaleParameter,
GroupQuantScaleParameter,
PackedColumnParameter,
PackedvLLMParameter,
RowvLLMParameter,
)
from sglang.srt.layers.quantization.base_config import (
FusedMoEMethodBase,
LinearMethodBase,
)
if TYPE_CHECKING:
from sglang.srt.layers.moe.token_dispatcher import (
StandardDispatchOutput,
)
from sglang.srt.layers.amx_utils import (
CPUQuantMethod,
_amx_process_weight_after_loading,
)
from .gptq import GPTQConfig
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)
class GPTQLinearIntelAMXMethod(LinearMethodBase):
"""Linear method 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 create_weights(
self,
layer: torch.nn.Module,
input_size_per_partition: int,
output_partition_sizes: list[int],
input_size: int,
output_size: int,
params_dtype: torch.dtype,
**extra_weight_attrs,
):
del output_size # Unused.
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 "
"weight shape. This can be caused by too large "
"tensor parallel size."
)
output_size_per_partition = sum(output_partition_sizes)
if output_size_per_partition % self.quant_config.pack_factor.numerator != 0:
raise ValueError(
"The output size is not aligned with the quantized "
"weight shape. This can be caused by too large "
"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:
group_size = input_size
scale_and_zero_size = input_size_per_partition // group_size
scale_and_zero_input_dim = 0
qweight = PackedvLLMParameter(
data=torch.empty(
input_size_per_partition // self.quant_config.pack_factor,
output_size_per_partition,
dtype=torch.int32,
),
input_dim=0,
output_dim=1,
packed_dim=0,
packed_factor=self.quant_config.pack_factor,
weight_loader=weight_loader,
)
g_idx = RowvLLMParameter(
data=torch.tensor(
[
i // self.quant_config.group_size
for i in range(input_size_per_partition)
],
dtype=torch.int32,
),
input_dim=0,
weight_loader=weight_loader,
)
qzeros_args = {
"data": torch.empty(
scale_and_zero_size,
output_size_per_partition // self.quant_config.pack_factor,
dtype=torch.int32,
),
"weight_loader": weight_loader,
}
weight_scale_args = {
"data": torch.empty(
scale_and_zero_size,
output_size_per_partition,
dtype=params_dtype,
),
"weight_loader": weight_loader,
}
if scale_and_zero_input_dim is None:
scales = ChannelQuantScaleParameter(output_dim=1, **weight_scale_args)
qzeros = PackedColumnParameter(
output_dim=1,
packed_dim=1,
packed_factor=self.quant_config.pack_factor,
**qzeros_args,
)
else:
scales = GroupQuantScaleParameter(
output_dim=1, input_dim=0, **weight_scale_args
)
qzeros = PackedvLLMParameter(
input_dim=0,
output_dim=1,
packed_dim=1,
packed_factor=self.quant_config.pack_factor,
**qzeros_args,
)
layer.register_parameter("qweight", qweight)
layer.register_parameter("g_idx", g_idx)
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 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.moe_runner_config: Optional[MoeRunnerConfig] = None
def create_weights(
self,
layer: torch.nn.Module,
num_experts: int,
hidden_size: int,
intermediate_size_per_partition: int,
params_dtype: torch.dtype,
**extra_weight_attrs,
):
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
if self.quant_config.group_size != -1:
scales_size13 = hidden_size // self.quant_config.group_size
w2_scales_size = intermediate_size_per_partition
scales_size2 = w2_scales_size // self.quant_config.group_size
strategy = FusedMoeWeightScaleSupported.GROUP.value
else:
scales_size13 = 1
scales_size2 = 1
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,
2 * intermediate_size_per_partition,
dtype=torch.int32,
),
requires_grad=False,
)
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,
hidden_size,
dtype=torch.int32,
),
requires_grad=False,
)
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,
scales_size13,
2 * intermediate_size_per_partition,
dtype=params_dtype,
),
requires_grad=False,
)
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,
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,
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,
),
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,
intermediate_size_per_partition,
dtype=torch.int32,
),
requires_grad=False,
)
layer.register_parameter("w2_g_idx", w2_g_idx)
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.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
True, # is_vnni
)
return StandardCombineInput(hidden_states=output)