134 lines
4.1 KiB
Python
Executable File
134 lines
4.1 KiB
Python
Executable File
# 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)
|