# 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 torch from ...adapter.adapter import Adapter, Lora from ...models._const import DEVICE, PLATFORM from ...nn_modules.qlinear import AWQuantLinear from ...quantization import FORMAT, METHOD from ...quantization.awq.utils.packing_utils import unpack_reorder_pack from ...utils.backend import BACKEND from ...utils.exllamav2 import ( ScratchSpace, exllamav2_awq_gemm_half_q_half, exllamav2_awq_make_q_matrix, exllamav2_awq_runtime_available, exllamav2_awq_runtime_error, ) from ...utils.logger import setup_logger log = setup_logger() class AwqExllamaV2Linear(AWQuantLinear): SUPPORTS_BACKENDS = [BACKEND.AWQ_EXLLAMA_V2] SUPPORTS_METHODS = [METHOD.AWQ] SUPPORTS_FORMATS = {FORMAT.GEMM: 80} SUPPORTS_BITS = [4] # TODO: intel is reporting v2 has accuracy issues with group-size == 16 for this kernel # disable for now until we can validate this issue: ref https://github.com/ModelCloud/GPTQModel/issues/1515 SUPPORTS_GROUP_SIZE = [-1, 32, 64, 128] SUPPORTS_DESC_ACT = [True, False] SUPPORTS_SYM = [True, False] SUPPORTS_SHARDS = True SUPPORTS_TRAINING = False SUPPORTS_AUTO_PADDING = False SUPPORTS_IN_FEATURES_DIVISIBLE_BY = [32] SUPPORTS_OUT_FEATURES_DIVISIBLE_BY = [32] SUPPORTS_DEVICES = [DEVICE.CUDA] # ROCm has broken accuracies issues SUPPORTS_PLATFORM = [PLATFORM.LINUX, PLATFORM.WIN32] SUPPORTS_PACK_DTYPES = [torch.int32] SUPPORTS_ADAPTERS = [Lora] SUPPORTS_DTYPES = [torch.float16, torch.bfloat16] REQUIRES_FORMAT_V2 = True # for transformers/optimum tests compat QUANT_TYPE = "awq_exllamav2" 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, **kwargs, ): 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, backend=kwargs.pop("backend", BACKEND.AWQ_EXLLAMA_V2), adapter=adapter, **kwargs) @classmethod def validate_once(cls): # Validate through the shared JIT loader so import-time checks match the # runtime build path used on first kernel access. if not exllamav2_awq_runtime_available(): return False, ImportError(exllamav2_awq_runtime_error()) return True, None def ext_make_q_matrix_awq(self, qweight, qzeros, scales, temp_dq) -> int: runtime_scales = scales if scales.dtype == torch.float16 else scales.to(dtype=torch.float16) return exllamav2_awq_make_q_matrix( qweight, None, None, None, None, None, qzeros, runtime_scales, None, temp_dq, ) def post_init(self, scratch_space: ScratchSpace): # if self.padded_infeatures != self.in_features: # self.qweight.resize_(self.padded_infeatures // self.pack_dtype_bits * self.bits, self.out_features) # self.qzeros.resize_( # math.ceil(self.padded_infeatures / self.group_size), # self.out_features // self.pack_dtype_bits * self.bits # ) # self.scales.resize_((math.ceil(self.padded_infeatures / self.group_size), self.out_features), ) # self.g_idx = torch.tensor([i // self.group_size for i in range(self.padded_infeatures)], dtype=torch.int32, # device=self.g_idx.device) if not exllamav2_awq_runtime_available(): raise ModuleNotFoundError("ExLlamaV2 AWQ torch.ops kernels are not properly installed. Error: " + exllamav2_awq_runtime_error()) # awq only accepts float16 self.scales = self.scales.to(dtype=torch.float16) assert self.qweight.device.type == "cuda" assert self.qweight.device.index is not None self.qweight, self.qzeros = unpack_reorder_pack( self.qweight, self.qzeros, self.bits ) temp_dq_size = self.temp_dq_size() temp_dq = scratch_space.get_slice(temp_dq_size) self.q_handle = self.ext_make_q_matrix_awq(self.qweight, self.qzeros, self.scales, temp_dq) super().post_init() def forward(self, x: torch.Tensor): if self.input_rows(x) == 0: return self.empty_linear_output(x) assert self.q_handle is not None, ( "module.post_init() must be called before module.forward(). " ) if not exllamav2_awq_runtime_available(): raise ModuleNotFoundError("ExLlamaV2 AWQ torch.ops kernels are not properly installed. Error: " + exllamav2_awq_runtime_error()) input_dtype = x.dtype out_shape = x.shape[:-1] + (self.out_features,) if input_dtype != torch.float16: x = x.to(dtype=torch.float16) x = x.view(-1, x.shape[-1]) out = torch.empty( (x.shape[0], self.out_features), dtype=torch.float16, device=x.device, ) exllamav2_awq_gemm_half_q_half(x, self.q_handle, out, False) if input_dtype != torch.float16: out = out.to(dtype=input_dtype) if self.bias is not None: out.add_(self.bias) if self.adapter: out = self.adapter.apply(x=x, out=out) return out.view(out_shape) def temp_dq_size(self): """ Returns the size of the temporary buffer required for the dq kernel. """ return self.in_features * self.out_features * 2 + 128 def temp_fwd_size(self, max_input_len, max_batch_size): """ Returns the size of the temporary buffer required for the fwd kernel. """ return self.out_features * max_input_len * max_batch_size * 4 + 128 def scratch_space_fixed(self, max_input_len=2048, max_batch_size=8): """ Returns the size of the fixed scratch space required for the kernel. """ return self.temp_dq_size() + self.temp_fwd_size(max_input_len, max_batch_size) def next_multiple(x, multiple): return ((x + multiple - 1) // multiple) * multiple __all__ = ["AwqExllamaV2Linear"]