# 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"