diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py index 20b507ece..30d8a82e7 100644 --- a/modules/sdnq/common.py +++ b/modules/sdnq/common.py @@ -6,7 +6,7 @@ import torch from modules import shared -sdnq_version = "0.2.2" +sdnq_version = "0.2.3" sdnq_keys = {"weight", "scale", "zero_point", "svd_up", "svd_down"} torch_version = torch.__version__[:4] diff --git a/modules/sdnq/kernel_wrappers.py b/modules/sdnq/kernel_wrappers.py index 406ecba33..a27d2d994 100644 --- a/modules/sdnq/kernel_wrappers.py +++ b/modules/sdnq/kernel_wrappers.py @@ -47,7 +47,7 @@ if os.environ.get("SDNQ_USE_TRITON_MM", None) is None: use_triton_mm = bool(not is_alchemist_or_igpu and (devices.backend in {"cuda", "rocm", "ipex", "xpu", "zluda"})) else: use_triton_mm = bool(os.environ.get("SDNQ_USE_TRITON_MM", "0").lower() not in {"0", "false", "no"}) -use_triton_scaled_mm = use_triton_mm and os.environ.get("SDNQ_USE_TRITON_SCALED_MM", "1").lower() not in {"0", "false", "no"} +use_triton_scaled_mm = bool(use_triton_mm and os.environ.get("SDNQ_USE_TRITON_SCALED_MM", "1").lower() not in {"0", "false", "no"}) if os.environ.get("SDNQ_USE_TENSORWISE_FP8_MM", None) is None: # row-wise FP8 only exist on H100 hardware, sdnq will use software row-wise with tensorwise hardware with this setting @@ -65,20 +65,6 @@ else: use_contiguous_fp8_mm = use_contiguous_fp16_mm -def fp_mm_torch_cuda(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor: - return torch.mm(a,b, out_dtype=out_dtype) - -def fp_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor: - if b.dtype == torch.float8_e4m3fn: - fp16_scale = 4 * b.shape[-2] - else: - fp16_scale = 65536 * b.shape[-2] - in_scale = fp16_scale**0.5 - a = a.to(dtype=torch.float32).div_(in_scale).to(dtype=torch.float16) - b = b.to(dtype=torch.float32).div_(in_scale).to(dtype=torch.float16) - return torch.mm(a,b).to(dtype=torch.float32).mul_(fp16_scale).to(dtype=out_dtype) - - int_mm_func = None fp_mm_func = None fp8_mm_func = None @@ -130,6 +116,26 @@ if ( use_triton_scaled_mm = False +if os.environ.get("SDNQ_INCLUDE_MM_KERNEL_IN_COMPILE", None) is None: + include_mm_kernel_in_compile = bool(not use_triton_scaled_mm) +else: + include_mm_kernel_in_compile = bool(os.environ.get("SDNQ_INCLUDE_MM_KERNEL_IN_COMPILE", "0").lower() not in {"0", "false", "no"}) + + +def fp_mm_torch_cuda(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor: + return torch.mm(a,b, out_dtype=out_dtype) + +def fp_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor: + if b.dtype == torch.float8_e4m3fn: + fp16_scale = 4 * b.shape[-2] + else: + fp16_scale = 65536 * b.shape[-2] + in_scale = fp16_scale**0.5 + a = a.to(dtype=torch.float32).div_(in_scale).to(dtype=torch.float16) + b = b.to(dtype=torch.float32).div_(in_scale).to(dtype=torch.float16) + return torch.mm(a,b).to(dtype=torch.float32).mul_(fp16_scale).to(dtype=out_dtype) + + if int_mm_func is None: def int_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.int32) -> torch.FloatTensor: return torch._int_mm(a,b).to(dtype=out_dtype) diff --git a/modules/sdnq/layers/linear/linear_fp16.py b/modules/sdnq/layers/linear/linear_fp16.py index d4b40c2cf..07afeb27f 100644 --- a/modules/sdnq/layers/linear/linear_fp16.py +++ b/modules/sdnq/layers/linear/linear_fp16.py @@ -3,7 +3,7 @@ import torch from ...common import compile_func -from ...kernel_wrappers import fp_scaled_mm_func +from ...kernel_wrappers import fp_scaled_mm_func, include_mm_kernel_in_compile from ...quant_utils import rotate_hadamard, get_hadamard from ...packed_float import unpack_float @@ -44,6 +44,29 @@ def get_fp16_matmul_inputs( return input, weight, input_scale, scale, bias, return_dtype, output_shape +def fp16_matmul( + input: torch.FloatTensor, + weight: torch.Tensor, + scale: torch.FloatTensor, + bias: torch.FloatTensor | None = None, + svd_up: torch.FloatTensor | None = None, + svd_down: torch.FloatTensor | None = None, + hadamard: torch.FloatTensor | None = None, + quantized_weight_shape: torch.Size | None = None, + weights_dtype: str | None = None, +) -> torch.FloatTensor: + input, weight, input_scale, scale, bias, return_dtype, output_shape = get_fp16_matmul_inputs( + input, weight, scale, + bias=bias, + svd_up=svd_up, + svd_down=svd_down, + hadamard=hadamard, + quantized_weight_shape=quantized_weight_shape, + weights_dtype=weights_dtype, + ) + return fp_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(output_shape) + + def quantized_linear_forward_fp16_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor: if torch.numel(input) / input.shape[-1] < 32: return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, zero_point=self.zero_point, svd_up=self.svd_up, svd_down=self.svd_down, skip_quantized_matmul=True), self.bias) @@ -58,7 +81,7 @@ def quantized_linear_forward_fp16_matmul(self, input: torch.FloatTensor) -> torc else: hadamard = None - input, weight, input_scale, scale, bias, return_dtype, output_shape = get_fp16_matmul_inputs( + return fp16_matmul( input, weight, scale, bias=self.bias, svd_up=self.svd_up, @@ -67,7 +90,9 @@ def quantized_linear_forward_fp16_matmul(self, input: torch.FloatTensor) -> torc quantized_weight_shape=quantized_weight_shape, weights_dtype=self.sdnq_dequantizer.weights_dtype, ) - return fp_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(output_shape) -get_fp16_matmul_inputs = compile_func(get_fp16_matmul_inputs) +if not include_mm_kernel_in_compile: + get_fp16_matmul_inputs = compile_func(get_fp16_matmul_inputs) +else: + fp16_matmul = compile_func(fp16_matmul) diff --git a/modules/sdnq/layers/linear/linear_fp8.py b/modules/sdnq/layers/linear/linear_fp8.py index 132d0cbf1..71ce7e404 100644 --- a/modules/sdnq/layers/linear/linear_fp8.py +++ b/modules/sdnq/layers/linear/linear_fp8.py @@ -3,7 +3,7 @@ import torch from ...common import compile_func -from ...kernel_wrappers import fp8_scaled_mm_func +from ...kernel_wrappers import fp8_scaled_mm_func, is_fp8_mm_supported, include_mm_kernel_in_compile from ...quant_utils import quantize_fp_mm, rotate_hadamard, get_hadamard from ...packed_float import unpack_float @@ -51,6 +51,29 @@ def get_fp8_matmul_inputs( return input, weight, input_scale, scale, bias, return_dtype, output_shape +def fp8_matmul( + input: torch.FloatTensor, + weight: torch.Tensor, + scale: torch.FloatTensor, + bias: torch.FloatTensor | None = None, + svd_up: torch.FloatTensor | None = None, + svd_down: torch.FloatTensor | None = None, + hadamard: torch.FloatTensor | None = None, + quantized_weight_shape: torch.Size | None = None, + weights_dtype: str | None = None, +) -> torch.FloatTensor: + input, weight, input_scale, scale, bias, return_dtype, output_shape = get_fp8_matmul_inputs( + input, weight, scale, + bias=bias, + svd_up=svd_up, + svd_down=svd_down, + hadamard=hadamard, + quantized_weight_shape=quantized_weight_shape, + weights_dtype=weights_dtype, + ) + return fp8_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(output_shape) + + def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor: if torch.numel(input) / input.shape[-1] < 32: return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, zero_point=self.zero_point, svd_up=self.svd_up, svd_down=self.svd_down, skip_quantized_matmul=True), self.bias) @@ -65,7 +88,7 @@ def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch else: hadamard = None - input, weight, input_scale, scale, bias, return_dtype, output_shape = get_fp8_matmul_inputs( + return fp8_matmul( input, weight, scale, bias=self.bias, svd_up=self.svd_up, @@ -74,7 +97,9 @@ def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch quantized_weight_shape=quantized_weight_shape, weights_dtype=self.sdnq_dequantizer.weights_dtype, ) - return fp8_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(output_shape) -get_fp8_matmul_inputs = compile_func(get_fp8_matmul_inputs) +if is_fp8_mm_supported and not include_mm_kernel_in_compile: + get_fp8_matmul_inputs = compile_func(get_fp8_matmul_inputs) +else: + fp8_matmul = compile_func(fp8_matmul) diff --git a/modules/sdnq/layers/linear/linear_int8.py b/modules/sdnq/layers/linear/linear_int8.py index 435db05d2..4a2c0d3ff 100644 --- a/modules/sdnq/layers/linear/linear_int8.py +++ b/modules/sdnq/layers/linear/linear_int8.py @@ -3,7 +3,7 @@ import torch from ...common import compile_func -from ...kernel_wrappers import int_scaled_mm_func +from ...kernel_wrappers import int_scaled_mm_func, include_mm_kernel_in_compile from ...quant_utils import quantize_int_mm, rotate_hadamard, get_hadamard from ...packed_int import unpack_int @@ -68,6 +68,31 @@ def get_int8_matmul_inputs( return input, weight, input_scale, scale, bias, return_dtype, output_shape +def int8_matmul( + input: torch.FloatTensor, + weight: torch.Tensor, + scale: torch.FloatTensor, + bias: torch.FloatTensor | None = None, + svd_up: torch.FloatTensor | None = None, + svd_down: torch.FloatTensor | None = None, + zero_point: torch.FloatTensor | None = None, + hadamard: torch.FloatTensor | None = None, + quantized_weight_shape: torch.Size | None = None, + weights_dtype: str | None = None, +) -> torch.FloatTensor: + input, weight, input_scale, scale, bias, return_dtype, output_shape = get_int8_matmul_inputs( + input, weight, scale, + bias=bias, + svd_up=svd_up, + svd_down=svd_down, + zero_point=zero_point, + hadamard=hadamard, + quantized_weight_shape=quantized_weight_shape, + weights_dtype=weights_dtype, + ) + return int_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(output_shape) + + def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor: if torch.numel(input) / input.shape[-1] < 32: return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, zero_point=self.zero_point, svd_up=self.svd_up, svd_down=self.svd_down, skip_quantized_matmul=True), self.bias) @@ -83,7 +108,7 @@ def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torc else: hadamard = None - input, weight, input_scale, scale, bias, return_dtype, output_shape = get_int8_matmul_inputs( + return int8_matmul( input, weight, scale, bias=self.bias, svd_up=self.svd_up, @@ -93,7 +118,9 @@ def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torc quantized_weight_shape=quantized_weight_shape, weights_dtype=self.sdnq_dequantizer.weights_dtype, ) - return int_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(output_shape) -get_int8_matmul_inputs = compile_func(get_int8_matmul_inputs) +if not include_mm_kernel_in_compile: + get_int8_matmul_inputs = compile_func(get_int8_matmul_inputs) +else: + int8_matmul = compile_func(int8_matmul) diff --git a/modules/sdnq/layers/linear/linear_uint8.py b/modules/sdnq/layers/linear/linear_uint8.py index 9bc9f3f4d..12fb23756 100644 --- a/modules/sdnq/layers/linear/linear_uint8.py +++ b/modules/sdnq/layers/linear/linear_uint8.py @@ -3,7 +3,7 @@ import torch from ...common import compile_func -from ...kernel_wrappers import int_scaled_mm_func +from ...kernel_wrappers import int_scaled_mm_func, include_mm_kernel_in_compile from ...quant_utils import quantize_uint_mm, rotate_hadamard, get_hadamard from ...packed_int import unpack_int @@ -73,6 +73,31 @@ def get_uint8_matmul_inputs( return input, weight, input_scale, scale, zero_bias, return_dtype, output_shape +def uint8_matmul( + input: torch.FloatTensor, + weight: torch.Tensor, + scale: torch.FloatTensor, + zero_point: torch.FloatTensor, + bias: torch.FloatTensor | None = None, + svd_up: torch.FloatTensor | None = None, + svd_down: torch.FloatTensor | None = None, + hadamard: torch.FloatTensor | None = None, + quantized_weight_shape: torch.Size | None = None, + weights_dtype: str | None = None, +) -> torch.FloatTensor: + input, weight, input_scale, scale, zero_bias, return_dtype, output_shape = get_uint8_matmul_inputs( + input, weight, + scale, zero_point, + bias=bias, + svd_up=svd_up, + svd_down=svd_down, + hadamard=hadamard, + quantized_weight_shape=quantized_weight_shape, + weights_dtype=weights_dtype, + ) + return int_scaled_mm_func(input, weight, input_scale, scale, bias=zero_bias, out_dtype=return_dtype).view(output_shape) + + def quantized_linear_forward_uint8_matmul(self, input: torch.FloatTensor) -> torch.FloatTensor: if torch.numel(input) / input.shape[-1] < 32: return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight, self.scale, zero_point=self.zero_point, svd_up=self.svd_up, svd_down=self.svd_down, skip_quantized_matmul=True), self.bias) @@ -87,7 +112,7 @@ def quantized_linear_forward_uint8_matmul(self, input: torch.FloatTensor) -> tor else: hadamard = None - input, weight, input_scale, scale, zero_bias, return_dtype, output_shape = get_uint8_matmul_inputs( + return uint8_matmul( input, weight, scale, zero_point, bias=self.bias, @@ -97,7 +122,9 @@ def quantized_linear_forward_uint8_matmul(self, input: torch.FloatTensor) -> tor quantized_weight_shape=quantized_weight_shape, weights_dtype=self.sdnq_dequantizer.weights_dtype, ) - return int_scaled_mm_func(input, weight, input_scale, scale, bias=zero_bias, out_dtype=return_dtype).view(output_shape) -get_uint8_matmul_inputs = compile_func(get_uint8_matmul_inputs) +if not include_mm_kernel_in_compile: + get_uint8_matmul_inputs = compile_func(get_uint8_matmul_inputs) +else: + uint8_matmul = compile_func(uint8_matmul)