SDNQ FP8 matmul support for Conv2d

modules/model_quant_sdnq.py

@@ -29,6 +29,7 @@ dtype_dict = {
     "float8_e5m2fnuz": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e5m2fnuz, "storage_dtype": torch.float8_e5m2fnuz, "is_unsigned": False, "is_integer": False},
 }
 
+use_tensorwise_fp8_matmul = torch_version < 2.5 or devices.backend in {"cpu", "openvino"} or (devices.backend == "cuda" and sys.platform == "win32" and torch_version <= 2.7 and torch.cuda.get_device_capability(devices.device) == (8,9))
 quantized_matmul_dtypes = ("int8", "int6", "int4", "int2", "float8_e4m3fn", "float8_e5m2")
 if devices.backend in {"cpu", "openvino"}:
     quantized_matmul_dtypes += ("float8_e4m3fnuz", "float8_e5m2fnuz")
@@ -49,7 +50,6 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
         is_conv_type = False
         is_conv_transpose_type = False
         is_linear_type = False
-        use_tensorwise_fp8_matmul = False
         result_shape = None
         if torch_dtype is None:
             torch_dtype = devices.dtype
@@ -65,10 +65,9 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
             group_channel_size = channel_size // layer.groups
             use_quantized_matmul = False
             if shared.opts.sdnq_use_quantized_matmul_conv:
-                use_quantized_matmul = dtype_dict[weights_dtype]["is_integer"] and weights_dtype in quantized_matmul_dtypes and group_channel_size >= 32 and output_channel_size >= 32
-                #if use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"]:
-                #    use_quantized_matmul = output_channel_size % 16 == 0 and group_channel_size % 16 == 0
-                #    use_tensorwise_fp8_matmul = torch_version < 2.5 or devices.backend in {"cpu", "openvino"} or (devices.backend == "cuda" and sys.platform == "win32" and torch_version <= 2.7 and torch.cuda.get_device_capability(devices.device) == (8,9))
+                use_quantized_matmul = weights_dtype in quantized_matmul_dtypes and group_channel_size >= 32 and output_channel_size >= 32
+                if use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"]:
+                    use_quantized_matmul = output_channel_size % 16 == 0 and group_channel_size % 16 == 0
                 if use_quantized_matmul:
                     result_shape = layer.weight.shape
                     layer.weight.data = layer.weight.reshape(output_channel_size, -1)
@@ -89,7 +88,6 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
                 use_quantized_matmul = weights_dtype in quantized_matmul_dtypes and channel_size >= 32 and output_channel_size >= 32
                 if use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"]:
                     use_quantized_matmul = output_channel_size % 16 == 0 and channel_size % 16 == 0
-                    use_tensorwise_fp8_matmul = torch_version < 2.5 or devices.backend in {"cpu", "openvino"} or (devices.backend == "cuda" and sys.platform == "win32" and torch_version <= 2.7 and torch.cuda.get_device_capability(devices.device) == (8,9))
 
         if group_size == 0:
             if is_linear_type:
@@ -203,7 +201,13 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz
                 layer.forward = quantized_linear_forward
         elif is_conv_type:
             if use_quantized_matmul:
-                layer.forward = quantized_conv2d_forward_int8_matmul
+                if dtype_dict[weights_dtype]["is_integer"]:
+                    layer.forward = quantized_conv2d_forward_int8_matmul
+                else:
+                    if use_tensorwise_fp8_matmul:
+                        layer.forward = quantized_conv2d_forward_fp8_matmul_tensorwise
+                    else:
+                        layer.forward = quantized_conv2d_forward_fp8_matmul
             else:
                 layer.forward = quantized_conv_forward
         elif is_conv_transpose_type:
@@ -466,6 +470,89 @@ def int8_matmul(
     return result
 
 
+def conv2d_fp8_matmul(
+    input: torch.FloatTensor,
+    weight: torch.ByteTensor,
+    bias: torch.FloatTensor,
+    scale: torch.FloatTensor,
+    result_shape: torch.Size,
+    weights_dtype: str,
+    reversed_padding_repeated_twice: List[int],
+    padding_mode: str, groups: int,
+    stride_h: int, stride_w: int,
+    padding_h: int, padding_w: int,
+    dilation_h: int, dilation_w: int,
+) -> torch.FloatTensor:
+    return_dtype = input.dtype
+    mm_output_shape, K_h, K_w = get_conv2d_shapes(input.shape, result_shape, stride_h, stride_w, padding_h, padding_w, dilation_h, dilation_w)
+    if padding_mode != "zeros":
+        input = torch.nn.functional.pad(input, reversed_padding_repeated_twice, mode=padding_mode)
+        padding_h = padding_w = 0
+
+    input, input_scale = quantize_fp8_matmul_input(
+        torch.nn.functional.unfold(
+            input, kernel_size=(K_h, K_w), padding=(padding_h, padding_w), stride=(stride_h, stride_w), dilation=(dilation_h, dilation_w)
+        ).transpose(1,2),
+    )
+
+    if groups == 1:
+        result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).reshape(mm_output_shape).permute(0,3,1,2)
+    else:
+        scale = scale.reshape(groups, 1, scale.shape[1] // groups)
+        input_scale = input_scale.reshape(groups, input_scale.shape[0] // groups, 1)
+        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)
+        result = result.permute(0,3,1,2)
+    return result
+
+
+def conv2d_fp8_matmul_tensorwise(
+    input: torch.FloatTensor,
+    weight: torch.ByteTensor,
+    bias: torch.FloatTensor,
+    scale: torch.FloatTensor,
+    result_shape: torch.Size,
+    weights_dtype: str,
+    reversed_padding_repeated_twice: List[int],
+    padding_mode: str, groups: int,
+    stride_h: int, stride_w: int,
+    padding_h: int, padding_w: int,
+    dilation_h: int, dilation_w: int,
+) -> torch.FloatTensor:
+    return_dtype = input.dtype
+    mm_output_shape, K_h, K_w = get_conv2d_shapes(input.shape, result_shape, stride_h, stride_w, padding_h, padding_w, dilation_h, dilation_w)
+    if padding_mode != "zeros":
+        input = torch.nn.functional.pad(input, reversed_padding_repeated_twice, mode=padding_mode)
+        padding_h = padding_w = 0
+
+    input, scale = quantize_fp8_matmul_input_tensorwise(
+        torch.nn.functional.unfold(
+            input, kernel_size=(K_h, K_w), padding=(padding_h, padding_w), stride=(stride_h, stride_w), dilation=(dilation_h, dilation_w)
+        ).transpose(1,2),
+        scale,
+    )
+
+    dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
+    if groups == 1:
+        result = decompress_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)
+    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 = decompress_symmetric(torch.cat(result, dim=-1), scale, return_dtype, mm_output_shape)
+    if bias is not None:
+        result.add_(bias)
+    return result.permute(0,3,1,2)
+
+
 def conv2d_int8_matmul(
     input: torch.FloatTensor,
     weight: torch.ByteTensor,
@@ -481,17 +568,12 @@ def conv2d_int8_matmul(
     dilation_h: int, dilation_w: int,
 ) -> torch.FloatTensor:
     return_dtype = input.dtype
-    batch_size, _, H_in, W_in = input.shape
-    C_out, _, K_h, K_w = result_shape
-    W_out = (W_in + 2 * padding_w - dilation_w * (K_w - 1) - 1) // stride_w + 1
-    H_out = (H_in + 2 * padding_h - dilation_h * (K_h - 1) - 1) // stride_h + 1
-    mm_output_shape =  (batch_size, H_out, W_out, C_out)
-
-    if compressed_weight_shape is not None:
-        weight = unpack_int_symetric(weight, compressed_weight_shape, weights_dtype, dtype=torch.int8, transpose=True)
+    mm_output_shape, K_h, K_w = get_conv2d_shapes(input.shape, result_shape, stride_h, stride_w, padding_h, padding_w, dilation_h, dilation_w)
     if padding_mode != "zeros":
         input = torch.nn.functional.pad(input, reversed_padding_repeated_twice, mode=padding_mode)
         padding_h = padding_w = 0
+    if compressed_weight_shape is not None:
+        weight = unpack_int_symetric(weight, compressed_weight_shape, weights_dtype, dtype=torch.int8, transpose=True)
 
     input, scale = quantize_int8_matmul_input(
         torch.nn.functional.unfold(
@@ -532,22 +614,66 @@ def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torc
     return int8_matmul(input, self.weight, self.bias, self.sdnq_decompressor.scale, getattr(self.sdnq_decompressor, "compressed_weight_shape", None), self.sdnq_decompressor.weights_dtype)
 
 
+def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor:
+    return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight), self.bias)
+
+
+def get_conv2d_args(stride, padding, dilation):
+    if isinstance(stride, int):
+        stride_h = stride_w = stride
+    else:
+        stride_h, stride_w = stride
+    if isinstance(padding, int):
+        padding_h = padding_w = padding
+    else:
+        padding_h, padding_w = padding
+    if isinstance(dilation, int):
+        dilation_h = dilation_w = dilation
+    else:
+        dilation_h, dilation_w = dilation
+    return stride_h, stride_w, padding_h, padding_w, dilation_h, dilation_w
+
+
+def get_conv2d_shapes(input_shape, result_shape, stride_h, stride_w, padding_h, padding_w, dilation_h, dilation_w):
+    batch_size, _, H_in, W_in = input_shape
+    C_out, _, K_h, K_w = result_shape
+    W_out = (W_in + 2 * padding_w - dilation_w * (K_w - 1) - 1) // stride_w + 1
+    H_out = (H_in + 2 * padding_h - dilation_h * (K_h - 1) - 1) // stride_h + 1
+    return (batch_size, H_out, W_out, C_out), K_h, K_w
+
+
+def quantized_conv2d_forward_fp8_matmul(self, input) -> torch.FloatTensor:
+    stride_h, stride_w, padding_h, padding_w, dilation_h, dilation_w = get_conv2d_args(self.stride, self.padding, self.dilation)
+    return conv2d_fp8_matmul(
+        input, self.weight, self.bias,
+        self.sdnq_decompressor.scale,
+        self.sdnq_decompressor.result_shape,
+        self.sdnq_decompressor.weights_dtype,
+        self._reversed_padding_repeated_twice,
+        self.padding_mode, self.groups,
+        stride_h, stride_w,
+        padding_h, padding_w,
+        dilation_h, dilation_w,
+    )
+
+
+def quantized_conv2d_forward_fp8_matmul_tensorwise(self, input) -> torch.FloatTensor:
+    stride_h, stride_w, padding_h, padding_w, dilation_h, dilation_w = get_conv2d_args(self.stride, self.padding, self.dilation)
+    return conv2d_fp8_matmul_tensorwise(
+        input, self.weight, self.bias,
+        self.sdnq_decompressor.scale,
+        self.sdnq_decompressor.result_shape,
+        self.sdnq_decompressor.weights_dtype,
+        self._reversed_padding_repeated_twice,
+        self.padding_mode, self.groups,
+        stride_h, stride_w,
+        padding_h, padding_w,
+        dilation_h, dilation_w,
+    )
+
+
 def quantized_conv2d_forward_int8_matmul(self, input) -> torch.FloatTensor:
-    if isinstance(self.stride, int):
-        stride_h = stride_w = self.stride
-    else:
-        stride_h, stride_w = self.stride
-
-    if isinstance(self.padding, int):
-        padding_h = padding_w = self.padding
-    else:
-        padding_h, padding_w = self.padding
-
-    if isinstance(self.dilation, int):
-        dilation_h = dilation_w = self.dilation
-    else:
-        dilation_h, dilation_w = self.dilation
-
+    stride_h, stride_w, padding_h, padding_w, dilation_h, dilation_w = get_conv2d_args(self.stride, self.padding, self.dilation)
     return conv2d_int8_matmul(
         input, self.weight, self.bias,
         self.sdnq_decompressor.scale,
@@ -562,10 +688,6 @@ def quantized_conv2d_forward_int8_matmul(self, input) -> torch.FloatTensor:
     )
 
 
-def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor:
-    return torch.nn.functional.linear(input, self.sdnq_decompressor(self.weight), self.bias)
-
-
 def quantized_conv_forward(self, input) -> torch.FloatTensor:
     return self._conv_forward(input, self.sdnq_decompressor(self.weight), self.bias)
 
@@ -927,10 +1049,12 @@ if shared.opts.sdnq_decompress_compile:
         decompress_symmetric_compiled = torch.compile(decompress_symmetric, fullgraph=True)
         decompress_packed_int_asymmetric_compiled = torch.compile(decompress_packed_int_asymmetric, fullgraph=True)
         decompress_packed_int_symmetric_compiled = torch.compile(decompress_packed_int_symmetric, fullgraph=True)
+        int8_matmul = torch.compile(int8_matmul, fullgraph=True)
         fp8_matmul = torch.compile(fp8_matmul, fullgraph=True)
         fp8_matmul_tensorwise = torch.compile(fp8_matmul_tensorwise, fullgraph=True)
-        int8_matmul = torch.compile(int8_matmul, fullgraph=True)
         conv2d_int8_matmul = torch.compile(conv2d_int8_matmul, fullgraph=True)
+        conv2d_fp8_matmul = torch.compile(conv2d_fp8_matmul, fullgraph=True)
+        conv2d_fp8_matmul_tensorwise = torch.compile(conv2d_fp8_matmul_tensorwise, fullgraph=True)
     except Exception as e:
         shared.log.warning(f"Quantization: type=sdnq Decompress using torch.compile is not available: {e}")
         decompress_asymmetric_compiled = decompress_asymmetric