[4/N] Quantization Refactor: AWQ schemes and Kernel call and weight init split (#21126)
This commit is contained in:
Yaochen Han
@@ -0,0 +1,255 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.moe import MoeRunner
|
||||
from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo
|
||||
from sglang.srt.layers.quantization.marlin_utils import (
|
||||
apply_awq_marlin_linear,
|
||||
awq_to_marlin_zero_points,
|
||||
marlin_make_empty_g_idx,
|
||||
marlin_make_workspace,
|
||||
marlin_moe_permute_scales,
|
||||
marlin_permute_scales,
|
||||
moe_awq_to_marlin_zero_points,
|
||||
)
|
||||
from sglang.srt.layers.quantization.utils import get_scalar_types, replace_parameter
|
||||
from sglang.srt.utils import is_hip, is_xpu
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.moe.token_dispatcher import (
|
||||
CombineInput,
|
||||
StandardDispatchOutput,
|
||||
)
|
||||
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
||||
|
||||
awq_marlin_moe_repack = None
|
||||
awq_marlin_repack = None
|
||||
|
||||
|
||||
def _unsupported_awq_dequantize(*args, **kwargs):
|
||||
raise RuntimeError("AWQ GPU kernels are unavailable on the current platform.")
|
||||
|
||||
|
||||
awq_dequantize = _unsupported_awq_dequantize
|
||||
|
||||
if is_xpu():
|
||||
try:
|
||||
from sgl_kernel import awq_dequantize
|
||||
except ImportError:
|
||||
pass
|
||||
elif is_hip():
|
||||
try:
|
||||
from sglang.srt.layers.quantization.awq.awq_triton import (
|
||||
awq_dequantize_triton as awq_dequantize,
|
||||
)
|
||||
except ImportError:
|
||||
pass
|
||||
else:
|
||||
try:
|
||||
from sglang.jit_kernel.awq_dequantize import awq_dequantize
|
||||
from sglang.jit_kernel.awq_marlin_repack import (
|
||||
awq_marlin_moe_repack,
|
||||
awq_marlin_repack,
|
||||
)
|
||||
from sglang.srt.utils.custom_op import register_custom_op_from_extern
|
||||
|
||||
awq_dequantize = register_custom_op_from_extern(
|
||||
awq_dequantize,
|
||||
fake_impl=lambda qweight, scales, qzeros: qweight.new_empty(
|
||||
qweight.shape[:-1] + (qweight.shape[-1] * 8,), dtype=scales.dtype
|
||||
),
|
||||
)
|
||||
except ImportError:
|
||||
try:
|
||||
from sglang.srt.layers.quantization.awq.awq_triton import (
|
||||
awq_dequantize_triton as awq_dequantize,
|
||||
)
|
||||
except ImportError:
|
||||
try:
|
||||
from sgl_kernel import awq_dequantize
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
_, scalar_types = get_scalar_types()
|
||||
|
||||
|
||||
class AWQLinearKernel:
|
||||
def __init__(self, quant_config: Optional["QuantizationConfig"] = None):
|
||||
self.quant_config = quant_config
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
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:
|
||||
qweight = layer.qweight
|
||||
scales = layer.scales
|
||||
qzeros = layer.qzeros
|
||||
pack_factor = self.quant_config.pack_factor
|
||||
out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,)
|
||||
reshaped_x = x.reshape(-1, x.shape[-1])
|
||||
out = awq_dequantize(qweight, scales, qzeros)
|
||||
out = torch.matmul(reshaped_x, out)
|
||||
|
||||
if bias is not None:
|
||||
out.add_(bias)
|
||||
return out.reshape(out_shape)
|
||||
|
||||
|
||||
class AWQMarlinLinearKernel:
|
||||
def __init__(self, quant_config: Optional["QuantizationConfig"] = None):
|
||||
self.quant_config = quant_config
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
device = layer.qweight.device
|
||||
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)
|
||||
|
||||
layer.workspace = marlin_make_workspace(device)
|
||||
|
||||
marlin_qweight = awq_marlin_repack(
|
||||
layer.qweight,
|
||||
size_k=layer.input_size_per_partition,
|
||||
size_n=layer.output_size_per_partition,
|
||||
num_bits=self.quant_config.quant_type.size_bits,
|
||||
)
|
||||
replace_parameter(layer, "qweight", marlin_qweight)
|
||||
|
||||
marlin_scales = marlin_permute_scales(
|
||||
layer.scales,
|
||||
size_k=layer.input_size_per_partition,
|
||||
size_n=layer.output_size_per_partition,
|
||||
group_size=self.quant_config.group_size,
|
||||
)
|
||||
replace_parameter(layer, "scales", marlin_scales)
|
||||
|
||||
marlin_zp = awq_to_marlin_zero_points(
|
||||
layer.qzeros,
|
||||
size_k=layer.num_groups,
|
||||
size_n=layer.output_size_per_partition,
|
||||
num_bits=self.quant_config.quant_type.size_bits,
|
||||
)
|
||||
replace_parameter(layer, "qzeros", marlin_zp)
|
||||
|
||||
layer.g_idx = marlin_make_empty_g_idx(device)
|
||||
layer.g_idx_sort_indices = marlin_make_empty_g_idx(device)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
return apply_awq_marlin_linear(
|
||||
input=x,
|
||||
weight=layer.qweight,
|
||||
weight_scale=layer.scales,
|
||||
weight_zp=layer.qzeros,
|
||||
g_idx=layer.g_idx,
|
||||
g_idx_sort_indices=layer.g_idx_sort_indices,
|
||||
workspace=layer.workspace,
|
||||
quant_type=self.quant_config.quant_type,
|
||||
output_size_per_partition=layer.output_size_per_partition,
|
||||
input_size_per_partition=layer.input_size_per_partition,
|
||||
bias=bias,
|
||||
)
|
||||
|
||||
|
||||
class AWQMoEKernel:
|
||||
def __init__(self, quant_config: Optional["QuantizationConfig"] = None):
|
||||
self.quant_config = quant_config
|
||||
self.runner: Optional[MoeRunner] = None
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
num_experts = layer.w13_qweight.shape[0]
|
||||
device = layer.w13_qweight.device
|
||||
|
||||
layer.w13_g_idx_sort_indices = torch.nn.Parameter(
|
||||
torch.empty((num_experts, 0), dtype=torch.int32, device=device),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w2_g_idx_sort_indices = torch.nn.Parameter(
|
||||
torch.empty((num_experts, 0), dtype=torch.int32, device=device),
|
||||
requires_grad=False,
|
||||
)
|
||||
|
||||
marlin_w13_qweight = awq_marlin_moe_repack(
|
||||
layer.w13_qweight,
|
||||
layer.w13_g_idx_sort_indices,
|
||||
size_k=layer.w13_qweight.shape[1],
|
||||
size_n=layer.w13_qweight.shape[2] * self.quant_config.pack_factor,
|
||||
num_bits=self.quant_config.weight_bits,
|
||||
)
|
||||
replace_parameter(layer, "w13_qweight", marlin_w13_qweight)
|
||||
|
||||
marlin_w2_qweight = awq_marlin_moe_repack(
|
||||
layer.w2_qweight,
|
||||
layer.w2_g_idx_sort_indices,
|
||||
size_k=layer.w2_qweight.shape[1],
|
||||
size_n=layer.w2_qweight.shape[2] * self.quant_config.pack_factor,
|
||||
num_bits=self.quant_config.weight_bits,
|
||||
)
|
||||
replace_parameter(layer, "w2_qweight", marlin_w2_qweight)
|
||||
|
||||
marlin_w13_scales = marlin_moe_permute_scales(
|
||||
s=layer.w13_scales,
|
||||
size_k=layer.intermediate_size_per_partition,
|
||||
size_n=layer.w13_scales.shape[2],
|
||||
group_size=self.quant_config.group_size,
|
||||
)
|
||||
replace_parameter(layer, "w13_scales", marlin_w13_scales)
|
||||
|
||||
marlin_w2_scales = marlin_moe_permute_scales(
|
||||
s=layer.w2_scales,
|
||||
size_k=layer.intermediate_size_per_partition,
|
||||
size_n=layer.w2_scales.shape[2],
|
||||
group_size=self.quant_config.group_size,
|
||||
)
|
||||
replace_parameter(layer, "w2_scales", marlin_w2_scales)
|
||||
|
||||
marlin_w13_zp = moe_awq_to_marlin_zero_points(
|
||||
layer.w13_qzeros,
|
||||
size_k=layer.w13_qzeros.shape[1],
|
||||
size_n=layer.w13_qzeros.shape[2] * self.quant_config.pack_factor,
|
||||
num_bits=self.quant_config.weight_bits,
|
||||
)
|
||||
replace_parameter(layer, "w13_qzeros", marlin_w13_zp)
|
||||
|
||||
marlin_w2_zp = moe_awq_to_marlin_zero_points(
|
||||
layer.w2_qzeros,
|
||||
size_k=layer.w2_qzeros.shape[1],
|
||||
size_n=layer.w2_qzeros.shape[2] * self.quant_config.pack_factor,
|
||||
num_bits=self.quant_config.weight_bits,
|
||||
)
|
||||
replace_parameter(layer, "w2_qzeros", marlin_w2_zp)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
dispatch_output: "StandardDispatchOutput",
|
||||
) -> "CombineInput":
|
||||
if self.runner is None:
|
||||
raise RuntimeError("moe runner is not initialized")
|
||||
|
||||
quant_info = MarlinMoeQuantInfo(
|
||||
w13_qweight=layer.w13_qweight,
|
||||
w2_qweight=layer.w2_qweight,
|
||||
w13_scales=layer.w13_scales,
|
||||
w2_scales=layer.w2_scales,
|
||||
w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices,
|
||||
w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices,
|
||||
w13_qzeros=layer.w13_qzeros,
|
||||
w2_qzeros=layer.w2_qzeros,
|
||||
weight_bits=self.quant_config.weight_bits,
|
||||
)
|
||||
return self.runner.run(dispatch_output, quant_info)
|
||||
@@ -0,0 +1,156 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
|
||||
NPUW4A16Int4DynamicMoEMethod,
|
||||
)
|
||||
from sglang.srt.layers.quantization.utils import replace_parameter
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
|
||||
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
||||
|
||||
import torch_npu
|
||||
|
||||
|
||||
class AWQAscendLinearKernel:
|
||||
def __init__(self, quant_config: Optional["QuantizationConfig"] = None):
|
||||
self.quant_config = quant_config
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
|
||||
qweight_tmp = torch.zeros_like(layer.qweight.data)
|
||||
qzeros_tmp = layer.qzeros.data
|
||||
qzeros_list = []
|
||||
shifts = [0, 4, 1, 5, 2, 6, 3, 7]
|
||||
|
||||
for i in range(0, self.quant_config.pack_factor):
|
||||
shift_num = shifts[i] * 4
|
||||
qzeros_list.append((qzeros_tmp.reshape(-1, 1) >> shift_num) & 0xF)
|
||||
qweight_tmp.bitwise_or_(
|
||||
((layer.qweight.data >> shift_num) & 0xF) << (4 * i)
|
||||
)
|
||||
|
||||
qweight_tmp.bitwise_xor_(0x88888888)
|
||||
|
||||
qzeros_tmp = torch.cat(qzeros_list, dim=-1).reshape(qzeros_tmp.shape[0], -1)
|
||||
qzeros_tmp = -(qzeros_tmp - 8)
|
||||
qzeros_tmp = qzeros_tmp.to(layer.scales.data.dtype)
|
||||
|
||||
layer.zeros = torch.nn.Parameter(qzeros_tmp, requires_grad=False)
|
||||
layer.weight = torch.nn.Parameter(qweight_tmp, requires_grad=False)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
qweight = layer.weight
|
||||
scales = layer.scales
|
||||
qzeros = layer.zeros
|
||||
pack_factor = self.quant_config.pack_factor
|
||||
out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,)
|
||||
reshaped_x = x.reshape(-1, x.shape[-1])
|
||||
|
||||
if bias is not None and bias.dtype == torch.bfloat16:
|
||||
bias = bias.float()
|
||||
|
||||
out = torch_npu.npu_weight_quant_batchmatmul(
|
||||
reshaped_x,
|
||||
qweight,
|
||||
antiquant_scale=scales,
|
||||
antiquant_offset=qzeros,
|
||||
antiquant_group_size=self.quant_config.group_size,
|
||||
bias=bias,
|
||||
)
|
||||
|
||||
return out.reshape(out_shape)
|
||||
|
||||
|
||||
class AWQAscendMoEKernel:
|
||||
def __init__(self, quant_config: Optional["QuantizationConfig"] = None):
|
||||
self.quant_config = quant_config
|
||||
self.kernel = NPUW4A16Int4DynamicMoEMethod()
|
||||
|
||||
@staticmethod
|
||||
def _register_or_replace_parameter(
|
||||
layer: torch.nn.Module, name: str, tensor: torch.Tensor
|
||||
) -> None:
|
||||
if hasattr(layer, name):
|
||||
replace_parameter(layer, name, tensor)
|
||||
else:
|
||||
layer.register_parameter(
|
||||
name, torch.nn.Parameter(tensor, requires_grad=False)
|
||||
)
|
||||
|
||||
def _convert_awq_weight_to_npu_layout(self, qweight: torch.Tensor) -> torch.Tensor:
|
||||
num_experts, input_size, _ = qweight.shape
|
||||
unpacked_weight = (
|
||||
self.kernel._unpack_from_int32(qweight.flatten(0, 1), 4)
|
||||
.view(num_experts, input_size, -1)
|
||||
.transpose(1, 2)
|
||||
.contiguous()
|
||||
.int()
|
||||
)
|
||||
return self.kernel._pack_to_int32(unpacked_weight)
|
||||
|
||||
def _convert_awq_qzeros_to_npu_offset(
|
||||
self, qzeros: torch.Tensor, dtype: torch.dtype
|
||||
) -> torch.Tensor:
|
||||
num_experts, num_groups, _ = qzeros.shape
|
||||
offset = (
|
||||
-self.kernel._unpack_from_int32(qzeros.flatten(0, 1), 4)
|
||||
.view(num_experts, num_groups, -1)
|
||||
.transpose(1, 2)
|
||||
.contiguous()
|
||||
)
|
||||
return offset.to(dtype)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
self._register_or_replace_parameter(
|
||||
layer,
|
||||
"w13_weight",
|
||||
self._convert_awq_weight_to_npu_layout(layer.w13_qweight.data),
|
||||
)
|
||||
self._register_or_replace_parameter(
|
||||
layer,
|
||||
"w2_weight",
|
||||
self._convert_awq_weight_to_npu_layout(layer.w2_qweight.data),
|
||||
)
|
||||
self._register_or_replace_parameter(
|
||||
layer,
|
||||
"w13_weight_scale",
|
||||
layer.w13_scales.data.transpose(1, 2).contiguous(),
|
||||
)
|
||||
self._register_or_replace_parameter(
|
||||
layer,
|
||||
"w2_weight_scale",
|
||||
layer.w2_scales.data.transpose(1, 2).contiguous(),
|
||||
)
|
||||
self._register_or_replace_parameter(
|
||||
layer,
|
||||
"w13_weight_offset",
|
||||
self._convert_awq_qzeros_to_npu_offset(
|
||||
layer.w13_qzeros.data, layer.w13_scales.data.dtype
|
||||
),
|
||||
)
|
||||
self._register_or_replace_parameter(
|
||||
layer,
|
||||
"w2_weight_offset",
|
||||
self._convert_awq_qzeros_to_npu_offset(
|
||||
layer.w2_qzeros.data, layer.w2_scales.data.dtype
|
||||
),
|
||||
)
|
||||
|
||||
self.kernel.process_weights_after_loading(layer)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
dispatch_output: "StandardDispatchOutput",
|
||||
) -> torch.Tensor:
|
||||
return self.kernel.apply(layer, dispatch_output)
|
||||
@@ -54,10 +54,7 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
WEIGHT_LOADER_V2_SUPPORTED = [
|
||||
"CompressedTensorsLinearMethod",
|
||||
"AWQMarlinLinearMethod",
|
||||
"AWQLinearMethod",
|
||||
"AWQLinearAscendMethod",
|
||||
"AWQLinearIntelAMXMethod",
|
||||
"GPTQMarlinLinearMethod",
|
||||
"Fp8LinearMethod",
|
||||
"BlockInt8LinearMethod",
|
||||
|
||||
@@ -17,8 +17,7 @@ class DummyConfig:
|
||||
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.awq import AWQConfig, AWQCPUConfig, AWQMarlinConfig
|
||||
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
|
||||
@@ -100,7 +99,7 @@ CPU_QUANTIZATION_METHODS = {
|
||||
"fp8": Fp8Config,
|
||||
"w8a8_int8": W8A8Int8Config,
|
||||
"compressed-tensors": CompressedTensorsConfig,
|
||||
"awq": CPUAWQConfig,
|
||||
"awq": AWQCPUConfig,
|
||||
"gptq": CPUGPTQConfig,
|
||||
}
|
||||
|
||||
|
||||
@@ -258,8 +258,8 @@ class AutoRoundConfig(QuantizationConfig):
|
||||
use_marlin = False
|
||||
if use_marlin:
|
||||
from sglang.srt.layers.quantization.awq import (
|
||||
AWQLinearMethod,
|
||||
AWQMarlinConfig,
|
||||
AWQMarlinLinearMethod,
|
||||
AWQMoEMethod,
|
||||
)
|
||||
|
||||
@@ -282,6 +282,7 @@ class AutoRoundConfig(QuantizationConfig):
|
||||
|
||||
if isinstance(layer, FusedMoE):
|
||||
if use_marlin:
|
||||
layer.scheme = quant_args_marlin.get_moe_scheme(layer)
|
||||
return AWQMoEMethod(quant_args_marlin)
|
||||
from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config
|
||||
|
||||
@@ -296,8 +297,10 @@ class AutoRoundConfig(QuantizationConfig):
|
||||
|
||||
if isinstance(layer, (LinearBase, ParallelLMHead)):
|
||||
if use_marlin:
|
||||
return AWQMarlinLinearMethod(quant_args_marlin)
|
||||
layer.scheme = quant_args_marlin.get_linear_scheme(layer)
|
||||
return AWQLinearMethod(quant_args_marlin)
|
||||
else:
|
||||
layer.scheme = quant_args.get_linear_scheme(layer)
|
||||
return AWQLinearMethod(quant_args)
|
||||
return None
|
||||
|
||||
|
||||
@@ -1,966 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import warnings
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
|
||||
npu_fused_experts,
|
||||
)
|
||||
from sglang.srt.layers.linear import LinearBase, set_weight_attrs
|
||||
from sglang.srt.layers.moe import (
|
||||
MoeRunner,
|
||||
MoeRunnerBackend,
|
||||
MoeRunnerConfig,
|
||||
get_moe_runner_backend,
|
||||
)
|
||||
from sglang.srt.layers.moe.moe_runner.marlin import MarlinMoeQuantInfo
|
||||
from sglang.srt.layers.parameter import GroupQuantScaleParameter, PackedvLLMParameter
|
||||
from sglang.srt.layers.quantization.base_config import (
|
||||
FusedMoEMethodBase,
|
||||
LinearMethodBase,
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from sglang.srt.layers.quantization.marlin_utils import (
|
||||
apply_awq_marlin_linear,
|
||||
awq_to_marlin_zero_points,
|
||||
check_marlin_supported,
|
||||
check_marlin_supports_layer,
|
||||
check_moe_marlin_supports_layer,
|
||||
marlin_make_empty_g_idx,
|
||||
marlin_make_workspace,
|
||||
marlin_moe_permute_scales,
|
||||
marlin_permute_scales,
|
||||
moe_awq_to_marlin_zero_points,
|
||||
verify_marlin_supported,
|
||||
verify_marlin_supports_shape,
|
||||
)
|
||||
from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
|
||||
from sglang.srt.layers.quantization.utils import get_scalar_types, replace_parameter
|
||||
from sglang.srt.utils.patch_torch import register_fake_if_exists
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.moe.token_dispatcher import (
|
||||
CombineInput,
|
||||
StandardDispatchOutput,
|
||||
)
|
||||
|
||||
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
|
||||
|
||||
_is_cuda = is_cuda()
|
||||
_is_hip = is_hip()
|
||||
_is_xpu = is_xpu()
|
||||
_is_npu = is_npu()
|
||||
|
||||
if _is_npu:
|
||||
import torch_npu
|
||||
|
||||
if _is_cuda:
|
||||
from sglang.jit_kernel.awq_dequantize import awq_dequantize
|
||||
from sglang.jit_kernel.awq_marlin_repack import (
|
||||
awq_marlin_moe_repack,
|
||||
awq_marlin_repack,
|
||||
)
|
||||
from sglang.srt.utils.custom_op import register_custom_op_from_extern
|
||||
|
||||
awq_dequantize = register_custom_op_from_extern(
|
||||
awq_dequantize,
|
||||
fake_impl=lambda qweight, scales, qzeros: qweight.new_empty(
|
||||
qweight.shape[:-1] + (qweight.shape[-1] * 8,), dtype=scales.dtype
|
||||
),
|
||||
)
|
||||
|
||||
elif _is_hip:
|
||||
from sglang.srt.layers.quantization.awq_triton import (
|
||||
awq_dequantize_triton as awq_dequantize,
|
||||
)
|
||||
|
||||
elif _is_xpu:
|
||||
from sgl_kernel import awq_dequantize
|
||||
|
||||
warnings.warn(f"XPU does not support fused_marlin_moe currently.")
|
||||
else:
|
||||
warnings.warn(f"Only CUDA, HIP and XPU support AWQ currently.")
|
||||
|
||||
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 AWQConfig(QuantizationConfig):
|
||||
"""Config class for AWQ.
|
||||
|
||||
Reference: https://arxiv.org/abs/2306.00978
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
weight_bits: int,
|
||||
group_size: int,
|
||||
zero_point: bool,
|
||||
modules_to_not_convert: Optional[List[str]] = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.weight_bits = weight_bits
|
||||
self.group_size = group_size
|
||||
self.zero_point = zero_point
|
||||
self.modules_to_not_convert = modules_to_not_convert or []
|
||||
|
||||
if self.weight_bits != 4:
|
||||
raise ValueError(
|
||||
"Currently, only 4-bit weight quantization is supported for "
|
||||
f"AWQ, but got {self.weight_bits} bits."
|
||||
)
|
||||
self.pack_factor = 32 // self.weight_bits
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"AWQConfig(weight_bits={self.weight_bits}, "
|
||||
f"group_size={self.group_size}, "
|
||||
f"zero_point={self.zero_point}, "
|
||||
f"modules_to_not_convert={self.modules_to_not_convert})"
|
||||
)
|
||||
|
||||
def get_scaled_act_names(self) -> List[str]:
|
||||
return []
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "awq"
|
||||
|
||||
def get_supported_act_dtypes(self) -> List[torch.dtype]:
|
||||
return [torch.float16] if not _is_npu else [torch.float16, torch.bfloat16]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
# The AWQ kernel only supports Turing or newer GPUs.
|
||||
if _is_npu:
|
||||
raise NotImplementedError(
|
||||
'NPU hardware does not support "get_min_capability" feature.'
|
||||
)
|
||||
else:
|
||||
return 75
|
||||
|
||||
@staticmethod
|
||||
def get_config_filenames() -> List[str]:
|
||||
return [
|
||||
"quant_config.json", # E.g., casperhansen/vicuna-7b-v1.5-awq
|
||||
# E.g., abhinavkulkarni/mosaicml-mpt-7b-instruct-w4-g128-awq
|
||||
"quantize_config.json",
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: Dict[str, Any]) -> AWQConfig:
|
||||
weight_bits = cls.get_from_keys(config, ["w_bit", "bits"])
|
||||
group_size = cls.get_from_keys(config, ["q_group_size", "group_size"])
|
||||
zero_point = cls.get_from_keys(config, ["zero_point"])
|
||||
modules_to_not_convert = cls.get_from_keys_or(
|
||||
config, ["modules_to_not_convert"], None
|
||||
)
|
||||
return cls(weight_bits, group_size, zero_point, modules_to_not_convert)
|
||||
|
||||
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 _is_npu:
|
||||
if isinstance(layer, LinearBase):
|
||||
if is_layer_skipped_awq(prefix, self.modules_to_not_convert):
|
||||
return UnquantizedLinearMethod()
|
||||
return AWQLinearAscendMethod(self)
|
||||
elif isinstance(layer, FusedMoE):
|
||||
return AWQMoEAscendMethod(self)
|
||||
return None
|
||||
|
||||
if isinstance(layer, LinearBase):
|
||||
if is_layer_skipped_awq(prefix, self.modules_to_not_convert):
|
||||
return UnquantizedLinearMethod()
|
||||
return AWQLinearMethod(self)
|
||||
return None
|
||||
|
||||
|
||||
class AWQMarlinConfig(QuantizationConfig):
|
||||
"""Config class for AWQ Marlin"""
|
||||
|
||||
# num_bits -> type
|
||||
TYPE_MAP = {
|
||||
4: scalar_types.uint4,
|
||||
8: scalar_types.uint8,
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
weight_bits: int,
|
||||
group_size: int,
|
||||
zero_point: bool,
|
||||
lm_head_quantized: bool,
|
||||
modules_to_not_convert: Optional[list[str]],
|
||||
full_config: dict[str, Any],
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if _is_hip:
|
||||
warnings.warn(f"HIP does not support fused_marlin_moe currently.")
|
||||
self.pack_factor = 32 // weight_bits # packed into int32
|
||||
self.group_size = group_size
|
||||
self.zero_point = zero_point
|
||||
self.lm_head_quantized = lm_head_quantized
|
||||
self.weight_bits = weight_bits
|
||||
self.modules_to_not_convert = modules_to_not_convert or []
|
||||
self.full_config = full_config
|
||||
|
||||
if self.weight_bits not in self.TYPE_MAP:
|
||||
raise ValueError(
|
||||
f"Unsupported num_bits = {self.weight_bits}. "
|
||||
f"Supported num_bits = {self.TYPE_MAP.keys()}"
|
||||
)
|
||||
|
||||
self.quant_type = self.TYPE_MAP[self.weight_bits]
|
||||
|
||||
verify_marlin_supported(
|
||||
self.quant_type, group_size=self.group_size, has_zp=self.zero_point
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"AWQMarlinConfig(quant_type={self.quant_type}, "
|
||||
f"group_size={self.group_size}, "
|
||||
f"zero_point={self.zero_point}, "
|
||||
f"lm_head_quantized={self.lm_head_quantized}, "
|
||||
f"modules_to_not_convert={self.modules_to_not_convert})"
|
||||
)
|
||||
|
||||
def get_scaled_act_names(self) -> List[str]:
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> str:
|
||||
return "awq_marlin"
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||||
return [torch.half, torch.bfloat16]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 80
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> list[str]:
|
||||
return ["quantize_config.json"]
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> AWQMarlinConfig:
|
||||
weight_bits = cls.get_from_keys(config, ["bits"])
|
||||
group_size = cls.get_from_keys(config, ["group_size"])
|
||||
zero_point = cls.get_from_keys(config, ["zero_point"])
|
||||
lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False)
|
||||
modules_to_not_convert = cls.get_from_keys_or(
|
||||
config, ["modules_to_not_convert"], None
|
||||
)
|
||||
return cls(
|
||||
weight_bits,
|
||||
group_size,
|
||||
zero_point,
|
||||
lm_head_quantized,
|
||||
modules_to_not_convert,
|
||||
config,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def override_quantization_method(cls, hf_quant_cfg, user_quant) -> Optional[str]:
|
||||
can_convert = cls.is_awq_marlin_compatible(hf_quant_cfg)
|
||||
is_valid_user_quant = (
|
||||
user_quant is None or user_quant == "marlin" or user_quant == "awq_marlin"
|
||||
)
|
||||
|
||||
if can_convert and is_valid_user_quant:
|
||||
msg = (
|
||||
"The model is convertible to {} during runtime."
|
||||
" Using {} kernel.".format(cls.get_name(), cls.get_name())
|
||||
)
|
||||
logger.info(msg)
|
||||
return cls.get_name()
|
||||
|
||||
if can_convert and user_quant == "awq":
|
||||
logger.info(
|
||||
"Detected that the model can run with awq_marlin"
|
||||
", however you specified quantization=awq explicitly,"
|
||||
" so forcing awq. Use quantization=awq_marlin for"
|
||||
" faster inference"
|
||||
)
|
||||
return None
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> Optional[QuantizeMethodBase]:
|
||||
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
|
||||
from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
|
||||
|
||||
if isinstance(layer, LinearBase) or (
|
||||
isinstance(layer, ParallelLMHead) and self.lm_head_quantized
|
||||
):
|
||||
if is_layer_skipped_awq(prefix, self.modules_to_not_convert):
|
||||
return UnquantizedLinearMethod()
|
||||
# Check if the layer is supported by AWQMarlin.
|
||||
if not check_marlin_supports_layer(layer, self.group_size):
|
||||
logger.warning_once(
|
||||
"Layer '%s' is not supported by AWQMarlin. Falling back to unoptimized AWQ kernels.", # noqa: E501
|
||||
prefix,
|
||||
)
|
||||
return AWQConfig.from_config(self.full_config).get_quant_method(
|
||||
layer, prefix
|
||||
)
|
||||
return AWQMarlinLinearMethod(self)
|
||||
elif isinstance(layer, FusedMoE):
|
||||
from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config
|
||||
|
||||
if not check_moe_marlin_supports_layer(layer, self.group_size):
|
||||
logger.warning_once(
|
||||
f"Layer '{prefix}' is not supported by AWQMoeMarlin. "
|
||||
"Falling back to Moe WNA16 kernels."
|
||||
)
|
||||
return MoeWNA16Config.from_config(self.full_config).get_quant_method(
|
||||
layer, prefix
|
||||
)
|
||||
return AWQMoEMethod(self)
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def is_awq_marlin_compatible(cls, quant_config: dict[str, Any]):
|
||||
# Extract data from quant config.
|
||||
quant_method = quant_config.get("quant_method", "").lower()
|
||||
num_bits = quant_config.get("bits")
|
||||
group_size = quant_config.get("group_size")
|
||||
zero_point = quant_config.get("zero_point")
|
||||
|
||||
if not _is_cuda:
|
||||
return False
|
||||
|
||||
if quant_method != "awq":
|
||||
return False
|
||||
|
||||
# If we cannot find the info needed in the config, cannot convert.
|
||||
if num_bits is None or group_size is None or zero_point is None:
|
||||
return False
|
||||
|
||||
if num_bits not in cls.TYPE_MAP:
|
||||
return False
|
||||
|
||||
return check_marlin_supported(
|
||||
quant_type=cls.TYPE_MAP[num_bits], group_size=group_size, has_zp=zero_point
|
||||
)
|
||||
|
||||
|
||||
class AWQLinearMethod(LinearMethodBase):
|
||||
"""Linear method for AWQ.
|
||||
|
||||
Args:
|
||||
quant_config: The AWQ quantization config.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: AWQConfig):
|
||||
self.quant_config = quant_config
|
||||
|
||||
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,
|
||||
):
|
||||
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 != 0:
|
||||
raise ValueError(
|
||||
"The output size is not aligned with the quantized "
|
||||
"weight shape. This can be caused by too large "
|
||||
"tensor parallel size."
|
||||
)
|
||||
|
||||
weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
qweight = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
input_size_per_partition,
|
||||
output_size_per_partition // self.quant_config.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
packed_dim=1,
|
||||
packed_factor=self.quant_config.pack_factor,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
qzeros = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
input_size_per_partition // self.quant_config.group_size,
|
||||
output_size_per_partition // self.quant_config.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
packed_dim=1,
|
||||
packed_factor=self.quant_config.pack_factor,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
scales = GroupQuantScaleParameter(
|
||||
data=torch.empty(
|
||||
input_size_per_partition // self.quant_config.group_size,
|
||||
output_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
layer.register_parameter("qweight", qweight)
|
||||
layer.register_parameter("qzeros", qzeros)
|
||||
layer.register_parameter("scales", scales)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
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:
|
||||
qweight = layer.qweight
|
||||
scales = layer.scales
|
||||
qzeros = layer.qzeros
|
||||
pack_factor = self.quant_config.pack_factor
|
||||
out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,)
|
||||
reshaped_x = x.reshape(-1, x.shape[-1])
|
||||
out = awq_dequantize(qweight, scales, qzeros)
|
||||
out = torch.matmul(reshaped_x, out)
|
||||
|
||||
if bias is not None:
|
||||
out.add_(bias)
|
||||
return out.reshape(out_shape)
|
||||
|
||||
|
||||
class AWQMarlinLinearMethod(LinearMethodBase):
|
||||
"""Linear method for AWQ Marlin.
|
||||
|
||||
Args:
|
||||
quant_config: The AWQ Marlin quantization config.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: AWQMarlinConfig) -> None:
|
||||
self.quant_config = quant_config
|
||||
|
||||
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,
|
||||
) -> None:
|
||||
del output_size
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
|
||||
# Normalize group_size
|
||||
if self.quant_config.group_size != -1:
|
||||
group_size = self.quant_config.group_size
|
||||
else:
|
||||
group_size = input_size
|
||||
|
||||
verify_marlin_supports_shape(
|
||||
output_size_per_partition=output_size_per_partition,
|
||||
input_size_per_partition=input_size_per_partition,
|
||||
input_size=input_size,
|
||||
group_size=group_size,
|
||||
)
|
||||
|
||||
qweight = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
input_size_per_partition,
|
||||
output_size_per_partition // self.quant_config.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
packed_dim=1,
|
||||
packed_factor=self.quant_config.pack_factor,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
num_groups = input_size_per_partition // group_size
|
||||
|
||||
qzeros = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
num_groups,
|
||||
output_size_per_partition // self.quant_config.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
packed_dim=1,
|
||||
packed_factor=self.quant_config.pack_factor,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
scales = GroupQuantScaleParameter(
|
||||
data=torch.empty(
|
||||
num_groups,
|
||||
output_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
layer.register_parameter("qweight", qweight)
|
||||
layer.register_parameter("qzeros", qzeros)
|
||||
layer.register_parameter("scales", scales)
|
||||
|
||||
layer.input_size_per_partition = input_size_per_partition
|
||||
layer.output_size_per_partition = output_size_per_partition
|
||||
layer.num_groups = num_groups
|
||||
|
||||
# TODO: Update this docs
|
||||
# Checkpoints are serialized in AutoAWQ format, which is different from the
|
||||
# marlin format. This function is called after the weights are loaded.
|
||||
# Here, we handle the repacking
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
device = layer.qweight.device
|
||||
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)
|
||||
|
||||
# Allocate marlin workspace
|
||||
layer.workspace = marlin_make_workspace(device)
|
||||
|
||||
# Repack weights from AWQ format to marlin format.
|
||||
marlin_qweight = awq_marlin_repack(
|
||||
layer.qweight,
|
||||
size_k=layer.input_size_per_partition,
|
||||
size_n=layer.output_size_per_partition,
|
||||
num_bits=self.quant_config.quant_type.size_bits,
|
||||
)
|
||||
replace_parameter(layer, "qweight", marlin_qweight)
|
||||
|
||||
# Permute scales from AWQ format to marlin format.
|
||||
marlin_scales = marlin_permute_scales(
|
||||
layer.scales,
|
||||
size_k=layer.input_size_per_partition,
|
||||
size_n=layer.output_size_per_partition,
|
||||
group_size=self.quant_config.group_size,
|
||||
)
|
||||
replace_parameter(layer, "scales", marlin_scales)
|
||||
|
||||
# Permute zero-points from AWQ format to marlin format.
|
||||
marlin_zp = awq_to_marlin_zero_points(
|
||||
layer.qzeros,
|
||||
size_k=layer.num_groups,
|
||||
size_n=layer.output_size_per_partition,
|
||||
num_bits=self.quant_config.quant_type.size_bits,
|
||||
)
|
||||
replace_parameter(layer, "qzeros", marlin_zp)
|
||||
|
||||
# Not-used
|
||||
layer.g_idx = marlin_make_empty_g_idx(device)
|
||||
layer.g_idx_sort_indices = marlin_make_empty_g_idx(device)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
return apply_awq_marlin_linear(
|
||||
input=x,
|
||||
weight=layer.qweight,
|
||||
weight_scale=layer.scales,
|
||||
weight_zp=layer.qzeros,
|
||||
g_idx=layer.g_idx,
|
||||
g_idx_sort_indices=layer.g_idx_sort_indices,
|
||||
workspace=layer.workspace,
|
||||
quant_type=self.quant_config.quant_type,
|
||||
output_size_per_partition=layer.output_size_per_partition,
|
||||
input_size_per_partition=layer.input_size_per_partition,
|
||||
bias=bias,
|
||||
)
|
||||
|
||||
|
||||
class AWQLinearAscendMethod(AWQLinearMethod):
|
||||
"""Linear method for AWQ on Ascend.
|
||||
|
||||
Args:
|
||||
quant_config: The AWQ quantization config.
|
||||
"""
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
|
||||
qweight_tmp = torch.zeros_like(layer.qweight.data)
|
||||
qzeros_tmp = layer.qzeros.data
|
||||
qzeros_list = []
|
||||
shifts = [0, 4, 1, 5, 2, 6, 3, 7]
|
||||
|
||||
for i in range(0, self.quant_config.pack_factor):
|
||||
shift_num = shifts[i] * 4
|
||||
qzeros_list.append((qzeros_tmp.reshape(-1, 1) >> shift_num) & 0xF)
|
||||
qweight_tmp.bitwise_or_(
|
||||
((layer.qweight.data >> shift_num) * (2 ** (4 * i))) & (0xF << (4 * i))
|
||||
)
|
||||
|
||||
qweight_tmp.bitwise_xor_(0x88888888)
|
||||
|
||||
qzeros_tmp = torch.cat(qzeros_list, dim=-1).reshape(qzeros_tmp.shape[0], -1)
|
||||
qzeros_tmp = -(qzeros_tmp - 8)
|
||||
qzeros_tmp = qzeros_tmp.to(layer.scales.data.dtype)
|
||||
|
||||
layer.zeros = torch.nn.Parameter(qzeros_tmp, requires_grad=False)
|
||||
layer.weight = torch.nn.Parameter(qweight_tmp, requires_grad=False)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
qweight = layer.weight
|
||||
scales = layer.scales
|
||||
qzeros = layer.zeros
|
||||
pack_factor = self.quant_config.pack_factor
|
||||
out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,)
|
||||
reshaped_x = x.reshape(-1, x.shape[-1])
|
||||
|
||||
if bias is not None and bias.dtype == torch.bfloat16:
|
||||
bias = bias.float()
|
||||
|
||||
out = torch_npu.npu_weight_quant_batchmatmul(
|
||||
reshaped_x,
|
||||
qweight,
|
||||
antiquant_scale=scales,
|
||||
antiquant_offset=qzeros,
|
||||
antiquant_group_size=self.quant_config.group_size,
|
||||
bias=bias,
|
||||
)
|
||||
|
||||
return out.reshape(out_shape)
|
||||
|
||||
|
||||
class AWQMoEMethod(FusedMoEMethodBase):
|
||||
|
||||
def __init__(self, quant_config: AWQMarlinConfig):
|
||||
self.quant_config = quant_config
|
||||
if self.quant_config.weight_bits != 4:
|
||||
raise ValueError("AWQMoEMethod only supports 4bit now.")
|
||||
self.quant_type = scalar_types.uint4
|
||||
|
||||
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,
|
||||
):
|
||||
# Delay the import to avoid circular dependency
|
||||
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
|
||||
|
||||
extra_weight_attrs.update(
|
||||
{
|
||||
"is_transposed": True,
|
||||
"quant_method": FusedMoeWeightScaleSupported.GROUP.value,
|
||||
}
|
||||
)
|
||||
|
||||
w13_qweight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
2 * intermediate_size_per_partition // self.quant_config.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_qweight", w13_qweight)
|
||||
set_weight_attrs(w13_qweight, extra_weight_attrs)
|
||||
|
||||
w2_qweight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
intermediate_size_per_partition,
|
||||
hidden_size // self.quant_config.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_qweight", w2_qweight)
|
||||
set_weight_attrs(w2_qweight, extra_weight_attrs)
|
||||
|
||||
num_groups_w13 = hidden_size // self.quant_config.group_size
|
||||
num_groups_w2 = intermediate_size_per_partition // self.quant_config.group_size
|
||||
|
||||
# WEIGHT_SCALES
|
||||
# Allocate 2 scales for w1 and w3 respectively.
|
||||
w13_scales = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
num_groups_w13,
|
||||
intermediate_size_per_partition * 2,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_scales", w13_scales)
|
||||
set_weight_attrs(w13_scales, extra_weight_attrs)
|
||||
|
||||
w2_scales = torch.nn.Parameter(
|
||||
torch.empty(num_experts, num_groups_w2, hidden_size, dtype=params_dtype),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_scales", w2_scales)
|
||||
set_weight_attrs(w2_scales, extra_weight_attrs)
|
||||
|
||||
# WEIGHT_ZERO_POINT
|
||||
# Allocate 2 zero points for w1 and w3 respectively.
|
||||
w13_qzeros = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
num_groups_w13,
|
||||
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)
|
||||
|
||||
w2_qzeros = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
num_groups_w2,
|
||||
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)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
num_experts = layer.w13_qweight.shape[0]
|
||||
device = layer.w13_qweight.device
|
||||
|
||||
layer.w13_g_idx_sort_indices = torch.nn.Parameter(
|
||||
torch.empty((num_experts, 0), dtype=torch.int32, device=device),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.w2_g_idx_sort_indices = torch.nn.Parameter(
|
||||
torch.empty((num_experts, 0), dtype=torch.int32, device=device),
|
||||
requires_grad=False,
|
||||
)
|
||||
|
||||
marlin_w13_qweight = awq_marlin_moe_repack(
|
||||
layer.w13_qweight,
|
||||
layer.w13_g_idx_sort_indices,
|
||||
size_k=layer.w13_qweight.shape[1],
|
||||
size_n=layer.w13_qweight.shape[2] * self.quant_config.pack_factor,
|
||||
num_bits=self.quant_config.weight_bits,
|
||||
)
|
||||
replace_parameter(layer, "w13_qweight", marlin_w13_qweight)
|
||||
|
||||
marlin_w2_qweight = awq_marlin_moe_repack(
|
||||
layer.w2_qweight,
|
||||
layer.w2_g_idx_sort_indices,
|
||||
size_k=layer.w2_qweight.shape[1],
|
||||
size_n=layer.w2_qweight.shape[2] * self.quant_config.pack_factor,
|
||||
num_bits=self.quant_config.weight_bits,
|
||||
)
|
||||
replace_parameter(layer, "w2_qweight", marlin_w2_qweight)
|
||||
|
||||
# hidden_size->intermediate_size
|
||||
marlin_w13_scales = marlin_moe_permute_scales(
|
||||
s=layer.w13_scales,
|
||||
size_k=layer.intermediate_size_per_partition,
|
||||
size_n=layer.w13_scales.shape[2],
|
||||
group_size=self.quant_config.group_size,
|
||||
)
|
||||
|
||||
replace_parameter(layer, "w13_scales", marlin_w13_scales)
|
||||
|
||||
marlin_w2_scales = marlin_moe_permute_scales(
|
||||
s=layer.w2_scales,
|
||||
size_k=layer.intermediate_size_per_partition,
|
||||
size_n=layer.w2_scales.shape[2],
|
||||
group_size=self.quant_config.group_size,
|
||||
)
|
||||
replace_parameter(layer, "w2_scales", marlin_w2_scales)
|
||||
|
||||
marlin_w13_zp = moe_awq_to_marlin_zero_points(
|
||||
layer.w13_qzeros,
|
||||
size_k=layer.w13_qzeros.shape[1],
|
||||
size_n=layer.w13_qzeros.shape[2] * self.quant_config.pack_factor,
|
||||
num_bits=self.quant_config.weight_bits,
|
||||
)
|
||||
replace_parameter(layer, "w13_qzeros", marlin_w13_zp)
|
||||
|
||||
marlin_w2_zp = moe_awq_to_marlin_zero_points(
|
||||
layer.w2_qzeros,
|
||||
size_k=layer.w2_qzeros.shape[1],
|
||||
size_n=layer.w2_qzeros.shape[2] * self.quant_config.pack_factor,
|
||||
num_bits=self.quant_config.weight_bits,
|
||||
)
|
||||
replace_parameter(layer, "w2_qzeros", marlin_w2_zp)
|
||||
|
||||
def create_moe_runner(
|
||||
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
|
||||
):
|
||||
assert get_moe_runner_backend().is_auto()
|
||||
self.moe_runner_config = moe_runner_config
|
||||
self.runner = MoeRunner(MoeRunnerBackend.MARLIN, moe_runner_config)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
dispatch_output: StandardDispatchOutput,
|
||||
) -> CombineInput:
|
||||
|
||||
quant_info = MarlinMoeQuantInfo(
|
||||
w13_qweight=layer.w13_qweight,
|
||||
w2_qweight=layer.w2_qweight,
|
||||
w13_scales=layer.w13_scales,
|
||||
w2_scales=layer.w2_scales,
|
||||
w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices,
|
||||
w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices,
|
||||
w13_qzeros=layer.w13_qzeros,
|
||||
w2_qzeros=layer.w2_qzeros,
|
||||
weight_bits=self.quant_config.weight_bits,
|
||||
)
|
||||
|
||||
return self.runner.run(dispatch_output, quant_info)
|
||||
|
||||
|
||||
class AWQMoEAscendMethod(AWQMoEMethod):
|
||||
def __init__(self, quant_config: AWQConfig):
|
||||
self.quant_config = quant_config
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
w13_qweight_tmp = torch.zeros_like(layer.w13_qweight.data)
|
||||
w2_qweight_tmp = torch.zeros_like(layer.w2_qweight.data)
|
||||
w13_qzeros_list = []
|
||||
w2_qzeros_list = []
|
||||
shifts = [0, 4, 1, 5, 2, 6, 3, 7]
|
||||
for i in range(0, self.quant_config.pack_factor):
|
||||
shift_num = shifts[i] * 4
|
||||
w13_qzeros_list.append(
|
||||
(layer.w13_qzeros.data.reshape(-1, 1) >> shift_num) & 0xF
|
||||
)
|
||||
w2_qzeros_list.append(
|
||||
(layer.w2_qzeros.data.reshape(-1, 1) >> shift_num) & 0xF
|
||||
)
|
||||
w13_qweight_tmp.bitwise_or_(
|
||||
((layer.w13_qweight.data >> shift_num) * (2 ** (4 * i)))
|
||||
& (0xF << (4 * i))
|
||||
)
|
||||
w2_qweight_tmp.bitwise_or_(
|
||||
((layer.w2_qweight.data >> shift_num) * (2 ** (4 * i)))
|
||||
& (0xF << (4 * i))
|
||||
)
|
||||
|
||||
w13_qweight_tmp.bitwise_xor_(0x88888888)
|
||||
w2_qweight_tmp.bitwise_xor_(0x88888888)
|
||||
|
||||
w13_qzeros_tmp = torch.cat(w13_qzeros_list, dim=-1).reshape(
|
||||
layer.w13_qzeros.shape[0], layer.w13_qzeros.shape[1], -1
|
||||
)
|
||||
w13_qzeros_tmp = -(w13_qzeros_tmp - 8)
|
||||
w13_qzeros_tmp = w13_qzeros_tmp.to(layer.w13_scales.data.dtype)
|
||||
w2_qzeros_tmp = torch.cat(w2_qzeros_list, dim=-1).reshape(
|
||||
layer.w2_qzeros.shape[0], layer.w2_qzeros.shape[1], -1
|
||||
)
|
||||
w2_qzeros_tmp = -(w2_qzeros_tmp - 8)
|
||||
w2_qzeros_tmp = w2_qzeros_tmp.to(layer.w2_scales.data.dtype)
|
||||
|
||||
layer.register_parameter(
|
||||
"w13_qzeros", torch.nn.Parameter(w13_qzeros_tmp, requires_grad=False)
|
||||
)
|
||||
layer.register_parameter(
|
||||
"w13_qweight", torch.nn.Parameter(w13_qweight_tmp, requires_grad=False)
|
||||
)
|
||||
layer.register_parameter(
|
||||
"w2_qzeros", torch.nn.Parameter(w2_qzeros_tmp, requires_grad=False)
|
||||
)
|
||||
layer.register_parameter(
|
||||
"w2_qweight", torch.nn.Parameter(w2_qweight_tmp, requires_grad=False)
|
||||
)
|
||||
|
||||
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
|
||||
topk_ids = topk_ids.to(torch.int32)
|
||||
topk_weights = topk_weights.to(x.dtype)
|
||||
output = npu_fused_experts(
|
||||
hidden_states=x,
|
||||
w13=layer.w13_qweight,
|
||||
w13_scale=layer.w13_scales,
|
||||
w13_offset=layer.w13_qzeros,
|
||||
w2=layer.w2_qweight,
|
||||
w2_scale=layer.w2_scales,
|
||||
w2_offset=layer.w2_qzeros,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
top_k=topk_ids.shape[1],
|
||||
use_wna16=True,
|
||||
)
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
|
||||
|
||||
# Register fake implementations for torch.compile support
|
||||
if _is_cuda:
|
||||
|
||||
@register_fake_if_exists("sgl_kernel::awq_marlin_repack")
|
||||
def _(b_q_weight, size_k, size_n, num_bits):
|
||||
return b_q_weight.new_empty(
|
||||
(size_k // 16, size_n * (num_bits // 2)), dtype=b_q_weight.dtype
|
||||
)
|
||||
@@ -0,0 +1,32 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from .awq import (
|
||||
AWQConfig,
|
||||
AWQCPUConfig,
|
||||
AWQLinearMethod,
|
||||
AWQMarlinConfig,
|
||||
AWQMoEMethod,
|
||||
)
|
||||
from .awq_triton import awq_dequantize_decomposition, awq_dequantize_triton
|
||||
from .schemes import (
|
||||
AWQAscendLinearScheme,
|
||||
AWQAscendMoEScheme,
|
||||
AWQLinearScheme,
|
||||
AWQMarlinLinearScheme,
|
||||
AWQMoEScheme,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"AWQConfig",
|
||||
"AWQCPUConfig",
|
||||
"AWQMarlinConfig",
|
||||
"AWQLinearMethod",
|
||||
"AWQMoEMethod",
|
||||
"AWQLinearScheme",
|
||||
"AWQMarlinLinearScheme",
|
||||
"AWQAscendLinearScheme",
|
||||
"AWQMoEScheme",
|
||||
"AWQAscendMoEScheme",
|
||||
"awq_dequantize_triton",
|
||||
"awq_dequantize_decomposition",
|
||||
]
|
||||
@@ -0,0 +1,484 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import warnings
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.linear import LinearBase
|
||||
from sglang.srt.layers.moe import MoeRunnerConfig
|
||||
from sglang.srt.layers.quantization.base_config import (
|
||||
FusedMoEMethodBase,
|
||||
LinearMethodBase,
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from sglang.srt.layers.quantization.marlin_utils import (
|
||||
check_marlin_supported,
|
||||
check_marlin_supports_layer,
|
||||
check_moe_marlin_supports_layer,
|
||||
verify_marlin_supported,
|
||||
)
|
||||
from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
|
||||
from sglang.srt.layers.quantization.utils import get_scalar_types
|
||||
from sglang.srt.utils.patch_torch import register_fake_if_exists
|
||||
|
||||
from .schemes import (
|
||||
AWQAscendLinearScheme,
|
||||
AWQAscendMoEScheme,
|
||||
AWQIntelAMXLinearScheme,
|
||||
AWQIntelAMXMoEScheme,
|
||||
AWQLinearScheme,
|
||||
AWQMarlinLinearScheme,
|
||||
AWQMoEScheme,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.moe.token_dispatcher import (
|
||||
CombineInput,
|
||||
StandardDispatchOutput,
|
||||
)
|
||||
|
||||
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
|
||||
|
||||
_is_cuda = is_cuda()
|
||||
_is_hip = is_hip()
|
||||
_is_xpu = is_xpu()
|
||||
_is_npu = is_npu()
|
||||
|
||||
if not (_is_cuda or _is_hip or _is_xpu or _is_npu):
|
||||
warnings.warn(f"Only CUDA, HIP and XPU support AWQ currently.")
|
||||
|
||||
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 AWQConfig(QuantizationConfig):
|
||||
"""Config class for AWQ.
|
||||
|
||||
Reference: https://arxiv.org/abs/2306.00978
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
weight_bits: int,
|
||||
group_size: int,
|
||||
zero_point: bool,
|
||||
modules_to_not_convert: Optional[List[str]] = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.weight_bits = weight_bits
|
||||
self.group_size = group_size
|
||||
self.zero_point = zero_point
|
||||
self.modules_to_not_convert = modules_to_not_convert or []
|
||||
|
||||
if self.weight_bits != 4:
|
||||
raise ValueError(
|
||||
"Currently, only 4-bit weight quantization is supported for "
|
||||
f"AWQ, but got {self.weight_bits} bits."
|
||||
)
|
||||
self.pack_factor = 32 // self.weight_bits
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"AWQConfig(weight_bits={self.weight_bits}, "
|
||||
f"group_size={self.group_size}, "
|
||||
f"zero_point={self.zero_point}, "
|
||||
f"modules_to_not_convert={self.modules_to_not_convert})"
|
||||
)
|
||||
|
||||
def get_scaled_act_names(self) -> List[str]:
|
||||
return []
|
||||
|
||||
def get_name(self) -> str:
|
||||
return "awq"
|
||||
|
||||
def get_supported_act_dtypes(self) -> List[torch.dtype]:
|
||||
return [torch.float16] if not _is_npu else [torch.float16, torch.bfloat16]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
# The AWQ kernel only supports Turing or newer GPUs.
|
||||
if _is_npu:
|
||||
raise NotImplementedError(
|
||||
'NPU hardware does not support "get_min_capability" feature.'
|
||||
)
|
||||
else:
|
||||
return 75
|
||||
|
||||
@staticmethod
|
||||
def get_config_filenames() -> List[str]:
|
||||
return [
|
||||
"quant_config.json", # E.g., casperhansen/vicuna-7b-v1.5-awq
|
||||
# E.g., abhinavkulkarni/mosaicml-mpt-7b-instruct-w4-g128-awq
|
||||
"quantize_config.json",
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: Dict[str, Any]) -> AWQConfig:
|
||||
weight_bits = cls.get_from_keys(config, ["w_bit", "bits"])
|
||||
group_size = cls.get_from_keys(config, ["q_group_size", "group_size"])
|
||||
zero_point = cls.get_from_keys(config, ["zero_point"])
|
||||
modules_to_not_convert = cls.get_from_keys_or(
|
||||
config, ["modules_to_not_convert"], None
|
||||
)
|
||||
return cls(weight_bits, group_size, zero_point, modules_to_not_convert)
|
||||
|
||||
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 _is_npu:
|
||||
if isinstance(layer, LinearBase):
|
||||
if is_layer_skipped_awq(prefix, self.modules_to_not_convert):
|
||||
return UnquantizedLinearMethod()
|
||||
layer.scheme = self.get_linear_scheme(layer)
|
||||
return AWQLinearMethod(self)
|
||||
elif isinstance(layer, FusedMoE):
|
||||
layer.scheme = self.get_moe_scheme(layer)
|
||||
return AWQMoEMethod(self)
|
||||
return None
|
||||
|
||||
if isinstance(layer, LinearBase):
|
||||
if is_layer_skipped_awq(prefix, self.modules_to_not_convert):
|
||||
return UnquantizedLinearMethod()
|
||||
layer.scheme = self.get_linear_scheme(layer)
|
||||
return AWQLinearMethod(self)
|
||||
return None
|
||||
|
||||
def get_linear_scheme(self, layer: torch.nn.Module):
|
||||
assert isinstance(layer, LinearBase)
|
||||
# TODO: move platform-specific AWQ scheme selection into the platform
|
||||
# plugin factory once quantization hooks are available there.
|
||||
if _is_npu:
|
||||
return AWQAscendLinearScheme(self)
|
||||
return AWQLinearScheme(self)
|
||||
|
||||
def get_moe_scheme(self, layer: torch.nn.Module):
|
||||
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
|
||||
|
||||
assert isinstance(layer, FusedMoE)
|
||||
# This is currently only reached by the NPU path in get_quant_method.
|
||||
if _is_npu:
|
||||
return AWQAscendMoEScheme(self)
|
||||
raise NotImplementedError("AWQConfig only supports MoE scheme on NPU.")
|
||||
|
||||
|
||||
class AWQCPUConfig(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()
|
||||
layer.scheme = self.get_linear_scheme(layer)
|
||||
return AWQLinearMethod(self)
|
||||
elif isinstance(layer, FusedMoE):
|
||||
layer.scheme = self.get_moe_scheme(layer)
|
||||
return AWQMoEMethod(self)
|
||||
return None
|
||||
|
||||
def get_linear_scheme(self, layer: torch.nn.Module):
|
||||
from sglang.srt.layers.linear import LinearBase
|
||||
|
||||
assert isinstance(layer, LinearBase)
|
||||
return AWQIntelAMXLinearScheme(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 AWQIntelAMXMoEScheme(self)
|
||||
|
||||
|
||||
class AWQMarlinConfig(QuantizationConfig):
|
||||
"""Config class for AWQ Marlin"""
|
||||
|
||||
# num_bits -> type
|
||||
TYPE_MAP = {
|
||||
4: scalar_types.uint4,
|
||||
8: scalar_types.uint8,
|
||||
}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
weight_bits: int,
|
||||
group_size: int,
|
||||
zero_point: bool,
|
||||
lm_head_quantized: bool,
|
||||
modules_to_not_convert: Optional[list[str]],
|
||||
full_config: dict[str, Any],
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if _is_hip:
|
||||
warnings.warn(f"HIP does not support fused_marlin_moe currently.")
|
||||
self.pack_factor = 32 // weight_bits # packed into int32
|
||||
self.group_size = group_size
|
||||
self.zero_point = zero_point
|
||||
self.lm_head_quantized = lm_head_quantized
|
||||
self.weight_bits = weight_bits
|
||||
self.modules_to_not_convert = modules_to_not_convert or []
|
||||
self.full_config = full_config
|
||||
|
||||
if self.weight_bits not in self.TYPE_MAP:
|
||||
raise ValueError(
|
||||
f"Unsupported num_bits = {self.weight_bits}. "
|
||||
f"Supported num_bits = {self.TYPE_MAP.keys()}"
|
||||
)
|
||||
|
||||
self.quant_type = self.TYPE_MAP[self.weight_bits]
|
||||
|
||||
verify_marlin_supported(
|
||||
self.quant_type, group_size=self.group_size, has_zp=self.zero_point
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"AWQMarlinConfig(quant_type={self.quant_type}, "
|
||||
f"group_size={self.group_size}, "
|
||||
f"zero_point={self.zero_point}, "
|
||||
f"lm_head_quantized={self.lm_head_quantized}, "
|
||||
f"modules_to_not_convert={self.modules_to_not_convert})"
|
||||
)
|
||||
|
||||
def get_scaled_act_names(self) -> List[str]:
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> str:
|
||||
return "awq_marlin"
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||||
return [torch.half, torch.bfloat16]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 80
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> list[str]:
|
||||
return ["quantize_config.json"]
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> AWQMarlinConfig:
|
||||
weight_bits = cls.get_from_keys(config, ["bits"])
|
||||
group_size = cls.get_from_keys(config, ["group_size"])
|
||||
zero_point = cls.get_from_keys(config, ["zero_point"])
|
||||
lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"], default=False)
|
||||
modules_to_not_convert = cls.get_from_keys_or(
|
||||
config, ["modules_to_not_convert"], None
|
||||
)
|
||||
return cls(
|
||||
weight_bits,
|
||||
group_size,
|
||||
zero_point,
|
||||
lm_head_quantized,
|
||||
modules_to_not_convert,
|
||||
config,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def override_quantization_method(cls, hf_quant_cfg, user_quant) -> Optional[str]:
|
||||
can_convert = cls.is_awq_marlin_compatible(hf_quant_cfg)
|
||||
is_valid_user_quant = (
|
||||
user_quant is None or user_quant == "marlin" or user_quant == "awq_marlin"
|
||||
)
|
||||
|
||||
if can_convert and is_valid_user_quant:
|
||||
msg = (
|
||||
"The model is convertible to {} during runtime."
|
||||
" Using {} kernel.".format(cls.get_name(), cls.get_name())
|
||||
)
|
||||
logger.info(msg)
|
||||
return cls.get_name()
|
||||
|
||||
if can_convert and user_quant == "awq":
|
||||
logger.info(
|
||||
"Detected that the model can run with awq_marlin"
|
||||
", however you specified quantization=awq explicitly,"
|
||||
" so forcing awq. Use quantization=awq_marlin for"
|
||||
" faster inference"
|
||||
)
|
||||
return None
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> Optional[QuantizeMethodBase]:
|
||||
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
|
||||
from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
|
||||
|
||||
if isinstance(layer, LinearBase) or (
|
||||
isinstance(layer, ParallelLMHead) and self.lm_head_quantized
|
||||
):
|
||||
if is_layer_skipped_awq(prefix, self.modules_to_not_convert):
|
||||
return UnquantizedLinearMethod()
|
||||
# Check if the layer is supported by AWQMarlin.
|
||||
if not check_marlin_supports_layer(layer, self.group_size):
|
||||
logger.warning_once(
|
||||
"Layer '%s' is not supported by AWQMarlin. Falling back to unoptimized AWQ kernels.", # noqa: E501
|
||||
prefix,
|
||||
)
|
||||
return AWQConfig.from_config(self.full_config).get_quant_method(
|
||||
layer, prefix
|
||||
)
|
||||
layer.scheme = self.get_linear_scheme(layer)
|
||||
return AWQLinearMethod(self)
|
||||
elif isinstance(layer, FusedMoE):
|
||||
from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config
|
||||
|
||||
if not check_moe_marlin_supports_layer(layer, self.group_size):
|
||||
logger.warning_once(
|
||||
f"Layer '{prefix}' is not supported by AWQMoeMarlin. "
|
||||
"Falling back to Moe WNA16 kernels."
|
||||
)
|
||||
return MoeWNA16Config.from_config(self.full_config).get_quant_method(
|
||||
layer, prefix
|
||||
)
|
||||
layer.scheme = self.get_moe_scheme(layer)
|
||||
return AWQMoEMethod(self)
|
||||
return None
|
||||
|
||||
def get_linear_scheme(self, layer: torch.nn.Module):
|
||||
return AWQMarlinLinearScheme(self)
|
||||
|
||||
def get_moe_scheme(self, layer: torch.nn.Module):
|
||||
return AWQMoEScheme(self)
|
||||
|
||||
@classmethod
|
||||
def is_awq_marlin_compatible(cls, quant_config: dict[str, Any]):
|
||||
# Extract data from quant config.
|
||||
quant_method = quant_config.get("quant_method", "").lower()
|
||||
num_bits = quant_config.get("bits")
|
||||
group_size = quant_config.get("group_size")
|
||||
zero_point = quant_config.get("zero_point")
|
||||
|
||||
if not _is_cuda:
|
||||
return False
|
||||
|
||||
if quant_method != "awq":
|
||||
return False
|
||||
|
||||
# If we cannot find the info needed in the config, cannot convert.
|
||||
if num_bits is None or group_size is None or zero_point is None:
|
||||
return False
|
||||
|
||||
if num_bits not in cls.TYPE_MAP:
|
||||
return False
|
||||
|
||||
return check_marlin_supported(
|
||||
quant_type=cls.TYPE_MAP[num_bits], group_size=group_size, has_zp=zero_point
|
||||
)
|
||||
|
||||
|
||||
class AWQLinearMethod(LinearMethodBase):
|
||||
"""Linear method for AWQ.
|
||||
|
||||
Args:
|
||||
quant_config: The AWQ quantization config.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: AWQConfig):
|
||||
self.quant_config = quant_config
|
||||
|
||||
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,
|
||||
):
|
||||
weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
layer.scheme.create_weights(
|
||||
layer=layer,
|
||||
input_size_per_partition=input_size_per_partition,
|
||||
output_partition_sizes=output_partition_sizes,
|
||||
input_size=input_size,
|
||||
output_size=output_size,
|
||||
params_dtype=params_dtype,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
layer.scheme.process_weights_after_loading(layer)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
return layer.scheme.apply_weights(layer, x, bias)
|
||||
|
||||
|
||||
class AWQMoEMethod(FusedMoEMethodBase):
|
||||
|
||||
def __init__(self, quant_config: AWQMarlinConfig):
|
||||
self.quant_config = quant_config
|
||||
self.quant_type = scalar_types.uint4
|
||||
if self.quant_config.weight_bits != 4:
|
||||
raise ValueError("AWQMoEMethod only supports 4bit now.")
|
||||
|
||||
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,
|
||||
):
|
||||
layer.scheme.create_weights(
|
||||
layer=layer,
|
||||
num_experts=num_experts,
|
||||
hidden_size=hidden_size,
|
||||
intermediate_size_per_partition=intermediate_size_per_partition,
|
||||
params_dtype=params_dtype,
|
||||
**extra_weight_attrs,
|
||||
)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
layer.scheme.process_weights_after_loading(layer)
|
||||
|
||||
def create_moe_runner(
|
||||
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
|
||||
):
|
||||
layer.scheme.create_moe_runner(layer, moe_runner_config)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
dispatch_output: StandardDispatchOutput,
|
||||
) -> CombineInput:
|
||||
return layer.scheme.apply_weights(layer, dispatch_output)
|
||||
|
||||
|
||||
# Register fake implementations for torch.compile support
|
||||
if _is_cuda:
|
||||
|
||||
@register_fake_if_exists("sgl_kernel::awq_marlin_repack")
|
||||
def _(b_q_weight, size_k, size_n, num_bits):
|
||||
return b_q_weight.new_empty(
|
||||
(size_k // 16, size_n * (num_bits // 2)), dtype=b_q_weight.dtype
|
||||
)
|
||||
@@ -0,0 +1,19 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from .awq_cpu import AWQIntelAMXLinearScheme, AWQIntelAMXMoEScheme
|
||||
from .awq_linear import AWQAscendLinearScheme, AWQLinearScheme
|
||||
from .awq_marlin import AWQMarlinLinearScheme
|
||||
from .awq_moe import AWQAscendMoEScheme, AWQMoEScheme
|
||||
from .awq_scheme import AWQLinearSchemeBase, AWQMoESchemeBase
|
||||
|
||||
__all__ = [
|
||||
"AWQLinearSchemeBase",
|
||||
"AWQMoESchemeBase",
|
||||
"AWQLinearScheme",
|
||||
"AWQAscendLinearScheme",
|
||||
"AWQIntelAMXLinearScheme",
|
||||
"AWQMarlinLinearScheme",
|
||||
"AWQMoEScheme",
|
||||
"AWQAscendMoEScheme",
|
||||
"AWQIntelAMXMoEScheme",
|
||||
]
|
||||
Executable → Regular
+37
-53
@@ -1,65 +1,29 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, List, Optional
|
||||
from typing import TYPE_CHECKING, 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,
|
||||
)
|
||||
from sglang.srt.layers.moe import MoeRunnerConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
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"]
|
||||
|
||||
|
||||
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."""
|
||||
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(
|
||||
@@ -75,7 +39,6 @@ class AWQLinearIntelAMXMethod(AWQLinearMethod):
|
||||
x: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
|
||||
return torch.ops.sgl_kernel.int4_scaled_mm_cpu(
|
||||
x,
|
||||
layer.qweight,
|
||||
@@ -85,8 +48,16 @@ class AWQLinearIntelAMXMethod(AWQLinearMethod):
|
||||
)
|
||||
|
||||
|
||||
class AWQMoEIntelAMXMethod(AWQMoEMethod):
|
||||
"""MoE method for AWQ on Intel CPU with AMX."""
|
||||
class AWQIntelAMXLinearScheme(AWQLinearScheme):
|
||||
"""Linear scheme for AWQ on Intel CPU with AMX."""
|
||||
|
||||
def _init_kernel(self, quant_config: "AWQConfig"):
|
||||
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(
|
||||
@@ -104,7 +75,7 @@ class AWQMoEIntelAMXMethod(AWQMoEMethod):
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
dispatch_output: StandardDispatchOutput,
|
||||
dispatch_output: "StandardDispatchOutput",
|
||||
) -> torch.Tensor:
|
||||
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
|
||||
|
||||
@@ -131,3 +102,16 @@ class AWQMoEIntelAMXMethod(AWQMoEMethod):
|
||||
True, # is_vnni
|
||||
)
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
|
||||
|
||||
class AWQIntelAMXMoEScheme(AWQMoEScheme):
|
||||
"""MoE scheme for AWQ on Intel CPU with AMX."""
|
||||
|
||||
def _init_kernel(self, quant_config: "AWQConfig"):
|
||||
return AWQIntelAMXMoEKernel(quant_config)
|
||||
|
||||
def create_moe_runner(
|
||||
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)
|
||||
@@ -0,0 +1,110 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, List, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.parameter import GroupQuantScaleParameter, PackedvLLMParameter
|
||||
|
||||
from .awq_scheme import AWQLinearSchemeBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.quantization.awq.awq import AWQConfig
|
||||
|
||||
__all__ = ["AWQLinearScheme", "AWQAscendLinearScheme"]
|
||||
|
||||
|
||||
class AWQLinearScheme(AWQLinearSchemeBase):
|
||||
def __init__(self, quant_config: "AWQConfig"):
|
||||
self.quant_config = quant_config
|
||||
self.kernel = self._init_kernel(quant_config)
|
||||
|
||||
def _init_kernel(self, quant_config: "AWQConfig"):
|
||||
from sglang.srt.hardware_backend.gpu.quantization.awq_kernels import (
|
||||
AWQLinearKernel,
|
||||
)
|
||||
|
||||
return AWQLinearKernel(quant_config)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
input_size_per_partition: int,
|
||||
output_partition_sizes: List[int],
|
||||
params_dtype: torch.dtype,
|
||||
weight_loader,
|
||||
**kwargs,
|
||||
):
|
||||
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 != 0:
|
||||
raise ValueError(
|
||||
"The output size is not aligned with the quantized "
|
||||
"weight shape. This can be caused by too large "
|
||||
"tensor parallel size."
|
||||
)
|
||||
|
||||
qweight = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
input_size_per_partition,
|
||||
output_size_per_partition // self.quant_config.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
packed_dim=1,
|
||||
packed_factor=self.quant_config.pack_factor,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
qzeros = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
input_size_per_partition // self.quant_config.group_size,
|
||||
output_size_per_partition // self.quant_config.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
packed_dim=1,
|
||||
packed_factor=self.quant_config.pack_factor,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
scales = GroupQuantScaleParameter(
|
||||
data=torch.empty(
|
||||
input_size_per_partition // self.quant_config.group_size,
|
||||
output_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
layer.register_parameter("qweight", qweight)
|
||||
layer.register_parameter("qzeros", qzeros)
|
||||
layer.register_parameter("scales", scales)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
self.kernel.process_weights_after_loading(layer)
|
||||
|
||||
def apply_weights(
|
||||
self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor]
|
||||
):
|
||||
return self.kernel.apply(layer, x, bias)
|
||||
|
||||
|
||||
class AWQAscendLinearScheme(AWQLinearScheme):
|
||||
def _init_kernel(self, quant_config: "AWQConfig"):
|
||||
from sglang.srt.hardware_backend.npu.quantization.awq_kernels import (
|
||||
AWQAscendLinearKernel,
|
||||
)
|
||||
|
||||
return AWQAscendLinearKernel(quant_config)
|
||||
@@ -0,0 +1,105 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
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
|
||||
|
||||
from .awq_scheme import AWQLinearSchemeBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.quantization.awq.awq import AWQMarlinConfig
|
||||
|
||||
__all__ = ["AWQMarlinLinearScheme"]
|
||||
|
||||
|
||||
class AWQMarlinLinearScheme(AWQLinearSchemeBase):
|
||||
def __init__(self, quant_config: "AWQMarlinConfig"):
|
||||
self.quant_config = quant_config
|
||||
self.kernel = AWQMarlinLinearKernel(quant_config)
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
input_size_per_partition: int,
|
||||
output_partition_sizes: list[int],
|
||||
input_size: int,
|
||||
params_dtype: torch.dtype,
|
||||
weight_loader,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
|
||||
group_size = (
|
||||
self.quant_config.group_size
|
||||
if self.quant_config.group_size != -1
|
||||
else input_size
|
||||
)
|
||||
|
||||
verify_marlin_supports_shape(
|
||||
output_size_per_partition=output_size_per_partition,
|
||||
input_size_per_partition=input_size_per_partition,
|
||||
input_size=input_size,
|
||||
group_size=group_size,
|
||||
)
|
||||
|
||||
qweight = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
input_size_per_partition,
|
||||
output_size_per_partition // self.quant_config.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
packed_dim=1,
|
||||
packed_factor=self.quant_config.pack_factor,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
num_groups = input_size_per_partition // group_size
|
||||
|
||||
qzeros = PackedvLLMParameter(
|
||||
data=torch.empty(
|
||||
num_groups,
|
||||
output_size_per_partition // self.quant_config.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
packed_dim=1,
|
||||
packed_factor=self.quant_config.pack_factor,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
scales = GroupQuantScaleParameter(
|
||||
data=torch.empty(
|
||||
num_groups,
|
||||
output_size_per_partition,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
|
||||
layer.register_parameter("qweight", qweight)
|
||||
layer.register_parameter("qzeros", qzeros)
|
||||
layer.register_parameter("scales", scales)
|
||||
|
||||
layer.input_size_per_partition = input_size_per_partition
|
||||
layer.output_size_per_partition = output_size_per_partition
|
||||
layer.num_groups = num_groups
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
self.kernel.process_weights_after_loading(layer)
|
||||
|
||||
def apply_weights(
|
||||
self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor]
|
||||
):
|
||||
return self.kernel.apply(layer, x, bias)
|
||||
@@ -0,0 +1,156 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.linear import set_weight_attrs
|
||||
from sglang.srt.layers.moe import (
|
||||
MoeRunner,
|
||||
MoeRunnerBackend,
|
||||
MoeRunnerConfig,
|
||||
get_moe_runner_backend,
|
||||
)
|
||||
|
||||
from .awq_scheme import AWQMoESchemeBase
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
|
||||
from sglang.srt.layers.quantization.awq.awq import AWQConfig, AWQMarlinConfig
|
||||
|
||||
__all__ = ["AWQMoEScheme", "AWQAscendMoEScheme"]
|
||||
|
||||
|
||||
class AWQMoEScheme(AWQMoESchemeBase):
|
||||
def __init__(self, quant_config: "AWQMarlinConfig"):
|
||||
self.quant_config = quant_config
|
||||
if self.quant_config.weight_bits != 4:
|
||||
raise ValueError("AWQMoEScheme only supports 4bit now.")
|
||||
self.kernel = self._init_kernel(quant_config)
|
||||
|
||||
def _init_kernel(self, quant_config: "AWQMarlinConfig"):
|
||||
from sglang.srt.hardware_backend.gpu.quantization.awq_kernels import (
|
||||
AWQMoEKernel,
|
||||
)
|
||||
|
||||
return AWQMoEKernel(quant_config)
|
||||
|
||||
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,
|
||||
):
|
||||
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
|
||||
|
||||
extra_weight_attrs.update(
|
||||
{
|
||||
"is_transposed": True,
|
||||
"quant_method": FusedMoeWeightScaleSupported.GROUP.value,
|
||||
}
|
||||
)
|
||||
|
||||
w13_qweight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
hidden_size,
|
||||
2 * intermediate_size_per_partition // self.quant_config.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_qweight", w13_qweight)
|
||||
set_weight_attrs(w13_qweight, extra_weight_attrs)
|
||||
|
||||
w2_qweight = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
intermediate_size_per_partition,
|
||||
hidden_size // self.quant_config.pack_factor,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_qweight", w2_qweight)
|
||||
set_weight_attrs(w2_qweight, extra_weight_attrs)
|
||||
|
||||
num_groups_w13 = hidden_size // self.quant_config.group_size
|
||||
num_groups_w2 = intermediate_size_per_partition // self.quant_config.group_size
|
||||
|
||||
w13_scales = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
num_groups_w13,
|
||||
intermediate_size_per_partition * 2,
|
||||
dtype=params_dtype,
|
||||
),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w13_scales", w13_scales)
|
||||
set_weight_attrs(w13_scales, extra_weight_attrs)
|
||||
|
||||
w2_scales = torch.nn.Parameter(
|
||||
torch.empty(num_experts, num_groups_w2, hidden_size, dtype=params_dtype),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.register_parameter("w2_scales", w2_scales)
|
||||
set_weight_attrs(w2_scales, extra_weight_attrs)
|
||||
|
||||
w13_qzeros = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
num_groups_w13,
|
||||
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)
|
||||
|
||||
w2_qzeros = torch.nn.Parameter(
|
||||
torch.empty(
|
||||
num_experts,
|
||||
num_groups_w2,
|
||||
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)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
self.kernel.process_weights_after_loading(layer)
|
||||
|
||||
def create_moe_runner(
|
||||
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
|
||||
):
|
||||
assert get_moe_runner_backend().is_auto()
|
||||
self.moe_runner_config = moe_runner_config
|
||||
self.kernel.runner = MoeRunner(MoeRunnerBackend.MARLIN, moe_runner_config)
|
||||
|
||||
def apply_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
dispatch_output: "StandardDispatchOutput",
|
||||
):
|
||||
return self.kernel.apply(layer, dispatch_output)
|
||||
|
||||
|
||||
class AWQAscendMoEScheme(AWQMoEScheme):
|
||||
def _init_kernel(self, quant_config: "AWQConfig"):
|
||||
from sglang.srt.hardware_backend.npu.quantization.awq_kernels import (
|
||||
AWQAscendMoEKernel,
|
||||
)
|
||||
|
||||
return AWQAscendMoEKernel(quant_config)
|
||||
|
||||
def create_moe_runner(
|
||||
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
|
||||
):
|
||||
self.moe_runner_config = moe_runner_config
|
||||
@@ -0,0 +1,54 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from abc import abstractmethod
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.moe import MoeRunnerConfig
|
||||
from sglang.srt.layers.quantization.base_scheme import BaseLinearScheme, BaseMoEScheme
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
|
||||
|
||||
__all__ = ["AWQLinearSchemeBase", "AWQMoESchemeBase"]
|
||||
|
||||
|
||||
class AWQLinearSchemeBase(BaseLinearScheme):
|
||||
@abstractmethod
|
||||
def create_weights(self, *args, **kwargs):
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module):
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def apply_weights(
|
||||
self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor]
|
||||
):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class AWQMoESchemeBase(BaseMoEScheme):
|
||||
@abstractmethod
|
||||
def create_weights(self, *args, **kwargs):
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def create_moe_runner(
|
||||
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
|
||||
):
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module):
|
||||
raise NotImplementedError
|
||||
|
||||
@abstractmethod
|
||||
def apply_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
dispatch_output: "StandardDispatchOutput",
|
||||
):
|
||||
raise NotImplementedError
|
||||
@@ -100,7 +100,7 @@ if _is_cuda:
|
||||
elif _is_cpu and _is_cpu_amx_available:
|
||||
pass
|
||||
elif _is_hip:
|
||||
from sglang.srt.layers.quantization.awq_triton import (
|
||||
from sglang.srt.layers.quantization.awq.awq_triton import (
|
||||
awq_dequantize_triton as awq_dequantize,
|
||||
)
|
||||
else:
|
||||
|
||||
@@ -79,14 +79,14 @@ def awq_dequantize_func():
|
||||
return awq_dequantize
|
||||
elif _is_hip:
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.srt.layers.quantization.awq_triton import (
|
||||
from sglang.srt.layers.quantization.awq.awq_triton import (
|
||||
awq_dequantize_triton as awq_dequantize,
|
||||
)
|
||||
|
||||
return debug_kernel_api(awq_dequantize, op_name="DeepseekCommon.awq_dequantize")
|
||||
elif _is_npu:
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.srt.layers.quantization.awq_triton import (
|
||||
from sglang.srt.layers.quantization.awq.awq_triton import (
|
||||
awq_dequantize_decomposition as awq_dequantize,
|
||||
)
|
||||
|
||||
|
||||
@@ -117,7 +117,7 @@ if _is_cuda:
|
||||
elif _is_cpu and _is_cpu_amx_available:
|
||||
pass
|
||||
elif _is_hip:
|
||||
from sglang.srt.layers.quantization.awq_triton import (
|
||||
from sglang.srt.layers.quantization.awq.awq_triton import (
|
||||
awq_dequantize_triton as awq_dequantize,
|
||||
)
|
||||
else:
|
||||
|
||||
@@ -97,7 +97,7 @@ if _is_cuda:
|
||||
elif _is_cpu and _is_cpu_amx_available:
|
||||
pass
|
||||
elif _is_hip:
|
||||
from sglang.srt.layers.quantization.awq_triton import (
|
||||
from sglang.srt.layers.quantization.awq.awq_triton import (
|
||||
awq_dequantize_triton as awq_dequantize,
|
||||
)
|
||||
else:
|
||||
|
||||
@@ -12,7 +12,7 @@ import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.quantization.awq_triton import (
|
||||
from sglang.srt.layers.quantization.awq.awq_triton import (
|
||||
AWQ_TRITON_SUPPORTED_GROUP_SIZES,
|
||||
awq_dequantize_triton,
|
||||
awq_gemm_triton,
|
||||
|
||||
Reference in New Issue
Block a user