chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,222 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import torch
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm import envs
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
pack_quantized_values_into_int32,
|
||||
unpack_quantized_values_into_int32,
|
||||
)
|
||||
from vllm.platforms import CpuArchEnum, current_platform
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
|
||||
|
||||
_CPUWNA16_SUPPORTED_QUANT_TYPES = (scalar_types.uint4, scalar_types.uint4b8)
|
||||
|
||||
|
||||
class CPUWNA16LinearKernel(MPLinearKernel):
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return -1
|
||||
|
||||
@classmethod
|
||||
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
|
||||
if not current_platform.is_cpu():
|
||||
return False, "CPUWNA16 only supported on CPU"
|
||||
|
||||
if c.weight_type not in _CPUWNA16_SUPPORTED_QUANT_TYPES:
|
||||
return (
|
||||
False,
|
||||
f"Quant type ({c.weight_type}) not supported by "
|
||||
"CPUWNA16, supported types are: "
|
||||
f"{_CPUWNA16_SUPPORTED_QUANT_TYPES}",
|
||||
)
|
||||
|
||||
if c.group_size != -1 and c.group_size % 2 != 0:
|
||||
return (
|
||||
False,
|
||||
f"Group size ({c.group_size}) not supported by "
|
||||
"CPUWNA16, supported group sizes are multiples of 2",
|
||||
)
|
||||
|
||||
if c.partition_weight_shape[0] % 32 != 0:
|
||||
return (
|
||||
False,
|
||||
f"Input size ({c.partition_weight_shape[0]}) not supported by "
|
||||
"CPUWNA16, supported sizes are multiples of 32",
|
||||
)
|
||||
|
||||
if c.partition_weight_shape[1] % 32 != 0:
|
||||
return (
|
||||
False,
|
||||
f"Output size ({c.partition_weight_shape[1]}) not supported by "
|
||||
"CPUWNA16, supported sizes are multiples of 32",
|
||||
)
|
||||
|
||||
return True, None
|
||||
|
||||
# note assumes that
|
||||
# `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
|
||||
# `weight_scale` is: {input_dim = 0, output_dim = 1}
|
||||
# `weight_zp` is: {input_dim = 0, output_dim = 1, packed_dim = 1}
|
||||
def _process_gptq_weights_w4a16(self, layer: torch.nn.Module):
|
||||
packed_weight = getattr(layer, self.w_q_name)
|
||||
bits = self.config.weight_type.mantissa
|
||||
pack_factor = 32 // bits
|
||||
p_w_k, _ = packed_weight.size()
|
||||
input_size = p_w_k * pack_factor
|
||||
isa_hint = _get_isa_hint(getattr(layer, self.w_s_name).dtype)
|
||||
layer.isa_hint = isa_hint
|
||||
|
||||
# convert input dim packed to output dim packed
|
||||
weight = unpack_quantized_values_into_int32(
|
||||
packed_weight, self.config.weight_type, 0
|
||||
)
|
||||
weight = pack_quantized_values_into_int32(weight, self.config.weight_type, 1)
|
||||
# make 16 output channel as a block and transpose to the make
|
||||
# the block contiguous
|
||||
weight = (
|
||||
weight.view(input_size, -1, 16 // pack_factor)
|
||||
.permute(1, 0, 2)
|
||||
.reshape(-1, input_size * 16 // pack_factor)
|
||||
.contiguous()
|
||||
)
|
||||
getattr(layer, self.w_q_name).data = weight
|
||||
|
||||
# note assumes that
|
||||
# `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
|
||||
# `weight_scale` is: {input_dim = 0, output_dim = 1}
|
||||
# `weight_zp` is: {input_dim = 0, output_dim = 1, packed_dim = 1}
|
||||
def _process_gptq_weights_w4a8(self, layer: torch.nn.Module):
|
||||
packed_weight = getattr(layer, self.w_q_name)
|
||||
scales = getattr(layer, self.w_s_name)
|
||||
group_num = scales.data.size(0)
|
||||
zp_output_size = scales.data.size(1) // 8
|
||||
if self.config.zero_points:
|
||||
assert self.w_zp_name
|
||||
packed_zp = getattr(layer, self.w_zp_name)
|
||||
else:
|
||||
# w4a8 kernel always requires zp, allocate a fake zp
|
||||
assert self.w_zp_name
|
||||
packed_zp = torch.nn.Parameter(
|
||||
torch.ones(group_num, zp_output_size, dtype=torch.int32) * -2004318072,
|
||||
requires_grad=False,
|
||||
)
|
||||
setattr(layer, self.w_zp_name, packed_zp)
|
||||
|
||||
# FIXME: some bugs in convert_weight_packed_scale_zp with GPTQ format,
|
||||
# repack to AWQ weight
|
||||
weight = unpack_quantized_values_into_int32(
|
||||
packed_weight, self.config.weight_type, 0
|
||||
)
|
||||
input_size, output_size = weight.size()
|
||||
weight = weight.view(input_size, output_size // 8, 8)
|
||||
weight = weight[:, :, (0, 2, 4, 6, 1, 3, 5, 7)].reshape(input_size, output_size)
|
||||
weight = pack_quantized_values_into_int32(
|
||||
weight, self.config.weight_type, 1
|
||||
).contiguous()
|
||||
|
||||
zp = unpack_quantized_values_into_int32(packed_zp, self.config.weight_type, 1)
|
||||
zp = zp.view(group_num, output_size // 8, 8)
|
||||
zp = zp[:, :, (0, 2, 4, 6, 1, 3, 5, 7)].reshape(group_num, output_size)
|
||||
zp = pack_quantized_values_into_int32(
|
||||
zp, self.config.weight_type, 1
|
||||
).contiguous()
|
||||
|
||||
blocked_w, blocked_zp, blocked_s = ops.convert_weight_packed_scale_zp(
|
||||
weight,
|
||||
zp,
|
||||
scales.data,
|
||||
ops.CPUQuantAlgo.AWQ,
|
||||
)
|
||||
|
||||
if layer.bias is not None:
|
||||
layer.bias.data = layer.bias.float()
|
||||
|
||||
packed_weight.data = blocked_w
|
||||
scales.data = blocked_s
|
||||
packed_zp.data = blocked_zp
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module):
|
||||
if (not self.config.zero_points) and (self.w_zp_name is not None):
|
||||
setattr(layer, self.w_zp_name, None)
|
||||
|
||||
if (not self.config.has_g_idx) and (self.w_gidx_name is not None):
|
||||
setattr(layer, self.w_gidx_name, None)
|
||||
|
||||
weights = getattr(layer, self.w_q_name)
|
||||
# Require GPTQ pack format
|
||||
assert weights.input_dim == weights.packed_dim
|
||||
|
||||
# Weights in CT format is [output_size, input_size]
|
||||
if weights.input_dim == 1:
|
||||
weights.data = weights.t()
|
||||
|
||||
# Scales in CT format is [output_size, group_num]
|
||||
scales = getattr(layer, self.w_s_name)
|
||||
if scales.output_dim == 0:
|
||||
scales.data = scales.t().contiguous()
|
||||
|
||||
# Zero points in CT format is [output_size, group_num]
|
||||
# Zero points in awq_marlin format is [output_size, group_num]
|
||||
if self.config.zero_points:
|
||||
assert self.w_zp_name
|
||||
zp = getattr(layer, self.w_zp_name)
|
||||
if zp.output_dim == 0:
|
||||
zp.data = zp.t().contiguous()
|
||||
|
||||
supports_amx = torch.cpu._is_amx_tile_supported()
|
||||
supports_riscv = current_platform.get_cpu_architecture() == CpuArchEnum.RISCV
|
||||
layer.use_w4a8 = (
|
||||
envs.VLLM_CPU_INT4_W4A8
|
||||
and not self.config.has_g_idx
|
||||
and self.config.act_type == torch.bfloat16
|
||||
and (supports_amx or supports_riscv)
|
||||
)
|
||||
# layer.use_w4a8 = False
|
||||
# AWQ format will be converted to GPTQ format in `AutoAWQMarlinLinearMethod`
|
||||
if layer.use_w4a8:
|
||||
self._process_gptq_weights_w4a8(layer)
|
||||
else:
|
||||
self._process_gptq_weights_w4a16(layer)
|
||||
|
||||
def apply_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
w_q, w_s, w_zp, w_gidx = self._get_weight_params(layer)
|
||||
if layer.use_w4a8:
|
||||
x = ops.int4_scaled_mm_cpu(
|
||||
x=x,
|
||||
w=w_q,
|
||||
w_zeros=w_zp,
|
||||
w_scales=w_s,
|
||||
bias=bias,
|
||||
)
|
||||
else:
|
||||
x = ops.cpu_gemm_wna16(
|
||||
input=x,
|
||||
q_weight=w_q,
|
||||
scales=w_s,
|
||||
zeros=w_zp,
|
||||
g_idx=w_gidx,
|
||||
bias=bias,
|
||||
pack_factor=8, # 32 // 4
|
||||
isa_hint=layer.isa_hint,
|
||||
)
|
||||
return x
|
||||
|
||||
|
||||
def _get_isa_hint(dtype: torch.dtype) -> str:
|
||||
supports_amx = torch.cpu._is_amx_tile_supported()
|
||||
if supports_amx and dtype in (torch.bfloat16,):
|
||||
return "amx"
|
||||
elif current_platform.get_cpu_architecture() == CpuArchEnum.RISCV:
|
||||
return "rvv"
|
||||
else:
|
||||
return "vec"
|
||||
Reference in New Issue
Block a user