SDNQ fuse bias into dequantizer with matmul

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

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)