| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import torch |
| 6 | from torch.nn.parameter import Parameter |
| 7 | |
| 8 | from vllm.logger import init_logger |
| 9 | from vllm.model_executor.layers.quantization.utils import replace_parameter |
| 10 | from vllm.platforms import current_platform |
| 11 | from vllm.scalar_type import scalar_types |
| 12 | |
| 13 | from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig |
| 14 | |
| 15 | _XPUWNA16_SUPPORTED_QUANT_TYPES = (scalar_types.uint4, scalar_types.uint4b8) |
| 16 | |
| 17 | logger = init_logger(__name__) |
| 18 | |
| 19 | |
| 20 | class XPUwNa16LinearKernel(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_xpu(): |
| 28 | return False, "XPUwNa16 only supported on XPU" |
| 29 | |
| 30 | if c.act_type != torch.bfloat16 and c.act_type != torch.float16: |
| 31 | return False, "XPUwNa16 only supports BF16/FP16 activations" |
| 32 | |
| 33 | if c.weight_type not in _XPUWNA16_SUPPORTED_QUANT_TYPES: |
| 34 | return ( |
| 35 | False, |
| 36 | f"Quant type ({c.weight_type}) not supported by " |
| 37 | "XPUwNa16, supported types are: " |
| 38 | f"{_XPUWNA16_SUPPORTED_QUANT_TYPES}", |
| 39 | ) |
| 40 | if c.group_size != -1 and c.group_size % 32 != 0: |
| 41 | return ( |
| 42 | False, |
| 43 | f"Group size ({c.group_size}) not supported by " |
| 44 | "XPUwNa16, supported group sizes are multiples of 32", |
| 45 | ) |
| 46 | |
| 47 | if c.partition_weight_shape[0] % 32 != 0: |
| 48 | return ( |
| 49 | False, |
| 50 | f"Input size ({c.partition_weight_shape[0]}) not supported by " |
| 51 | "XPUwNa16, supported sizes are multiples of 32", |
| 52 | ) |
| 53 | |
| 54 | return True, None |
| 55 | |
| 56 | def process_weights_after_loading(self, layer: torch.nn.Module): |
| 57 | # Default name since marlin requires empty parameter for zp, |
| 58 | # TODO: remove this requirement from marlin (allow optional tensors) |
| 59 | if self.w_zp_name is None: |
| 60 | self.w_zp_name = "w_zp" |
| 61 | |
| 62 | need_transpose = False |
| 63 | qweight_shape = getattr(layer, self.w_q_name).shape |
| 64 | scale_shape = getattr(layer, self.w_s_name).shape |
| 65 | # gptq marlin and compressed tensors wna16 expect different default |
| 66 | # layouts for weight and scale, so we check the shapes to determine |
| 67 | # if we need to transpose |
| 68 | if qweight_shape[0] != scale_shape[0]: |
| 69 | need_transpose = True |
| 70 | |
| 71 | if need_transpose: |
| 72 | getattr(layer, self.w_q_name).data = ( |
| 73 | getattr(layer, self.w_q_name).data.t().contiguous() |
| 74 | ) |
| 75 | getattr(layer, self.w_s_name).data = getattr(layer, self.w_s_name).data |
| 76 | else: |
| 77 | getattr(layer, self.w_s_name).data = ( |
| 78 | getattr(layer, self.w_s_name).data.t().contiguous() |
| 79 | ) |
| 80 | |
| 81 | if self.config.zero_points: |
| 82 | # (FIXME): maybe zero points should also be transposed. |
| 83 | getattr(layer, self.w_zp_name).data = ( |
| 84 | getattr(layer, self.w_zp_name).data.t().contiguous() |
| 85 | ) |
| 86 | else: |
| 87 | weight_zero_point = torch.Tensor([8]).to(torch.int8).to("xpu") |
| 88 | setattr( |
| 89 | layer, self.w_zp_name, Parameter(weight_zero_point, requires_grad=False) |
| 90 | ) |
| 91 | |
| 92 | def apply_weights( |
| 93 | self, |
| 94 | layer: torch.nn.Module, |
| 95 | x: torch.Tensor, |
| 96 | bias: torch.Tensor | None = None, |
| 97 | ) -> torch.Tensor: |
| 98 | reshaped_x = x.reshape(-1, x.shape[-1]) |
| 99 | w_q, w_s, w_zp = self._get_weight_params(layer) |
| 100 | out = torch.ops._xpu_C.int4_gemm_w4a16( |
| 101 | reshaped_x, |
| 102 | w_q.t(), |
| 103 | bias if bias is not None else None, |
| 104 | w_s, |
| 105 | w_zp, |
| 106 | self.config.group_size, |
| 107 | None, # Retained by the external XPU op ABI. |
| 108 | ) |
| 109 | return out |
| 110 | |
| 111 | |
| 112 | class XPUW4A8IntLinearKernel(MPLinearKernel): |
| 113 | """XPU kernel for W4A8 integer quantization using oneDNN int4_gemm_w4a8. |
| 114 | |
| 115 | Weights are symmetric group-quantized int4 packed as uint4. |
| 116 | Activations are dynamically quantized per-token to symmetric int8. |
| 117 | """ |
| 118 | |
| 119 | @classmethod |
| 120 | def get_min_capability(cls) -> int: |
| 121 | return -1 |
| 122 | |
| 123 | @classmethod |
| 124 | def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]: |
| 125 | if not current_platform.is_xpu(): |
| 126 | return False, "XPUW4A8Int only supported on XPU" |
| 127 | if c.act_type not in (torch.bfloat16, torch.float16): |
| 128 | return False, "XPUW4A8Int requires BF16/FP16 activations" |
| 129 | if c.weight_type != scalar_types.int4: |
| 130 | return ( |
| 131 | False, |
| 132 | f"XPUW4A8Int requires int4 weights, got {c.weight_type}", |
| 133 | ) |
| 134 | if c.zero_points: |
| 135 | return False, "XPUW4A8Int only supports symmetric weight quantization" |
| 136 | if c.group_size != -1 and c.group_size % 32 != 0: |
| 137 | return ( |
| 138 | False, |
| 139 | f"Group size ({c.group_size}) not supported by XPUW4A8Int, " |
| 140 | "must be a multiple of 32", |
| 141 | ) |
| 142 | in_size, out_size = c.partition_weight_shape |
| 143 | if in_size % 8 != 0 or out_size % 8 != 0: |
| 144 | return ( |
| 145 | False, |
| 146 | f"in/out sizes ({in_size}, {out_size}) must be multiples of 8", |
| 147 | ) |
| 148 | |
| 149 | if c.act_type != torch.float16: |
| 150 | logger.warning_once( |
| 151 | "XPUW4A8IntLinearKernel is running with model dtype %s, " |
| 152 | "but int4_gemm_w4a8 produces float16 output. Recommend " |
| 153 | "setting --dtype float16 for best performance.", |
| 154 | c.act_type, |
| 155 | ) |
| 156 | |
| 157 | return True, None |
| 158 | |
| 159 | def _pack_int4_weight(self, w: torch.Tensor) -> torch.Tensor: |
| 160 | # w is [N, K] int8 with values in [-8, 7] |
| 161 | w_u4 = w.to(torch.int32) + 8 # shift to [0, 15] |
| 162 | w_u4 = w_u4.reshape(w.shape[0], w.shape[1] // 8, 8) # [N, K/8, 8] |
| 163 | shifts = torch.arange(0, 32, 4, dtype=torch.int32, device=w.device) |
| 164 | packed = ((w_u4 & 0xF) << shifts[None, None, :]).sum(dim=2).to(torch.int32) |
| 165 | return packed |
| 166 | |
| 167 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 168 | layer.weight_scale.data = layer.weight_scale.data.t().contiguous() |
| 169 | |
| 170 | device = layer.weight_packed.device |
| 171 | # TODO: support asymmetric quantization |
| 172 | weight_zero_point = torch.tensor([8], dtype=torch.int8, device=device) |
| 173 | layer.weight_zero_point = Parameter(weight_zero_point, requires_grad=False) |
| 174 | |
| 175 | # weight_packed is [out, in] int8, signed int4 values in [-8, 7] |
| 176 | w = layer.weight_packed.data # [out, in] |
| 177 | |
| 178 | # TODO: implement asym case |
| 179 | packed = self._pack_int4_weight(w) # [out, in/8] packed uint4 |
| 180 | |
| 181 | replace_parameter( |
| 182 | layer, |
| 183 | self.w_q_name, |
| 184 | torch.nn.Parameter(packed, requires_grad=False), |
| 185 | ) |
| 186 | |
| 187 | # Free the original unpacked int8 weight (still registered as "weight") |
| 188 | # to avoid double-storing both int8 [N, K] and int32 [N, K/8] in memory. |
| 189 | layer.register_parameter("weight", None) |
| 190 | |
| 191 | def apply_weights( |
| 192 | self, |
| 193 | layer: torch.nn.Module, |
| 194 | x: torch.Tensor, |
| 195 | bias: torch.Tensor | None = None, |
| 196 | ) -> torch.Tensor: |
| 197 | reshaped_x = x.reshape(-1, x.shape[-1]) # [M, K] |
| 198 | from vllm._xpu_ops import xpu_ops as ops |
| 199 | |
| 200 | # TODO: static and asymmetric quantization case |
| 201 | # Common code for CompressedTensorsW4A8Int does not read act symmetry data |
| 202 | quant_x, x_scale, x_zero = ops.dynamic_per_token_int8_quant_ref( |
| 203 | reshaped_x, True, 8 |
| 204 | ) |
| 205 | |
| 206 | out = torch.ops._xpu_C.int4_gemm_w4a8( |
| 207 | quant_x, |
| 208 | x_scale, |
| 209 | x_zero, |
| 210 | layer.weight_packed.t(), |
| 211 | layer.weight_scale, |
| 212 | layer.weight_zero_point, |
| 213 | self.config.group_size, |
| 214 | None, # Retained by the external XPU op ABI. |
| 215 | bias, |
| 216 | ) |
| 217 | |
| 218 | return out.to(x.dtype) |