From d31df8c1eb37c4e02e5433f594c403fbe800c4ce Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sat, 14 Jun 2025 22:10:10 +0300 Subject: [PATCH] SDNQ fuse bias into dequantizer with matmul --- modules/sdnq/dequantizer.py | 29 ++++++++++++++++++-------- modules/sdnq/forward.py | 41 ++++++++++++++++++++++--------------- 2 files changed, 45 insertions(+), 25 deletions(-) diff --git a/modules/sdnq/dequantizer.py b/modules/sdnq/dequantizer.py index 1f2b36df1..be59019a3 100644 --- a/modules/sdnq/dequantizer.py +++ b/modules/sdnq/dequantizer.py @@ -24,6 +24,10 @@ def dequantize_symmetric(input: torch.CharTensor, scale: torch.FloatTensor, dtyp return result +def dequantize_symmetric_with_bias(input: torch.CharTensor, bias: torch.FloatTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor: + return torch.addcmul(bias, input.to(dtype=scale.dtype), scale).to(dtype=dtype).reshape(result_shape) + + def dequantize_packed_int_asymmetric(input: 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"](input, shape), scale, zero_point, dtype, result_shape) @@ -57,7 +61,7 @@ class AsymmetricWeightsDequantizer(torch.nn.Module): return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"]) def forward(self, weight, **kwargs): # pylint: disable=unused-argument - return dequantize_asymmetric(weight, self.scale, self.zero_point, self.result_dtype, self.result_shape) + return dequantize_asymmetric_compiled(weight, self.scale, self.zero_point, self.result_dtype, self.result_shape) class SymmetricWeightsDequantizer(torch.nn.Module): @@ -81,7 +85,7 @@ class SymmetricWeightsDequantizer(torch.nn.Module): return weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"]) def forward(self, weight, skip_quantized_matmul=False, **kwargs): # pylint: disable=unused-argument - return dequantize_symmetric(weight, self.scale, self.result_dtype, self.result_shape, skip_quantized_matmul=skip_quantized_matmul) + return dequantize_symmetric_compiled(weight, self.scale, self.result_dtype, self.result_shape, skip_quantized_matmul=skip_quantized_matmul) class PackedINTAsymmetricWeightsDequantizer(torch.nn.Module): @@ -108,7 +112,7 @@ class PackedINTAsymmetricWeightsDequantizer(torch.nn.Module): return packed_int_function_dict[self.weights_dtype]["pack"](weight.to(dtype=dtype_dict[self.weights_dtype]["torch_dtype"])) def forward(self, weight, **kwargs): # pylint: disable=unused-argument - return dequantize_packed_int_asymmetric(weight, self.scale, self.zero_point, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype) + 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) class PackedINTSymmetricWeightsDequantizer(torch.nn.Module): @@ -134,7 +138,7 @@ class PackedINTSymmetricWeightsDequantizer(torch.nn.Module): return pack_int_symetric(weight, self.weights_dtype) def forward(self, weight, skip_quantized_matmul=False, **kwargs): # pylint: disable=unused-argument - return dequantize_packed_int_symmetric(weight, self.scale, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype, skip_quantized_matmul=skip_quantized_matmul) + return dequantize_packed_int_symmetric_compiled(weight, self.scale, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype, skip_quantized_matmul=skip_quantized_matmul) dequantizer_dict = { @@ -164,9 +168,18 @@ dequantizer_dict = { if shared.opts.sdnq_dequantize_compile: try: torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit) - dequantize_asymmetric = torch.compile(dequantize_asymmetric, fullgraph=True) - dequantize_symmetric = torch.compile(dequantize_symmetric, fullgraph=True) - dequantize_packed_int_asymmetric = torch.compile(dequantize_packed_int_asymmetric, fullgraph=True) - dequantize_packed_int_symmetric = torch.compile(dequantize_packed_int_symmetric, fullgraph=True) + dequantize_asymmetric_compiled = torch.compile(dequantize_asymmetric, fullgraph=True) + dequantize_symmetric_compiled = torch.compile(dequantize_symmetric, fullgraph=True) + dequantize_packed_int_asymmetric_compiled = torch.compile(dequantize_packed_int_asymmetric, fullgraph=True) + dequantize_packed_int_symmetric_compiled = torch.compile(dequantize_packed_int_symmetric, fullgraph=True) except Exception as e: shared.log.warning(f"Quantization: type=sdnq Dequantize using torch.compile is not available: {e}") + dequantize_asymmetric_compiled = dequantize_asymmetric + dequantize_symmetric_compiled = dequantize_symmetric + dequantize_packed_int_asymmetric_compiled = dequantize_packed_int_asymmetric + dequantize_packed_int_symmetric_compiled = dequantize_packed_int_symmetric +else: + dequantize_asymmetric_compiled = dequantize_asymmetric + dequantize_symmetric_compiled = dequantize_symmetric + dequantize_packed_int_asymmetric_compiled = dequantize_packed_int_asymmetric + dequantize_packed_int_symmetric_compiled = dequantize_packed_int_symmetric diff --git a/modules/sdnq/forward.py b/modules/sdnq/forward.py index be4a957ea..9caa12f7d 100644 --- a/modules/sdnq/forward.py +++ b/modules/sdnq/forward.py @@ -5,7 +5,7 @@ import torch from modules import shared from .common import conv_types, conv_transpose_types -from .dequantizer import dequantize_symmetric +from .dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias from .packed_int import unpack_int_symetric @@ -94,10 +94,10 @@ def fp8_matmul_tensorwise( output_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) - result = dequantize_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, return_dtype, output_shape) if bias is not None: - result.add_(bias) - return result + 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), bias, scale, return_dtype, output_shape) + else: + return dequantize_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, return_dtype, output_shape) def int8_matmul( @@ -114,10 +114,10 @@ def int8_matmul( output_shape = list(input.shape) output_shape[-1] = weight.shape[-1] input, scale = quantize_int8_matmul_input(input, scale) - result = dequantize_symmetric(torch._int_mm(input, weight), scale, return_dtype, output_shape) if bias is not None: - result.add_(bias) - return result + return dequantize_symmetric_with_bias(torch._int_mm(input, weight), bias, scale, return_dtype, output_shape) + else: + return dequantize_symmetric(torch._int_mm(input, weight), scale, return_dtype, output_shape) def process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation): @@ -192,11 +192,14 @@ def conv_fp8_matmul( weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) result = [] - for i in range(groups): - result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=None, out_dtype=return_dtype)) - result = torch.cat(result, dim=-1).reshape(mm_output_shape) if bias is not None: - result.add_(bias) + bias = bias.reshape(groups, bias.shape[0] // groups) + for i in range(groups): + result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=bias[i], out_dtype=return_dtype)) + else: + for i in range(groups): + result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=None, out_dtype=return_dtype)) + result = torch.cat(result, dim=-1).reshape(mm_output_shape) if conv_type == 1: result = result.transpose(1,2) @@ -224,16 +227,18 @@ def conv_fp8_matmul_tensorwise( dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32) if groups == 1: - result = dequantize_symmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, return_dtype, mm_output_shape) + result = torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype) else: weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) result = [] for i in range(groups): result.append(torch._scaled_mm(input[i], weight[i], scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype)) - result = dequantize_symmetric(torch.cat(result, dim=-1), scale, return_dtype, mm_output_shape) + result = torch.cat(result, dim=-1) if bias is not None: - result.add_(bias) + dequantize_symmetric_with_bias(result, bias, scale, return_dtype, mm_output_shape) + else: + dequantize_symmetric(result, scale, return_dtype, mm_output_shape) if conv_type == 1: result = result.transpose(1,2) @@ -264,16 +269,18 @@ def conv_int8_matmul( weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8, transpose=True) if groups == 1: - result = dequantize_symmetric(torch._int_mm(input, weight), scale, return_dtype, mm_output_shape) + result = torch._int_mm(input, weight) else: weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1) input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1) result = [] for i in range(groups): result.append(torch._int_mm(input[i], weight[i])) - result = dequantize_symmetric(torch.cat(result, dim=-1), scale, return_dtype, mm_output_shape) + result = torch.cat(result, dim=-1) if bias is not None: - result.add_(bias) + result = dequantize_symmetric_with_bias(result, bias, scale, return_dtype, mm_output_shape) + else: + result = dequantize_symmetric(result, scale, return_dtype, mm_output_shape) if conv_type == 1: result = result.transpose(1,2)