From 6b67a9d0c45f76034aeb4736022df885537cef8e Mon Sep 17 00:00:00 2001 From: Disty0 Date: Tue, 30 Sep 2025 01:58:13 +0300 Subject: [PATCH] SDNQ add check_mats to matmul --- modules/rocm.py | 8 ++++---- modules/sdnq/__init__.py | 14 ++++++++------ modules/sdnq/layers/conv/conv_fp8.py | 2 ++ modules/sdnq/layers/conv/conv_fp8_tensorwise.py | 2 ++ modules/sdnq/layers/conv/conv_int8.py | 2 ++ modules/sdnq/layers/linear/forward.py | 15 +++++++++++++++ modules/sdnq/layers/linear/linear_fp8.py | 2 ++ .../sdnq/layers/linear/linear_fp8_tensorwise.py | 2 ++ modules/sdnq/layers/linear/linear_int8.py | 2 ++ 9 files changed, 39 insertions(+), 10 deletions(-) diff --git a/modules/rocm.py b/modules/rocm.py index 000eb039f..1b03bfd32 100644 --- a/modules/rocm.py +++ b/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) diff --git a/modules/sdnq/__init__.py b/modules/sdnq/__init__.py index 82689efca..4f35009ec 100644 --- a/modules/sdnq/__init__.py +++ b/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, diff --git a/modules/sdnq/layers/conv/conv_fp8.py b/modules/sdnq/layers/conv/conv_fp8.py index 2d1407a47..7a83bb3a8 100644 --- a/modules/sdnq/layers/conv/conv_fp8.py +++ b/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: diff --git a/modules/sdnq/layers/conv/conv_fp8_tensorwise.py b/modules/sdnq/layers/conv/conv_fp8_tensorwise.py index bc968e519..e3b005622 100644 --- a/modules/sdnq/layers/conv/conv_fp8_tensorwise.py +++ b/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: diff --git a/modules/sdnq/layers/conv/conv_int8.py b/modules/sdnq/layers/conv/conv_int8.py index 6c59f6d60..260acd170 100644 --- a/modules/sdnq/layers/conv/conv_int8.py +++ b/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) diff --git a/modules/sdnq/layers/linear/forward.py b/modules/sdnq/layers/linear/forward.py index 20224204c..eff9b0428 100644 --- a/modules/sdnq/layers/linear/forward.py +++ b/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) diff --git a/modules/sdnq/layers/linear/linear_fp8.py b/modules/sdnq/layers/linear/linear_fp8.py index 63eee40b4..d1845b1e7 100644 --- a/modules/sdnq/layers/linear/linear_fp8.py +++ b/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) diff --git a/modules/sdnq/layers/linear/linear_fp8_tensorwise.py b/modules/sdnq/layers/linear/linear_fp8_tensorwise.py index 0343087c5..a35eb0721 100644 --- a/modules/sdnq/layers/linear/linear_fp8_tensorwise.py +++ b/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: diff --git a/modules/sdnq/layers/linear/linear_int8.py b/modules/sdnq/layers/linear/linear_int8.py index bcc4d775c..356179b97 100644 --- a/modules/sdnq/layers/linear/linear_int8.py +++ b/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: