| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Zentorch int4 weight-quantized linear kernels for AMD Zen CPUs. |
| 4 | |
| 5 | Selected by ``choose_mp_linear_kernel`` ahead of the generic oneDNN-backed |
| 6 | ``CPUWNA16LinearKernel``. When ``can_implement`` rejects a layer, the selector |
| 7 | falls through to the next kernel in ``_POSSIBLE_KERNELS[PlatformEnum.CPU]``. |
| 8 | """ |
| 9 | |
| 10 | import torch |
| 11 | |
| 12 | from vllm import envs |
| 13 | from vllm.logger import init_logger |
| 14 | from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op |
| 15 | from vllm.platforms import current_platform |
| 16 | from vllm.scalar_type import scalar_types |
| 17 | |
| 18 | from .cpu import _CPUWNA16_SUPPORTED_QUANT_TYPES, CPUWNA16LinearKernel |
| 19 | from .MPLinearKernel import MPLinearLayerConfig |
| 20 | |
| 21 | logger = init_logger(__name__) |
| 22 | |
| 23 | |
| 24 | def _import_unpack_from_int32(): |
| 25 | """Import compressed-tensors' ``unpack_from_int32`` across versions.""" |
| 26 | try: |
| 27 | from compressed_tensors.compressors.pack_quantized.helpers import ( |
| 28 | unpack_from_int32, |
| 29 | ) |
| 30 | except ImportError: |
| 31 | from compressed_tensors.compressors.quantized_compressors.pack_quantized import ( # type: ignore[import-not-found] # noqa: E501 |
| 32 | unpack_from_int32, |
| 33 | ) |
| 34 | return unpack_from_int32 |
| 35 | |
| 36 | |
| 37 | class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel): |
| 38 | """Int4 kernel backed by ``zentorch_dynamic_qlinear`` (DA8W4) or |
| 39 | ``zentorch_woq_linear`` (W4A16).""" |
| 40 | |
| 41 | @classmethod |
| 42 | def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]: |
| 43 | if not current_platform.is_zen_cpu(): |
| 44 | return False, "ZentorchWNA16 requires an AMD Zen CPU." |
| 45 | |
| 46 | # Not super(): its N/K multiple-of-32 rule is a oneDNN packing |
| 47 | # constraint the zentorch repack does not share. |
| 48 | if c.weight_type not in _CPUWNA16_SUPPORTED_QUANT_TYPES: |
| 49 | return ( |
| 50 | False, |
| 51 | f"Quant type ({c.weight_type}) not supported by " |
| 52 | "CPUWNA16, supported types are: " |
| 53 | f"{_CPUWNA16_SUPPORTED_QUANT_TYPES}", |
| 54 | ) |
| 55 | |
| 56 | # -1 is per-channel, one group spanning K. Anything else has to divide |
| 57 | # K, which is stricter than the parent's multiple-of-2 rule. |
| 58 | if c.group_size != -1: |
| 59 | in_features = c.partition_weight_shape[0] |
| 60 | if c.group_size <= 0 or in_features % c.group_size != 0: |
| 61 | return ( |
| 62 | False, |
| 63 | f"Group size ({c.group_size}) must divide input size " |
| 64 | f"({in_features})", |
| 65 | ) |
| 66 | |
| 67 | if not has_zentorch_op(["zentorch_woq_repack_weight", "zentorch_woq_linear"]): |
| 68 | return ( |
| 69 | False, |
| 70 | "torch.ops.zentorch.{zentorch_woq_repack_weight, " |
| 71 | "zentorch_woq_linear} are not registered.", |
| 72 | ) |
| 73 | |
| 74 | return True, None |
| 75 | |
| 76 | def _zentorch_woq_eligible(self, layer: torch.nn.Module) -> bool: |
| 77 | """Eligibility predicate for the zentorch W4A16 GPTQ fast path. |
| 78 | |
| 79 | Constraints (any failure -> ``cpu_gemm_wna16`` path via ``super()`` |
| 80 | with ``layer`` untouched). |
| 81 | """ |
| 82 | weight_packed = getattr(layer, self.w_q_name, None) |
| 83 | weight_scale = getattr(layer, self.w_s_name, None) |
| 84 | if weight_packed is None or weight_scale is None: |
| 85 | return False |
| 86 | |
| 87 | bits = self.config.weight_type.mantissa |
| 88 | pack_factor = torch.iinfo(weight_packed.dtype).bits // bits |
| 89 | # 4-bit -> 8 values per int32; |
| 90 | if pack_factor != 8: |
| 91 | return False |
| 92 | |
| 93 | # GPTQ-only. AWQ packs along the output dim instead. |
| 94 | in_dim = getattr(weight_packed, "input_dim", None) |
| 95 | pk_dim = getattr(weight_packed, "packed_dim", None) |
| 96 | if in_dim is None or pk_dim is None or in_dim != pk_dim: |
| 97 | return False |
| 98 | |
| 99 | is_ct_format = in_dim == pk_dim == 1 |
| 100 | if not is_ct_format: |
| 101 | return False |
| 102 | |
| 103 | if weight_packed.dim() != 2 or weight_scale.dim() != 2: |
| 104 | return False |
| 105 | |
| 106 | # 4-bit -> 8 values per int32; in_features must be divisible by num_groups. |
| 107 | in_features = weight_packed.shape[1] * 8 |
| 108 | num_groups = weight_scale.shape[1] |
| 109 | return num_groups > 0 and in_features % num_groups == 0 |
| 110 | |
| 111 | def _maybe_process_da8w4_weights(self, layer: torch.nn.Module) -> bool: |
| 112 | """Repack ``layer`` into the DA8W4 layout, or leave it untouched. |
| 113 | |
| 114 | Returns False without mutating the layer when it is not eligible, so |
| 115 | the caller falls through to the W4A16 paths. |
| 116 | """ |
| 117 | if not envs.VLLM_CPU_INT4_W4A8: |
| 118 | return False |
| 119 | |
| 120 | if not has_zentorch_op(["zentorch_dynamic_qlinear"]): |
| 121 | return False |
| 122 | |
| 123 | # DA8W4 is symmetric-only, and the kernel rejects f32 activations. |
| 124 | if self.config.zero_points or self.config.weight_type == scalar_types.uint4: |
| 125 | return False |
| 126 | if self.config.act_type != torch.bfloat16: |
| 127 | return False |
| 128 | |
| 129 | if not self._zentorch_woq_eligible(layer): |
| 130 | return False |
| 131 | |
| 132 | weight_q = getattr(layer, self.w_q_name) |
| 133 | weight_s = getattr(layer, self.w_s_name) |
| 134 | weight_packed = weight_q.data if hasattr(weight_q, "data") else weight_q |
| 135 | weight_scale = weight_s.data if hasattr(weight_s, "data") else weight_s |
| 136 | |
| 137 | bits = self.config.weight_type.mantissa |
| 138 | pack_factor = torch.iinfo(weight_packed.dtype).bits // bits |
| 139 | out_features, num_groups = weight_scale.shape[0], weight_scale.shape[1] |
| 140 | in_features = weight_packed.shape[1] * pack_factor |
| 141 | |
| 142 | # AOCL sym_quant requires K/G to be a multiple of 4; K must be even to |
| 143 | # pack two s4 values per byte. |
| 144 | if (in_features // num_groups) % 4 != 0 or in_features % 2 != 0: |
| 145 | return False |
| 146 | |
| 147 | if self.w_zp_name is not None: |
| 148 | setattr(layer, self.w_zp_name, None) |
| 149 | |
| 150 | unpack_from_int32 = _import_unpack_from_int32() |
| 151 | |
| 152 | weight_unpacked = unpack_from_int32( |
| 153 | weight_packed, |
| 154 | bits, |
| 155 | torch.Size([out_features, in_features]), |
| 156 | packed_dim=weight_q.packed_dim, |
| 157 | ) |
| 158 | # The WOQ repack packs 8 int4 per int32, which on little-endian is the |
| 159 | # same byte stream as s4 [N, K/2]; the kernel takes either view. |
| 160 | layer._zentorch_da8w4_packed = ( |
| 161 | torch.ops.zentorch.zentorch_woq_repack_weight.default( |
| 162 | weight_unpacked.to(torch.int8).contiguous() |
| 163 | ).view(torch.int8) |
| 164 | ) |
| 165 | # CT stores scales as [N, G]; DA8W4 wants per-group {G, N}. |
| 166 | layer._zentorch_da8w4_scale = weight_scale.t().to(torch.bfloat16).contiguous() |
| 167 | |
| 168 | for param_name in (self.w_q_name, self.w_s_name): |
| 169 | param = getattr(layer, param_name, None) |
| 170 | if param is not None and hasattr(param, "data"): |
| 171 | param.data = torch.empty(0) |
| 172 | |
| 173 | layer._zentorch_kind = "compressed_tensors_w4a8_da8w4" |
| 174 | layer._zentorch_da8w4 = True |
| 175 | layer._zentorch_processed_weights = True |
| 176 | logger.info_once( |
| 177 | "[zen_cpu] Using zentorch_dynamic_qlinear for DA8W4 (W4A8) " |
| 178 | "(weight_type=%s, group_size=%d)", |
| 179 | self.config.weight_type, |
| 180 | in_features // num_groups, |
| 181 | ) |
| 182 | return True |
| 183 | |
| 184 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 185 | """Repack CT GPTQ weights into the zentorch DA8W4 or WOQ layout. |
| 186 | |
| 187 | Falls back to ``CPUWNA16LinearKernel.process_weights_after_loading`` |
| 188 | via ``super()`` when the layer doesn't satisfy |
| 189 | ``_zentorch_woq_eligible``. |
| 190 | |
| 191 | On success, ``layer._zentorch_processed_weights`` is set to ``True`` |
| 192 | """ |
| 193 | if getattr(layer, "_zentorch_processed_weights", False): |
| 194 | return |
| 195 | |
| 196 | if self._maybe_process_da8w4_weights(layer): |
| 197 | return |
| 198 | |
| 199 | if not self._zentorch_woq_eligible(layer): |
| 200 | logger.info_once( |
| 201 | "[zen_cpu] ZentorchWNA16 fast path not eligible for this " |
| 202 | "layer (AWQ pack layout or non-int32 storage); " |
| 203 | "falling back to CPUWNA16LinearKernel (cpu_gemm_wna16)." |
| 204 | ) |
| 205 | super().process_weights_after_loading(layer) |
| 206 | return |
| 207 | |
| 208 | if (not self.config.zero_points) and (self.w_zp_name is not None): |
| 209 | setattr(layer, self.w_zp_name, None) |
| 210 | |
| 211 | weight_q = getattr(layer, self.w_q_name) |
| 212 | weight_s = getattr(layer, self.w_s_name) |
| 213 | weight_packed = weight_q.data if hasattr(weight_q, "data") else weight_q |
| 214 | weight_scale = weight_s.data if hasattr(weight_s, "data") else weight_s |
| 215 | |
| 216 | bits = self.config.weight_type.mantissa |
| 217 | pack_factor = torch.iinfo(weight_packed.dtype).bits // bits |
| 218 | out_features, num_groups = weight_scale.shape[0], weight_scale.shape[1] |
| 219 | in_features = weight_packed.shape[1] * pack_factor |
| 220 | original_shape = torch.Size([out_features, in_features]) |
| 221 | unpack_from_int32 = _import_unpack_from_int32() |
| 222 | repack_op = torch.ops.zentorch.zentorch_woq_repack_weight.default |
| 223 | |
| 224 | weight_unpacked = unpack_from_int32( |
| 225 | weight_packed, |
| 226 | bits, |
| 227 | original_shape, |
| 228 | packed_dim=weight_q.packed_dim, |
| 229 | ) |
| 230 | |
| 231 | zp_param = ( |
| 232 | getattr(layer, self.w_zp_name, None) if self.w_zp_name is not None else None |
| 233 | ) |
| 234 | needs_unsigned_offset = self.config.weight_type == scalar_types.uint4 |
| 235 | |
| 236 | if needs_unsigned_offset: |
| 237 | weight_unpacked = (weight_unpacked.to(torch.int32) + 8).clamp(0, 15) |
| 238 | repacked = repack_op(weight_unpacked.to(torch.int8).contiguous()) |
| 239 | |
| 240 | if zp_param is None: |
| 241 | zp_tc = None |
| 242 | else: |
| 243 | zp_tensor = zp_param.data if hasattr(zp_param, "data") else zp_param |
| 244 | zp = unpack_from_int32( |
| 245 | zp_tensor, |
| 246 | bits, |
| 247 | (out_features, num_groups), |
| 248 | packed_dim=zp_param.packed_dim, |
| 249 | ) |
| 250 | if needs_unsigned_offset: |
| 251 | zp = (zp.to(torch.int32) + 8).clamp(0, 15) |
| 252 | zp_tc = zp.to(torch.int8).t().contiguous() |
| 253 | |
| 254 | layer._zentorch_woq_packed = repacked.t() |
| 255 | layer._zentorch_woq_scale = weight_scale.t().contiguous() |
| 256 | layer._zentorch_woq_zero_point = zp_tc |
| 257 | |
| 258 | for param_name in (self.w_q_name, self.w_s_name, self.w_zp_name): |
| 259 | if param_name is None: |
| 260 | continue |
| 261 | param = getattr(layer, param_name, None) |
| 262 | if param is None: |
| 263 | continue |
| 264 | if hasattr(param, "data"): |
| 265 | param.data = torch.empty(0) |
| 266 | else: |
| 267 | setattr(layer, param_name, torch.empty(0)) |
| 268 | |
| 269 | layer._zentorch_kind = "compressed_tensors_w4a16_gptq" |
| 270 | layer._zentorch_processed_weights = True |
| 271 | logger.info_once( |
| 272 | "[zen_cpu] Using zentorch_woq_linear for W4A16 GPTQ " |
| 273 | "(weight_type=%s, has_zp=%s)", |
| 274 | self.config.weight_type, |
| 275 | zp_tc is not None, |
| 276 | ) |
| 277 | |
| 278 | def apply_weights( |
| 279 | self, |
| 280 | layer: torch.nn.Module, |
| 281 | x: torch.Tensor, |
| 282 | bias: torch.Tensor | None = None, |
| 283 | ) -> torch.Tensor: |
| 284 | if getattr(layer, "_zentorch_da8w4", False): |
| 285 | # The kernel reads bias through a raw pointer, so it must be |
| 286 | # contiguous; the weight and scales already are. |
| 287 | return torch.ops.zentorch.zentorch_dynamic_qlinear.default( |
| 288 | x if x.dtype == torch.bfloat16 else x.to(torch.bfloat16), |
| 289 | layer._zentorch_da8w4_packed, |
| 290 | layer._zentorch_da8w4_scale, |
| 291 | bias.contiguous() if bias is not None else None, |
| 292 | ) |
| 293 | |
| 294 | if getattr(layer, "_zentorch_processed_weights", False): |
| 295 | return torch.ops.zentorch.zentorch_woq_linear.default( |
| 296 | x, |
| 297 | layer._zentorch_woq_packed, |
| 298 | layer._zentorch_woq_scale, |
| 299 | layer._zentorch_woq_zero_point, |
| 300 | bias, |
| 301 | ) |
| 302 | return super().apply_weights(layer, x, bias) |
| 303 | |
| 304 | |
| 305 | __all__ = ["ZentorchWNA16LinearKernel"] |