# SPDX-FileCopyrightText: 2024-2025 ModelCloud.ai # SPDX-FileCopyrightText: 2024-2025 qubitium@modelcloud.ai # SPDX-License-Identifier: Apache-2.0 # Contact: qubitium@modelcloud.ai, x.com/qubitium import math import torch from ...adapter.adapter import Adapter from ...looper.linear_mode import LinearMode from ...quantization import FORMAT, METHOD from ...quantization.awq.utils.packing_utils import ( dequantize_gemm, reverse_awq_order, unpack_awq, ) from ...utils.backend import BACKEND from ...utils.logger import setup_logger from ...utils.torch import TORCH_HAS_FUSED_OPS from . import AWQuantLinear from .torch_fused import Int4PackedOp, TorchFusedLinear, pack_scales_and_zeros log = setup_logger() class TorchFusedAwqLinear(AWQuantLinear): """Torch fused AWQ variant based on GPTQ fused kernels via CPU int4 packing.""" QUANT_TYPE = "torch_fused_awq" SUPPORTS_BACKENDS = [BACKEND.AWQ_TORCH_FUSED] SUPPORTS_METHODS = [METHOD.AWQ] SUPPORTS_FORMATS = {FORMAT.GEMM: 20} # inherit from torch fused SUPPORTS_BITS = TorchFusedLinear.SUPPORTS_BITS SUPPORTS_GROUP_SIZE = TorchFusedLinear.SUPPORTS_GROUP_SIZE SUPPORTS_DESC_ACT = TorchFusedLinear.SUPPORTS_DESC_ACT SUPPORTS_SYM = TorchFusedLinear.SUPPORTS_SYM SUPPORTS_SHARDS = TorchFusedLinear.SUPPORTS_SHARDS SUPPORTS_TRAINING = TorchFusedLinear.SUPPORTS_TRAINING SUPPORTS_AUTO_PADDING = TorchFusedLinear.SUPPORTS_AUTO_PADDING SUPPORTS_IN_FEATURES_DIVISIBLE_BY = TorchFusedLinear.SUPPORTS_IN_FEATURES_DIVISIBLE_BY SUPPORTS_OUT_FEATURES_DIVISIBLE_BY = TorchFusedLinear.SUPPORTS_OUT_FEATURES_DIVISIBLE_BY SUPPORTS_DEVICES = TorchFusedLinear.SUPPORTS_DEVICES SUPPORTS_PLATFORM = TorchFusedLinear.SUPPORTS_PLATFORM SUPPORTS_PACK_DTYPES = TorchFusedLinear.SUPPORTS_PACK_DTYPES SUPPORTS_ADAPTERS = TorchFusedLinear.SUPPORTS_ADAPTERS REQUIRES_FORMAT_V2 = False SUPPORTS_DTYPES = [torch.float32, torch.float16, torch.bfloat16] def __init__( self, bits: int, group_size: int, sym: bool, desc_act: bool, in_features: int, out_features: int, bias: bool = False, pack_dtype: torch.dtype = torch.int32, adapter: Adapter = None, register_buffers: bool = True, **kwargs, ): kwargs.setdefault("backend", BACKEND.AWQ_TORCH_FUSED) super().__init__( bits=bits, group_size=group_size, sym=sym, desc_act=desc_act, in_features=in_features, out_features=out_features, bias=bias, pack_dtype=pack_dtype, adapter=adapter, # Skip base buffer init, we need to manually init buffers for awq register_buffers=False, **kwargs, ) self.linear_mode = None # Create awq buffers if register_buffers: # AWQ packs each input row into pack_factor-wide columns for int4 lanes. pack_cols = max(1, self.out_features // self.pack_factor) qweight_shape = (self.in_features, pack_cols) group_size = max(int(self.group_size), 1) # Each group holds group_size input rows; ceil ensures remaining rows are included. group_rows = max(1, math.ceil(self.in_features / group_size)) self.register_buffer( "qweight", torch.zeros(qweight_shape, dtype=self.pack_dtype), ) self.register_buffer( "qzeros", torch.zeros((group_rows, pack_cols), dtype=self.pack_dtype), ) self.register_buffer( "scales", torch.zeros((group_rows, self.out_features), dtype=torch.float16), ) if bias: self.register_buffer("bias", torch.zeros(self.out_features, dtype=torch.float16)) else: self.bias = None def post_init(self): super().post_init() self.optimize() def optimize(self): if self.optimized: return super().optimize() def prepare_awq_fused_tensors(self, need_zeros: bool = True): self.scales.to(torch.float16).contiguous() iweight, izeros = unpack_awq(self.qweight, self.qzeros, self.bits) iweight, izeros = reverse_awq_order(iweight, izeros, self.bits) max_val = (1 << self.bits) - 1 iweight = torch.bitwise_and(iweight, max_val) if izeros is None: raise RuntimeError("AWQ fused kernel requires zero points.") izeros = torch.bitwise_and(izeros, max_val) if need_zeros: zero_offset = 1 << (self.bits - 1) zeros = (zero_offset - izeros.reshape_as(self.scales)) * self.scales gptq_qweight = self.pack_awq_qweight(iweight) gptq_qzeros = self.pack_awq_qzeros(izeros) return gptq_qweight, gptq_qzeros, self.scales, zeros if need_zeros else None def pack_awq_qweight(self, iweight: torch.Tensor) -> torch.Tensor: in_features, out_features = iweight.shape pack_factor = int(self.pack_factor) if in_features % pack_factor != 0: raise ValueError( f"AWQ in_features={in_features} must be divisible by pack_factor={pack_factor}." ) rows = iweight.view(in_features // pack_factor, pack_factor, out_features) packed = torch.zeros( (rows.shape[0], out_features), dtype=self.pack_dtype, device=iweight.device, ) shifts = range(0, pack_factor * self.bits, self.bits) for lane, shift in enumerate(shifts): packed |= rows[:, lane, :].to(torch.int32) << shift return packed.contiguous() def pack_awq_qzeros(self, izeros: torch.Tensor) -> torch.Tensor: pack_factor = int(self.pack_factor) if izeros.shape[1] % pack_factor != 0: raise ValueError( f"AWQ qzeros dimension {izeros.shape[1]} must be divisible by pack_factor={pack_factor}." ) cols = izeros.view(izeros.shape[0], izeros.shape[1] // pack_factor, pack_factor) packed = torch.zeros( (cols.shape[0], cols.shape[1]), dtype=self.pack_dtype, device=izeros.device, ) shifts = range(0, pack_factor * self.bits, self.bits) for lane, shift in enumerate(shifts): packed |= cols[:, :, lane].to(torch.int32) << shift return packed.contiguous() def transform_cpu_awq(self, dtype): # Unpack AWQ weights directly to integer form iweight, izeros = unpack_awq(self.qweight, self.qzeros, self.bits) iweight, izeros = reverse_awq_order(iweight, izeros, self.bits) max_val = (1 << self.bits) - 1 iweight = torch.bitwise_and(iweight, max_val) izeros = torch.bitwise_and(izeros, max_val) # Compute zeros: (zero_offset - izeros) * scales zero_offset = 1 << (self.bits - 1) zeros = (zero_offset - izeros.reshape_as(self.scales)) * self.scales # AWQ has no g_idx — weights are already in natural order, just transpose # iweight: (in_features, out_features) -> (out_features, in_features) for int4pack weight = iweight.t().contiguous() self.qweight = torch.ops.aten._convert_weight_to_int4pack_for_cpu(weight.int(), 1).contiguous() self.scales = self.scales.to(dtype=dtype).contiguous() self.qzeros = zeros.to(device=self.qweight.device, dtype=dtype).contiguous() self.scales_and_zeros = pack_scales_and_zeros(self.scales, self.qzeros) def transform_xpu_awq(self, dtype): # Unpack AWQ weights directly to integer form iweight, izeros = unpack_awq(self.qweight, self.qzeros, self.bits) iweight, izeros = reverse_awq_order(iweight, izeros, self.bits) max_val = (1 << self.bits) - 1 iweight = torch.bitwise_and(iweight, max_val) izeros = torch.bitwise_and(izeros, max_val) self.scales = self.scales.to(dtype).contiguous() # AWQ has no g_idx — pack weight directly for XPU without reordering # iweight: (in_features, out_features) -> transpose for XPU packing weight = iweight.t().contiguous() packed = torch.zeros(weight.shape[0], weight.shape[1] // self.pack_factor, dtype=torch.int32, device=weight.device) for col in range(weight.shape[1] // self.pack_factor): for i in range(self.pack_factor): packed_col = weight[:, col * self.pack_factor + i].to(torch.int32) packed[:, col] |= packed_col << (i * self.bits) self.qweight = packed.contiguous() self.qzeros = izeros.contiguous() def transform_cpu(self, dtype): self.transform_cpu_awq(dtype) @torch.no_grad() def _fused_op_forward(self, x): # AWQ has no g_idx reordering — skip ret_idx if x.device.type == "xpu": out = torch.ops.aten._weight_int4pack_mm_with_scales_and_zeros( x, self.qweight, self.group_size, self.scales, self.qzeros ) elif x.device.type == "cpu": out = self.torch_fused_op(x) else: raise NotImplementedError return out def awq_weight_dequantize(self, device, dtype): return dequantize_gemm( qweight=self.qweight, qzeros=self.qzeros, scales=self.scales, bits=self.bits, group_size=self.group_size, ).to(device=device, dtype=dtype) def transform(self, dtype, device): if device == "cpu": self.transform_cpu(dtype) elif device == "xpu": self.transform_xpu_awq(dtype) else: raise NotImplementedError( "TorchFusedAwqLinear only supports fused transforms on CPU or XPU devices." ) def forward(self, x: torch.Tensor): out_shape = x.shape[:-1] + (self.out_features,) x_flat = x.reshape(-1, x.shape[-1]) input_dtype = x_flat.dtype self.assert_supported_dtype(input_dtype) if input_dtype == torch.float32: x_flat = x_flat.to(torch.bfloat16) if ( not self.training and not x_flat.requires_grad and self.linear_mode is None and TORCH_HAS_FUSED_OPS ): self.transform(x_flat.dtype, x_flat.device.type) self.linear_mode = LinearMode.INFERENCE if x_flat.device.type == "cpu": self.torch_fused_op = Int4PackedOp( self.qweight, self.scales_and_zeros, self.group_size ).eval() import torch._inductor.config as config config.freezing = True config.max_autotune = True elif self.linear_mode is None: self.linear_mode = LinearMode.TRAIN if self.linear_mode == LinearMode.INFERENCE: # log.debug("awq calling fused op") out = self._fused_op_forward(x_flat) else: # log.debug("awq dense path") weight = self.awq_weight_dequantize(device=x_flat.device, dtype=x_flat.dtype) out = torch.matmul(x_flat, weight) if self.bias is not None: out.add_(self.bias) if self.adapter: out = self.adapter.apply(x=x_flat, out=out) if input_dtype == torch.float32: out = out.to(torch.float32) return out.reshape(out_shape) def assert_supported_dtype(self, dtype: torch.dtype): if dtype not in self.SUPPORTS_DTYPES: supported = ", ".join(str(d) for d in self.SUPPORTS_DTYPES) raise TypeError( f"{self.__class__.__name__} only supports input dtypes [{supported}], but received {dtype}." ) __all__ = ["TorchFusedAwqLinear"]