| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | from vllm import _custom_ops as ops |
| 7 | from vllm import envs |
| 8 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 9 | pack_quantized_values_into_int32, |
| 10 | unpack_quantized_values_into_int32, |
| 11 | ) |
| 12 | from vllm.platforms import CpuArchEnum, current_platform |
| 13 | from vllm.scalar_type import scalar_types |
| 14 | |
| 15 | from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig |
| 16 | |
| 17 | _CPUWNA16_SUPPORTED_QUANT_TYPES = (scalar_types.uint4, scalar_types.uint4b8) |
| 18 | |
| 19 | |
| 20 | class CPUWNA16LinearKernel(MPLinearKernel): |
| 21 | @classmethod |
| 22 | def get_min_capability(cls) -> int: |
| 23 | return -1 |
| 24 | |
| 25 | @classmethod |
| 26 | def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]: |
| 27 | if not current_platform.is_cpu(): |
| 28 | return False, "CPUWNA16 only supported on CPU" |
| 29 | |
| 30 | if c.weight_type not in _CPUWNA16_SUPPORTED_QUANT_TYPES: |
| 31 | return ( |
| 32 | False, |
| 33 | f"Quant type ({c.weight_type}) not supported by " |
| 34 | "CPUWNA16, supported types are: " |
| 35 | f"{_CPUWNA16_SUPPORTED_QUANT_TYPES}", |
| 36 | ) |
| 37 | |
| 38 | if c.group_size != -1 and c.group_size % 2 != 0: |
| 39 | return ( |
| 40 | False, |
| 41 | f"Group size ({c.group_size}) not supported by " |
| 42 | "CPUWNA16, supported group sizes are multiples of 2", |
| 43 | ) |
| 44 | |
| 45 | if c.partition_weight_shape[0] % 32 != 0: |
| 46 | return ( |
| 47 | False, |
| 48 | f"Input size ({c.partition_weight_shape[0]}) not supported by " |
| 49 | "CPUWNA16, supported sizes are multiples of 32", |
| 50 | ) |
| 51 | |
| 52 | if c.partition_weight_shape[1] % 32 != 0: |
| 53 | return ( |
| 54 | False, |
| 55 | f"Output size ({c.partition_weight_shape[1]}) not supported by " |
| 56 | "CPUWNA16, supported sizes are multiples of 32", |
| 57 | ) |
| 58 | |
| 59 | return True, None |
| 60 | |
| 61 | # note assumes that |
| 62 | # `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0} |
| 63 | # `weight_scale` is: {input_dim = 0, output_dim = 1} |
| 64 | # `weight_zp` is: {input_dim = 0, output_dim = 1, packed_dim = 1} |
| 65 | def _process_gptq_weights_w4a16(self, layer: torch.nn.Module): |
| 66 | packed_weight = getattr(layer, self.w_q_name) |
| 67 | bits = self.config.weight_type.mantissa |
| 68 | pack_factor = 32 // bits |
| 69 | p_w_k, _ = packed_weight.size() |
| 70 | input_size = p_w_k * pack_factor |
| 71 | isa_hint = _get_isa_hint(getattr(layer, self.w_s_name).dtype) |
| 72 | layer.isa_hint = isa_hint |
| 73 | |
| 74 | # convert input dim packed to output dim packed |
| 75 | weight = unpack_quantized_values_into_int32( |
| 76 | packed_weight, self.config.weight_type, 0 |
| 77 | ) |
| 78 | weight = pack_quantized_values_into_int32(weight, self.config.weight_type, 1) |
| 79 | # make 16 output channel as a block and transpose to the make |
| 80 | # the block contiguous |
| 81 | weight = ( |
| 82 | weight.view(input_size, -1, 16 // pack_factor) |
| 83 | .permute(1, 0, 2) |
| 84 | .reshape(-1, input_size * 16 // pack_factor) |
| 85 | .contiguous() |
| 86 | ) |
| 87 | getattr(layer, self.w_q_name).data = weight |
| 88 | |
| 89 | # note assumes that |
| 90 | # `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0} |
| 91 | # `weight_scale` is: {input_dim = 0, output_dim = 1} |
| 92 | # `weight_zp` is: {input_dim = 0, output_dim = 1, packed_dim = 1} |
| 93 | def _process_gptq_weights_w4a8(self, layer: torch.nn.Module): |
| 94 | packed_weight = getattr(layer, self.w_q_name) |
| 95 | scales = getattr(layer, self.w_s_name) |
| 96 | group_num = scales.data.size(0) |
| 97 | zp_output_size = scales.data.size(1) // 8 |
| 98 | if self.config.zero_points: |
| 99 | assert self.w_zp_name |
| 100 | packed_zp = getattr(layer, self.w_zp_name) |
| 101 | else: |
| 102 | # w4a8 kernel always requires zp, allocate a fake zp |
| 103 | assert self.w_zp_name |
| 104 | packed_zp = torch.nn.Parameter( |
| 105 | torch.ones(group_num, zp_output_size, dtype=torch.int32) * -2004318072, |
| 106 | requires_grad=False, |
| 107 | ) |
| 108 | setattr(layer, self.w_zp_name, packed_zp) |
| 109 | |
| 110 | # FIXME: some bugs in convert_weight_packed_scale_zp with GPTQ format, |
| 111 | # repack to AWQ weight |
| 112 | weight = unpack_quantized_values_into_int32( |
| 113 | packed_weight, self.config.weight_type, 0 |
| 114 | ) |
| 115 | input_size, output_size = weight.size() |
| 116 | weight = weight.view(input_size, output_size // 8, 8) |
| 117 | weight = weight[:, :, (0, 2, 4, 6, 1, 3, 5, 7)].reshape(input_size, output_size) |
| 118 | weight = pack_quantized_values_into_int32( |
| 119 | weight, self.config.weight_type, 1 |
| 120 | ).contiguous() |
| 121 | |
| 122 | zp = unpack_quantized_values_into_int32(packed_zp, self.config.weight_type, 1) |
| 123 | zp = zp.view(group_num, output_size // 8, 8) |
| 124 | zp = zp[:, :, (0, 2, 4, 6, 1, 3, 5, 7)].reshape(group_num, output_size) |
| 125 | zp = pack_quantized_values_into_int32( |
| 126 | zp, self.config.weight_type, 1 |
| 127 | ).contiguous() |
| 128 | |
| 129 | blocked_w, blocked_zp, blocked_s = ops.convert_weight_packed_scale_zp( |
| 130 | weight, |
| 131 | zp, |
| 132 | scales.data, |
| 133 | ops.CPUQuantAlgo.AWQ, |
| 134 | ) |
| 135 | |
| 136 | if layer.bias is not None: |
| 137 | layer.bias.data = layer.bias.float() |
| 138 | |
| 139 | packed_weight.data = blocked_w |
| 140 | scales.data = blocked_s |
| 141 | packed_zp.data = blocked_zp |
| 142 | |
| 143 | def process_weights_after_loading(self, layer: torch.nn.Module): |
| 144 | if (not self.config.zero_points) and (self.w_zp_name is not None): |
| 145 | setattr(layer, self.w_zp_name, None) |
| 146 | |
| 147 | weights = getattr(layer, self.w_q_name) |
| 148 | # Require GPTQ pack format |
| 149 | assert weights.input_dim == weights.packed_dim |
| 150 | |
| 151 | # Weights in CT format is [output_size, input_size] |
| 152 | if weights.input_dim == 1: |
| 153 | weights.data = weights.t() |
| 154 | |
| 155 | # Scales in CT format is [output_size, group_num] |
| 156 | scales = getattr(layer, self.w_s_name) |
| 157 | if scales.output_dim == 0: |
| 158 | scales.data = scales.t().contiguous() |
| 159 | |
| 160 | # Zero points in CT format is [output_size, group_num] |
| 161 | # Zero points in awq_marlin format is [output_size, group_num] |
| 162 | if self.config.zero_points: |
| 163 | assert self.w_zp_name |
| 164 | zp = getattr(layer, self.w_zp_name) |
| 165 | if zp.output_dim == 0: |
| 166 | zp.data = zp.t().contiguous() |
| 167 | |
| 168 | supports_amx = torch.cpu._is_amx_tile_supported() |
| 169 | supports_riscv = current_platform.get_cpu_architecture() == CpuArchEnum.RISCV |
| 170 | layer.use_w4a8 = ( |
| 171 | envs.VLLM_CPU_INT4_W4A8 |
| 172 | and self.config.act_type == torch.bfloat16 |
| 173 | and (supports_amx or supports_riscv) |
| 174 | ) |
| 175 | # layer.use_w4a8 = False |
| 176 | # AWQ format will be converted to GPTQ format in `AutoAWQMarlinLinearMethod` |
| 177 | if layer.use_w4a8: |
| 178 | self._process_gptq_weights_w4a8(layer) |
| 179 | else: |
| 180 | self._process_gptq_weights_w4a16(layer) |
| 181 | |
| 182 | def apply_weights( |
| 183 | self, |
| 184 | layer: torch.nn.Module, |
| 185 | x: torch.Tensor, |
| 186 | bias: torch.Tensor | None = None, |
| 187 | ) -> torch.Tensor: |
| 188 | w_q, w_s, w_zp = self._get_weight_params(layer) |
| 189 | if layer.use_w4a8: |
| 190 | x = ops.int4_scaled_mm_cpu( |
| 191 | x=x, |
| 192 | w=w_q, |
| 193 | w_zeros=w_zp, |
| 194 | w_scales=w_s, |
| 195 | bias=bias, |
| 196 | ) |
| 197 | else: |
| 198 | x = ops.cpu_gemm_wna16( |
| 199 | input=x, |
| 200 | q_weight=w_q, |
| 201 | scales=w_s, |
| 202 | zeros=w_zp, |
| 203 | bias=bias, |
| 204 | pack_factor=8, # 32 // 4 |
| 205 | isa_hint=layer.isa_hint, |
| 206 | ) |
| 207 | return x |
| 208 | |
| 209 | |
| 210 | def _get_isa_hint(dtype: torch.dtype) -> str: |
| 211 | supports_amx = torch.cpu._is_amx_tile_supported() |
| 212 | if supports_amx and dtype in (torch.bfloat16,): |
| 213 | return "amx" |
| 214 | elif current_platform.get_cpu_architecture() == CpuArchEnum.RISCV: |
| 215 | return "rvv" |
| 216 | elif ( |
| 217 | dtype == torch.bfloat16 |
| 218 | and current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC |
| 219 | ): |
| 220 | return "vsx" |
| 221 | else: |
| 222 | return "vec" |