Cleanup SDNQ and skip transpose on packed int8 matmul

modules/sdnq/__init__.py

@@ -160,8 +160,7 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
 
         if use_quantized_matmul:
             scale = scale.transpose(0,1)
-            if dtype_dict[weights_dtype]["num_bits"] == 8:
-                layer.weight.data = layer.weight.transpose(0,1)
+            layer.weight.data = layer.weight.transpose(0,1)
             if not dtype_dict[weights_dtype]["is_integer"]:
                 stride = layer.weight.stride()
                 if stride[0] > stride[1] and stride[1] == 1:

modules/sdnq/dequantizer.py

@@ -4,7 +4,7 @@ import torch
 from modules import shared
 
 from .common import dtype_dict, use_torch_compile
-from .packed_int import pack_int_symetric, unpack_int_symetric, packed_int_function_dict
+from .packed_int import pack_int_symetric, unpack_int_symetric, pack_int_asymetric, unpack_int_asymetric
 
 
 def dequantize_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor:
@@ -15,10 +15,9 @@ def dequantize_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, ze
 
 
 def dequantize_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
+    result = weight.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype)
     if skip_quantized_matmul:
-        result = weight.transpose(0,1).to(dtype=scale.dtype).mul_(scale.transpose(0,1)).to(dtype=dtype)
-    else:
-        result = weight.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype)
+        result = result.t()
     if result_shape is not None:
         result = result.reshape(result_shape)
     return result
@@ -29,14 +28,11 @@ def dequantize_symmetric_with_bias(weight: torch.CharTensor, scale: torch.FloatT
 
 
 def dequantize_packed_int_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str) -> torch.FloatTensor:
-    return dequantize_asymmetric(packed_int_function_dict[weights_dtype]["unpack"](weight, shape), scale, zero_point, dtype, result_shape)
+    return dequantize_asymmetric(unpack_int_asymetric(weight, shape, weights_dtype), scale, zero_point, dtype, result_shape)
 
 
 def dequantize_packed_int_symmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
-    if skip_quantized_matmul:
-        return dequantize_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale.transpose(0,1), dtype, result_shape)
-    else:
-        return dequantize_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape)
+    return dequantize_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape, skip_quantized_matmul=skip_quantized_matmul)
 
 
 class AsymmetricWeightsDequantizer(torch.nn.Module):
@@ -115,7 +111,7 @@ class PackedINTAsymmetricWeightsDequantizer(torch.nn.Module):
         self.register_buffer("zero_point", zero_point)
 
     def pack_weight(self, weight: torch.Tensor) -> torch.Tensor:
-        return packed_int_function_dict[self.weights_dtype]["pack"](weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"]))
+        return pack_int_asymetric(weight, self.weights_dtype)
 
     def forward(self, weight, **kwargs): # pylint: disable=unused-argument
         return dequantize_packed_int_asymmetric_compiled(weight, self.scale, self.zero_point, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype)

modules/sdnq/layers/linear/linear_int8.py

@@ -28,7 +28,7 @@ def int8_matmul(
     weights_dtype: str,
 ) -> torch.FloatTensor:
     if quantized_weight_shape is not None:
-        weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8, transpose=True)
+        weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8)
     return_dtype = input.dtype
     output_shape = list(input.shape)
     output_shape[-1] = weight.shape[-1]

modules/sdnq/packed_int.py

@@ -11,13 +11,18 @@ def pack_int_symetric(tensor: torch.CharTensor, weights_dtype: str) -> torch.Byt
     return packed_int_function_dict[weights_dtype]["pack"](tensor.sub_(dtype_dict[weights_dtype]["min"]).to(dtype=dtype_dict[weights_dtype]["storage_dtype"]))
 
 
-def unpack_int_symetric(packed_tensor: torch.ByteTensor, shape: torch.Size, weights_dtype: str, dtype: Optional[torch.dtype] = None, transpose: Optional[bool] = False) -> torch.CharTensor:
+def pack_int_asymetric(tensor: torch.CharTensor, weights_dtype: str) -> torch.ByteTensor:
+    return packed_int_function_dict[weights_dtype]["pack"](tensor.to(dtype=dtype_dict[weights_dtype]["storage_dtype"]))
+
+
+def unpack_int_symetric(packed_tensor: torch.ByteTensor, shape: torch.Size, weights_dtype: str, dtype: Optional[torch.dtype] = None) -> torch.CharTensor:
     if dtype is None:
         dtype = dtype_dict[weights_dtype]["torch_dtype"]
-    result = packed_int_function_dict[weights_dtype]["unpack"](packed_tensor, shape).to(dtype=dtype).add_(dtype_dict[weights_dtype]["min"])
-    if transpose:
-        result = result.transpose(0,1)
-    return result
+    return packed_int_function_dict[weights_dtype]["unpack"](packed_tensor, shape).to(dtype=dtype).add_(dtype_dict[weights_dtype]["min"])
+
+
+def unpack_int_asymetric(packed_tensor: torch.ByteTensor, shape: torch.Size, weights_dtype: str) -> torch.CharTensor:
+    return packed_int_function_dict[weights_dtype]["unpack"](packed_tensor, shape)
 
 
 def pack_uint7(tensor: torch.ByteTensor) -> torch.ByteTensor: