SDNQ add check_mats to matmul

modules/rocm.py

@@ -248,8 +248,8 @@ if sys.platform == "win32":
             import torch
             import numpy as np
 
-            cholesky_ex_gpu = torch.linalg.cholesky_ex
-            @wraps(cholesky_ex_gpu)
+            original_cholesky_ex = torch.linalg.cholesky_ex
+            @wraps(original_cholesky_ex)
             def cholesky_ex(A: torch.Tensor, upper=False, check_errors=False, out=None) -> torch.return_types.linalg_cholesky_ex:
                 assert not check_errors
                 return_device = A.device
@@ -261,8 +261,8 @@ if sys.platform == "win32":
                 return torch.return_types.linalg_cholesky_ex((L, info), {})
             torch.linalg.cholesky_ex = cholesky_ex
 
-            cholesky_gpu = torch.linalg.cholesky
-            @wraps(cholesky_gpu)
+            original_cholesky = torch.linalg.cholesky
+            @wraps(original_cholesky)
             def cholesky(A: torch.Tensor, upper=False, out=None) -> torch.Tensor:
                 return_device = A.device
                 L = torch.from_numpy(np.linalg.cholesky(A.to("cpu").numpy(), upper=upper)).to(return_device)

modules/sdnq/__init__.py

@@ -172,12 +172,14 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
         if use_quantized_matmul and not re_quantize_for_matmul:
             scale.transpose_(0,1)
             layer.weight.transpose_(0,1)
-            if not dtype_dict[weights_dtype]["is_integer"]:
-                weight_stride = layer.weight.stride()
-                if not (weight_stride[0] == 1 and weight_stride[1] > 1):
-                    layer.weight.data = layer.weight.t().contiguous().t()
-                if not use_tensorwise_fp8_matmul:
-                    scale = scale.to(torch.float32)
+            weight_stride = layer.weight.stride()
+            if not (weight_stride[0] == 1 and weight_stride[1] > 1):
+                if devices.backend != "ipex":
+                    layer.weight.data = layer.weight.t_().contiguous().t_()
+            elif devices.backend == "ipex":
+                layer.weight.data = layer.weight.t_().contiguous().t_()
+            if not use_tensorwise_fp8_matmul and not dtype_dict[weights_dtype]["is_integer"]:
+                scale = scale.to(torch.float32)
 
         layer.sdnq_dequantizer = dequantizer_dict[weights_dtype](
             scale=scale,

modules/sdnq/layers/conv/conv_fp8.py

@@ -6,6 +6,7 @@ import torch
 
 from ...common import compile_func # noqa: TID252
 from ..linear.linear_fp8 import quantize_fp8_matmul_input # noqa: TID252
+from ..linear.forward import check_mats # noqa: TID252
 from .forward import get_conv_args, process_conv_input
 
 
@@ -23,6 +24,7 @@ def conv_fp8_matmul(
     return_dtype = input.dtype
     input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
     input, input_scale = quantize_fp8_matmul_input(input)
+    input, weight = check_mats(input, weight)
 
     if groups == 1:
         if bias is not None and bias.dtype != torch.bfloat16:

modules/sdnq/layers/conv/conv_fp8_tensorwise.py

@@ -7,6 +7,7 @@ import torch
 from ...common import compile_func # noqa: TID252
 from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252
 from ..linear.linear_fp8_tensorwise import quantize_fp8_matmul_input_tensorwise # noqa: TID252
+from ..linear.forward import check_mats # noqa: TID252
 from .forward import get_conv_args, process_conv_input
 
 
@@ -24,6 +25,7 @@ def conv_fp8_matmul_tensorwise(
     return_dtype = input.dtype
     input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
     input, scale = quantize_fp8_matmul_input_tensorwise(input, scale)
+    input, weight = check_mats(input, weight)
     dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
 
     if groups == 1:

modules/sdnq/layers/conv/conv_int8.py

@@ -8,6 +8,7 @@ from ...common import compile_func # noqa: TID252
 from ...packed_int import unpack_int_symetric # noqa: TID252
 from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252
 from ..linear.linear_int8 import quantize_int8_matmul_input # noqa: TID252
+from ..linear.forward import check_mats # noqa: TID252
 from .forward import get_conv_args, process_conv_input
 
 
@@ -29,6 +30,7 @@ def conv_int8_matmul(
     input, scale = quantize_int8_matmul_input(input, scale)
     if quantized_weight_shape is not None:
         weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8)
+    input, weight = check_mats(input, weight)
 
     if groups == 1:
         result = torch._int_mm(input, weight)

modules/sdnq/layers/linear/forward.py

@@ -1,7 +1,22 @@
 # pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
 
+from typing import Tuple
+
 import torch
 
 
+def check_mats(input: torch.Tensor, weight: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    input_stride = input.stride()
+    if not (input_stride[0] > input_stride[1] and input_stride[1] == 1):
+        input = input.contiguous()
+    weight_stride = weight.stride()
+    if not (weight_stride[0] == 1 and weight_stride[1] > 1):
+        if weight.device.type != "xpu":
+            weight = weight.t().contiguous().t()
+    elif weight.device.type == "xpu":
+        weight = weight.t().contiguous().t()
+    return input, weight
+
+
 def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor:
     return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight), self.bias)

modules/sdnq/layers/linear/linear_fp8.py

@@ -6,6 +6,7 @@ import torch
 
 from ...common import compile_func # noqa: TID252
 from ...dequantizer import quantize_fp8 # noqa: TID252
+from .forward import check_mats
 
 
 def quantize_fp8_matmul_input(input: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]:
@@ -23,6 +24,7 @@ def fp8_matmul(
     return_dtype = input.dtype
     output_shape = (*input.shape[:-1], weight.shape[-1])
     input, input_scale = quantize_fp8_matmul_input(input)
+    input, weight = check_mats(input, weight)
     if bias is not None and bias.dtype != torch.bfloat16:
         bias = bias.to(dtype=torch.bfloat16)
     return torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=torch.bfloat16).view(output_shape).to(return_dtype)

modules/sdnq/layers/linear/linear_fp8_tensorwise.py

@@ -6,6 +6,7 @@ import torch
 
 from ...common import compile_func # noqa: TID252
 from ...dequantizer import quantize_fp8, dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252
+from .forward import check_mats
 
 
 def quantize_fp8_matmul_input_tensorwise(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]:
@@ -27,6 +28,7 @@ def fp8_matmul_tensorwise(
     output_shape = (*input.shape[:-1], weight.shape[-1])
     dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
     input, scale = quantize_fp8_matmul_input_tensorwise(input, scale)
+    input, weight = check_mats(input, weight)
     if bias is not None:
         return dequantize_symmetric_with_bias(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, bias, return_dtype, output_shape)
     else:

modules/sdnq/layers/linear/linear_int8.py

@@ -7,6 +7,7 @@ import torch
 from ...common import compile_func # noqa: TID252
 from ...packed_int import unpack_int_symetric # noqa: TID252
 from ...dequantizer import quantize_int8, dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252
+from .forward import check_mats
 
 
 def quantize_int8_matmul_input(input: torch.FloatTensor, scale: torch.FloatTensor) -> Tuple[torch.CharTensor, torch.FloatTensor]:
@@ -31,6 +32,7 @@ def int8_matmul(
     return_dtype = input.dtype
     output_shape = (*input.shape[:-1], weight.shape[-1])
     input, scale = quantize_int8_matmul_input(input, scale)
+    input, weight = check_mats(input, weight)
     if bias is not None:
         return dequantize_symmetric_with_bias(torch._int_mm(input, weight), scale, bias, return_dtype, output_shape)
     else: