SDNQ add fused Triton kernels

modules/sdnq/common.py

@@ -1,11 +1,10 @@
 # pylint: disable=redefined-builtin,no-member,protected-access
 
 import os
-import sys
 import json
 import torch
 
-from modules import shared, devices
+from modules import shared
 
 sdnq_version = "0.2.2"
 sdnq_keys = {"weight", "scale", "zero_point", "svd_up", "svd_down"}
@@ -336,106 +335,12 @@ weights_dtype_order = [
     "uint16", "float16_e1m15fnu", "float16_e2m14fnu", "float16_e3m13fnu", "float16_e4m12fnu", "float16_e5m11fnu",
 ]
 
+
 use_torch_compile = shared.opts.sdnq_dequantize_compile # this setting requires a full restart of the webui to apply
 
 def check_torch_compile() -> bool: # dynamo can be disabled after startup
     return use_torch_compile and not torch._dynamo.config.disable # pylint: disable=protected-access
 
-
-if devices.backend == "rocm":
-    gfx_version = devices.get_hip_agent().gfx_version
-    is_rdna2_and_older = bool(gfx_version < 0x940 or (gfx_version < 0x1100 and gfx_version >= 0x1000))
-else:
-    is_rdna2_and_older = False
-
-if os.environ.get("SDNQ_ALLOW_FP8_MM", None) is None:
-    if devices.backend == "cuda":
-        is_fp8_mm_supported = bool(torch.cuda.get_device_capability(devices.device) >= (8,9))
-    elif devices.backend == "rocm":
-        gfx_version = devices.get_hip_agent().gfx_version
-        is_fp8_mm_supported = bool(gfx_version >= 0x1200 or (gfx_version >= 0x940 and gfx_version < 0x1000))
-    else:
-        is_fp8_mm_supported = False
-else:
-    is_fp8_mm_supported = os.environ.get("SDNQ_ALLOW_FP8_MM", "0").lower() not in {"0", "false", "no"}
-
-if os.environ.get("SDNQ_ALLOW_FP8_COMPILE", None) is None:
-    if devices.backend == "cuda" and "linux" in sys.platform:
-        is_fp8_compile_supported = bool(torch.cuda.get_device_capability(devices.device) >= (8,9)) # triton has no e4m3 conversions before sm_89
-    else:
-        is_fp8_compile_supported = True
-else:
-    is_fp8_compile_supported = os.environ.get("SDNQ_ALLOW_FP8_COMPILE", "0").lower() not in {"0", "false", "no"}
-
-if os.environ.get("SDNQ_USE_OPENVINO_MM", None) is None:
-    use_openvino_mm = bool(devices.backend in {"cpu", "openvino"})
-else:
-    use_openvino_mm = bool(os.environ.get("SDNQ_USE_OPENVINO_MM", "0").lower() not in {"0", "false", "no"})
-
-if os.environ.get("SDNQ_USE_TRITON_MM", None) is None:
-    use_triton_mm = bool(is_rdna2_and_older or devices.backend == "zluda")
-else:
-    use_triton_mm = bool(os.environ.get("SDNQ_USE_TRITON_MM", "0").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
-    use_tensorwise_fp8_matmul = bool(devices.backend != "cuda" or (devices.backend == "cuda" and torch.cuda.get_device_capability(devices.device) < (9,0)))
-else:
-    use_tensorwise_fp8_matmul = os.environ.get("SDNQ_USE_TENSORWISE_FP8_MM", "0").lower() not in {"0", "false", "no"}
-
-
-fp_mm_func = None
-int_mm_func = None
-
-if use_openvino_mm:
-    try:
-        from .kernels.openvino_mm import openvino_int_mm, openvino_fp_mm
-        int_mm_func = openvino_int_mm
-        fp_mm_func = openvino_fp_mm
-    except Exception:
-        use_openvino_mm = False
-elif use_triton_mm:
-    try:
-        from .kernels.triton_mm import sdnq_triton_mm
-        int_mm_func = sdnq_triton_mm
-        fp_mm_func = sdnq_triton_mm
-    except Exception:
-        use_triton_mm = False
-
-if fp_mm_func is None and os.environ.get("SDNQ_USE_TRITON_MM", "1").lower() not in {"0", "false", "no"}:
-    try:
-        from .triton_mm import sdnq_triton_mm
-        fp_mm_func = sdnq_triton_mm
-    except Exception:
-        use_triton_mm = False
-
-if int_mm_func is None:
-    int_mm_func = torch._int_mm
-if fp_mm_func is None:
-    if devices.backend == "cuda":
-        def fp_mm_torch(x: torch.Tensor, y: torch.Tensor) -> torch.FloatTensor:
-            return torch.mm(x,y, out_dtype=torch.float32)
-    else:
-        def fp_mm_torch(x: torch.Tensor, y: torch.Tensor) -> torch.FloatTensor:
-            if y.dtype == torch.float8_e4m3fn:
-                fp16_scale = 4 * y.shape[-2]
-            else:
-                fp16_scale = 65536 * y.shape[-2]
-            in_scale = fp16_scale**0.5
-            x = x.to(dtype=torch.float32).div_(in_scale).to(dtype=torch.float16)
-            y = y.to(dtype=torch.float32).div_(in_scale).to(dtype=torch.float16)
-            return torch.mm(x,y).to(dtype=torch.float32).mul_(fp16_scale)
-    fp_mm_func = fp_mm_torch
-
-
-if os.environ.get("SDNQ_USE_CONTIGUOUS_MM", None) is None:
-    use_contiguous_int8_mm = bool(use_openvino_mm or is_rdna2_and_older or devices.backend in {"ipex", "mps", "openvino", "zluda"})
-    use_contiguous_fp16_mm = bool(use_contiguous_int8_mm or devices.backend == "rocm")
-else:
-    use_contiguous_int8_mm = bool(os.environ.get("SDNQ_USE_CONTIGUOUS_MM", "0").lower() not in {"0", "false", "no"})
-    use_contiguous_fp16_mm = use_contiguous_int8_mm
-
-
 if use_torch_compile:
     if hasattr(torch._dynamo.config, "recompile_limit"):
         torch._dynamo.config.recompile_limit = max(8192, getattr(torch._dynamo.config, "recompile_limit", 0))

modules/sdnq/dequantizer.py

@@ -5,7 +5,8 @@ from dataclasses import dataclass
 import torch
 
 from modules import devices
-from .common import dtype_dict, compile_func, use_contiguous_int8_mm, use_contiguous_fp16_mm, use_tensorwise_fp8_matmul, is_fp8_compile_supported
+from .common import dtype_dict, compile_func
+from .kernel_wrappers import use_contiguous_int8_mm, use_contiguous_fp16_mm, use_tensorwise_fp8_matmul, is_fp8_compile_supported
 from .quant_utils import quantize_int_mm, quantize_uint_mm, quantize_fp_mm, rotate_hadamard, get_hadamard
 from .packed_int import unpack_int
 from .packed_float import unpack_float

modules/sdnq/forward.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Callable
 
-from .common import dtype_dict, embedding_types, conv_types, conv_transpose_types, use_tensorwise_fp8_matmul
+from .common import dtype_dict, embedding_types, conv_types, conv_transpose_types
 
 
 def get_forward_func(layer_class_name: str, quantized_matmul_dtype: str, use_quantized_matmul: bool) -> Callable: # pylint: disable=inconsistent-return-statements
@@ -20,12 +20,8 @@ def get_forward_func(layer_class_name: str, quantized_matmul_dtype: str, use_qua
                     return quantized_conv_forward_int8_matmul
             else:
                 if dtype_dict[quantized_matmul_dtype]["num_bits"] == 8:
-                    if use_tensorwise_fp8_matmul:
-                        from .layers.conv.conv_fp8 import quantized_conv_forward_fp8_matmul
-                        return quantized_conv_forward_fp8_matmul
-                    else:
-                        from .layers.conv.conv_fp8_scaled import quantized_conv_forward_fp8_scaled_matmul
-                        return quantized_conv_forward_fp8_scaled_matmul
+                    from .layers.conv.conv_fp8 import quantized_conv_forward_fp8_matmul
+                    return quantized_conv_forward_fp8_matmul
                 else:
                     from .layers.conv.conv_fp16 import quantized_conv_forward_fp16_matmul
                     return quantized_conv_forward_fp16_matmul
@@ -53,12 +49,8 @@ def get_forward_func(layer_class_name: str, quantized_matmul_dtype: str, use_qua
                     return quantized_linear_forward_int8_matmul
             else:
                 if dtype_dict[quantized_matmul_dtype]["num_bits"] == 8:
-                    if use_tensorwise_fp8_matmul:
-                        from .layers.linear.linear_fp8 import quantized_linear_forward_fp8_matmul
-                        return quantized_linear_forward_fp8_matmul
-                    else:
-                        from .layers.linear.linear_fp8_scaled import quantized_linear_forward_fp8_scaled_matmul
-                        return quantized_linear_forward_fp8_scaled_matmul
+                    from .layers.linear.linear_fp8 import quantized_linear_forward_fp8_matmul
+                    return quantized_linear_forward_fp8_matmul
                 else:
                     from .layers.linear.linear_fp16 import quantized_linear_forward_fp16_matmul
                     return quantized_linear_forward_fp16_matmul

modules/sdnq/kernel_wrappers.py

@@ -0,0 +1,155 @@
+# pylint: disable=protected-access
+
+import os
+import sys
+import torch
+
+from modules import devices
+
+
+if os.environ.get("SDNQ_ALLOW_FP8_MM", None) is None:
+    if devices.backend == "cuda":
+        is_fp8_mm_supported = bool(torch.cuda.get_device_capability(devices.device) >= (8,9))
+    elif devices.backend == "rocm":
+        gfx_version = devices.get_hip_agent().gfx_version
+        is_fp8_mm_supported = bool(gfx_version >= 0x1200 or (gfx_version >= 0x940 and gfx_version < 0x1000))
+    else:
+        is_fp8_mm_supported = False
+else:
+    is_fp8_mm_supported = os.environ.get("SDNQ_ALLOW_FP8_MM", "0").lower() not in {"0", "false", "no"}
+
+if os.environ.get("SDNQ_ALLOW_FP8_COMPILE", None) is None:
+    if devices.backend == "cuda" and "linux" in sys.platform:
+        is_fp8_compile_supported = bool(torch.cuda.get_device_capability(devices.device) >= (8,9)) # triton has no e4m3 conversions before sm_89
+    else:
+        is_fp8_compile_supported = True
+else:
+    is_fp8_compile_supported = os.environ.get("SDNQ_ALLOW_FP8_COMPILE", "0").lower() not in {"0", "false", "no"}
+
+if devices.backend == "rocm":
+    gfx_version = devices.get_hip_agent().gfx_version
+    is_rdna2_and_older = bool(gfx_version < 0x940 or (gfx_version < 0x1100 and gfx_version >= 0x1000))
+else:
+    is_rdna2_and_older = False
+
+if os.environ.get("SDNQ_USE_OPENVINO_MM", None) is None:
+    use_openvino_mm = bool(devices.backend in {"cpu", "openvino"})
+else:
+    use_openvino_mm = bool(os.environ.get("SDNQ_USE_OPENVINO_MM", "0").lower() not in {"0", "false", "no"})
+
+if os.environ.get("SDNQ_USE_TRITON_MM", None) is None:
+    use_triton_mm = bool(is_rdna2_and_older or devices.backend == "zluda")
+else:
+    use_triton_mm = bool(os.environ.get("SDNQ_USE_TRITON_MM", "0").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
+    use_tensorwise_fp8_matmul = bool(devices.backend != "cuda" or (devices.backend == "cuda" and torch.cuda.get_device_capability(devices.device) < (9,0)))
+else:
+    use_tensorwise_fp8_matmul = os.environ.get("SDNQ_USE_TENSORWISE_FP8_MM", "0").lower() not in {"0", "false", "no"}
+
+if os.environ.get("SDNQ_USE_CONTIGUOUS_MM", None) is None:
+    use_contiguous_int8_mm = bool(use_openvino_mm or is_rdna2_and_older or devices.backend in {"ipex", "mps", "openvino", "zluda"})
+    use_contiguous_fp16_mm = bool(use_contiguous_int8_mm or devices.backend == "rocm")
+else:
+    use_contiguous_int8_mm = bool(os.environ.get("SDNQ_USE_CONTIGUOUS_MM", "0").lower() not in {"0", "false", "no"})
+    use_contiguous_fp16_mm = use_contiguous_int8_mm
+
+
+int_mm_func = None
+fp_mm_func = None
+fp8_mm_func = None
+int_scaled_mm_func = None
+fp_scaled_mm_func = None
+fp8_scaled_mm_func = None
+
+if use_openvino_mm:
+    try:
+        from .kernels.openvino_mm import openvino_int_mm, openvino_fp_mm
+        int_mm_func = openvino_int_mm
+        fp_mm_func = openvino_fp_mm
+        fp8_mm_func = openvino_fp_mm
+    except Exception:
+        use_openvino_mm = False
+elif use_triton_mm:
+    try:
+        from .kernels.triton_mm import sdnq_triton_mm
+        from .kernels.triton_scaled_mm import sdnq_scaled_mm
+        int_mm_func = sdnq_triton_mm
+        fp_mm_func = sdnq_triton_mm
+        fp8_mm_func = sdnq_triton_mm
+        int_scaled_mm_func = sdnq_scaled_mm
+        fp_scaled_mm_func = sdnq_scaled_mm
+        fp8_scaled_mm_func = sdnq_scaled_mm
+        use_tensorwise_fp8_matmul = False
+    except Exception:
+        use_triton_mm = False
+
+if fp_mm_func is None and os.environ.get("SDNQ_USE_TRITON_MM", "1").lower() not in {"0", "false", "no"}:
+    try:
+        from .kernels.triton_mm import sdnq_triton_mm
+        from .kernels.triton_scaled_mm import sdnq_scaled_mm
+        fp_mm_func = sdnq_triton_mm
+        fp_scaled_mm_func = sdnq_scaled_mm
+    except Exception:
+        use_triton_mm = False
+
+
+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)
+    int_mm_func = int_mm_torch
+
+if fp_mm_func is None:
+    if devices.backend == "cuda":
+        def fp_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
+            return torch.mm(a,b, out_dtype=out_dtype)
+    else:
+        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)
+    fp_mm_func = fp_mm_torch
+
+if fp8_mm_func is None:
+    def fp8_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
+        dummy_input_scale = torch.ones(1, device=a.device, dtype=torch.float32)
+        return torch._scaled_mm(a, b, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=out_dtype)
+    fp8_mm_func = fp8_mm_torch
+
+
+if int_scaled_mm_func is None:
+    def int_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
+        if bias is None:
+            return int_mm_func(a,b).to(dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype)
+        else:
+            return torch.addcmul(bias, int_mm_func(a,b).to(dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype)
+    int_scaled_mm_func = int_scaled_mm_torch
+
+if fp_scaled_mm_func is None:
+    def fp_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
+        if bias is None:
+            return fp_mm_func(a,b).to(dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype)
+        else:
+            return torch.addcmul(bias, fp_mm_func(a,b).to(dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype)
+    fp_scaled_mm_func = fp_scaled_mm_torch
+
+if fp8_scaled_mm_func is None:
+    if use_tensorwise_fp8_matmul:
+        def fp8_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
+            if bias is None:
+                return fp8_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype)
+            else:
+                return torch.addcmul(bias, fp8_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype)
+    else:
+        def fp8_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
+            if bias is not None and bias.ndim != 1:
+                return torch._scaled_mm(a, b, scale_a=scale_a, scale_b=scale_b, bias=None, out_dtype=out_dtype).add_(bias)
+            else:
+                return torch._scaled_mm(a, b, scale_a=scale_a, scale_b=scale_b, bias=bias.to(dtype=out_dtype) if bias is not None else None, out_dtype=out_dtype)
+    fp8_scaled_mm_func = fp8_scaled_mm_torch

modules/sdnq/kernels/openvino_mm.py

@@ -15,25 +15,25 @@ for ov_device in core.get_available_devices():
     core.set_property(ov_device, {ov_hints.execution_mode: ov_hints.ExecutionMode.ACCURACY})
 
 
-def ov_mm(A: torch.Tensor, B: torch.Tensor, infer_request: ov.InferRequest, out_name: str) -> torch.FloatTensor:
+def ov_mm(A: torch.Tensor, B: torch.Tensor, infer_request: ov.InferRequest, out_name: str, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
     C = torch.empty((A.shape[0], B.shape[-1]), device="cpu", dtype=torch.float32)
     infer_request.set_tensor("A", ov.Tensor(A.detach().contiguous().to("cpu").numpy(), shared_memory=True))
     infer_request.set_tensor("B", ov.Tensor(B.detach().contiguous().to("cpu").numpy(), shared_memory=True))
     infer_request.set_tensor(out_name, ov.Tensor(C.numpy(), shared_memory=True))
     infer_request.infer()
-    C = C.to(A.device)
+    C = C.to(A.device, dtype=out_dtype)
     return C
 
 
 @torch.library.custom_op("sdnq::openvino_int_mm", mutates_args=())
-def openvino_int_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor) -> torch.Tensor:
+def openvino_int_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
     if "GPU" not in OV_DEVICE:
         cache_key = (OV_DEVICE, "int8", Tensor_A.shape, Tensor_B.shape)
     else:
         cache_key = (OV_DEVICE, "int8", None, None)
     infer_request, out_name = OV_COMPILED_CACHE.get(cache_key, (None, None))
     if infer_request is not None:
-        return ov_mm(Tensor_A, Tensor_B, infer_request, out_name)
+        return ov_mm(Tensor_A, Tensor_B, infer_request, out_name, out_dtype=out_dtype)
 
     if "GPU" not in OV_DEVICE:
         shape_a = ov.Shape(Tensor_A.shape)
@@ -75,15 +75,15 @@ def openvino_int_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor) -> torch.Ten
     out_name = ov_model.outputs[0]
 
     OV_COMPILED_CACHE[cache_key] = (infer_request, out_name)
-    return ov_mm(Tensor_A, Tensor_B, infer_request, out_name)
+    return ov_mm(Tensor_A, Tensor_B, infer_request, out_name, out_dtype=out_dtype)
 
 @openvino_int_mm.register_fake
-def openvino_int_mm_fake(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
-    return torch.mm(A.to(dtype=torch.float32), B.to(dtype=torch.float32))
+def openvino_int_mm_fake(A: torch.Tensor, B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
+    return torch.mm(A.to(dtype=torch.float32), B.to(dtype=torch.float32)).to(dtype=out_dtype)
 
 
 @torch.library.custom_op("sdnq::openvino_fp_mm", mutates_args=())
-def openvino_fp_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor) -> torch.Tensor:
+def openvino_fp_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
     mm_dtype = "fp16" if Tensor_B.dtype == torch.float16 else "fp8"
     if mm_dtype == "fp8":
         Tensor_A = Tensor_A.to(dtype=torch.float16)
@@ -94,7 +94,7 @@ def openvino_fp_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor) -> torch.Tens
         cache_key = (OV_DEVICE, mm_dtype, None, None)
     infer_request, out_name = OV_COMPILED_CACHE.get(cache_key, (None, None))
     if infer_request is not None:
-        return ov_mm(Tensor_A, Tensor_B, infer_request, out_name)
+        return ov_mm(Tensor_A, Tensor_B, infer_request, out_name, out_dtype=out_dtype)
 
     if "GPU" not in OV_DEVICE:
         shape_a = ov.Shape(Tensor_A.shape)
@@ -135,8 +135,8 @@ def openvino_fp_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor) -> torch.Tens
     out_name = ov_model.outputs[0]
 
     OV_COMPILED_CACHE[cache_key] = (infer_request, out_name)
-    return ov_mm(Tensor_A, Tensor_B, infer_request, out_name)
+    return ov_mm(Tensor_A, Tensor_B, infer_request, out_name, out_dtype=out_dtype)
 
 @openvino_fp_mm.register_fake
-def openvino_fp_mm_fake(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
-    return torch.mm(A.to(dtype=torch.float32), B.to(dtype=torch.float32))
+def openvino_fp_mm_fake(A: torch.Tensor, B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
+    return torch.mm(A.to(dtype=torch.float32), B.to(dtype=torch.float32)).to(dtype=out_dtype)

modules/sdnq/kernels/triton_mm.py

@@ -18,14 +18,15 @@ matmul_configs = [
 ]
 
 
-@triton.autotune(configs=matmul_configs, key=["b_is_contiguous", "M_AT", "N_AT", "K_AT", "a_dtype", "out_dtype"], cache_results=True)
+@triton.autotune(configs=matmul_configs, key=["b_is_contiguous", "bias_ndim", "M_AT", "N_AT", "K_AT", "a_dtype", "out_dtype"], cache_results=True)
 @triton.jit
 def sdnq_triton_mm_kernel(
-    a_ptr, b_ptr, c_ptr,
+    a_ptr, b_ptr, c_ptr, bias_ptr,
     M: tl.constexpr,
     N: tl.constexpr,
     K: tl.constexpr,
     b_is_contiguous: tl.constexpr,
+    bias_ndim: tl.constexpr,
     M_AT: tl.constexpr, # pylint: disable=unused-argument
     N_AT: tl.constexpr, # pylint: disable=unused-argument
     K_AT: tl.constexpr, # pylint: disable=unused-argument
@@ -60,6 +61,7 @@ def sdnq_triton_mm_kernel(
     tl.assume(BLOCK_SIZE_K > 0)
     tl.assume(GROUP_SIZE_M > 0)
     tl.assume(b_is_contiguous == 0 or b_is_contiguous == 1) # pylint: disable=consider-using-in
+    tl.assume(bias_ndim >= 0 and bias_ndim <= 2) # pylint: disable=consider-using-in
 
     a_desc = tl.make_tensor_descriptor(base=a_ptr, shape=(M, K), strides=(K, 1), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_K))
     if b_is_contiguous:
@@ -70,7 +72,7 @@ def sdnq_triton_mm_kernel(
         b_ptrs = b_ptr + (offs_k[:, None] + offs_bn[None, :] * K)
 
     off_k = 0
-    accumulator_dtype = tl.int32 if a_ptr.type.element_ty == tl.int8 else tl.float32
+    accumulator_dtype: tl.constexpr = tl.int32 if a_ptr.type.element_ty == tl.int8 else tl.float32
     accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=accumulator_dtype)
     for _ in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K)):
         a = a_desc.load([off_m, off_k])
@@ -82,14 +84,33 @@ def sdnq_triton_mm_kernel(
         accumulator = tl.dot(a, b, accumulator, out_dtype=accumulator_dtype)
         off_k += BLOCK_SIZE_K
 
+    if bias_ndim == 1:
+        accumulator = accumulator.to(tl.float32)
+        bias_desc = tl.make_tensor_descriptor(base=bias_ptr, shape=(N,), strides=(1,), block_shape=(BLOCK_SIZE_N,))
+        bias = bias_desc.load([off_n])[None, :].to(tl.float32)
+        accumulator += bias
+    elif bias_ndim == 2:
+        accumulator = accumulator.to(tl.float32)
+        bias_desc = tl.make_tensor_descriptor(base=bias_ptr, shape=(M, N), strides=(N, 1), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N))
+        bias = bias_desc.load([off_m, off_n]).to(tl.float32)
+        accumulator += bias
+
     accumulator = accumulator.to(c_ptr.type.element_ty)
     c_desc = tl.make_tensor_descriptor(base=c_ptr, shape=(M, N), strides=(N, 1), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N))
     c_desc.store([off_m, off_n], accumulator)
 
 
-def sdnq_triton_mm(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype | None = None) -> torch.Tensor:
+def sdnq_triton_mm(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    bias: torch.FloatTensor | None = None,
+    out_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
     assert a.shape[1] == b.shape[0], "Incompatible dimensions"
     assert a.is_contiguous(), "Matrix A must be contiguous"
+    if bias is not None:
+        assert bias.is_contiguous(), "Bias must be contiguous"
+        assert bias.ndim in {1, 2}, "Bias must be 1D or 2D"
     M, K = a.shape
     K, N = b.shape
     if out_dtype is None:
@@ -98,9 +119,10 @@ def sdnq_triton_mm(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype | No
     def grid(META):
         return (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), )
     sdnq_triton_mm_kernel[grid](
-        a, b, c,
+        a, b, c, bias,
         M, N, K,
         (1 if b.is_contiguous() else 0),
+        (0 if bias is None else bias.ndim),
         math.ceil(M / min_block_size),
         math.ceil(N / min_block_size),
         math.ceil(K / min_block_size),

modules/sdnq/kernels/triton_scaled_mm.py

@@ -0,0 +1,147 @@
+"""
+W4A8 fallback with Triton.
+This is intended as a template for future INT4 MM kernels as Triton has no support for INT4 hardware yet.
+"""
+
+import os
+import math
+import torch
+
+import triton
+import triton.language as tl
+
+
+min_block_size = int(os.environ.get("SDNQ_TRITON_MM_MIN_BLOCK_SIZE", "256"))
+matmul_configs = [
+    triton.Config({"BLOCK_SIZE_M": BM, "BLOCK_SIZE_N": BN, "BLOCK_SIZE_K": BK, "GROUP_SIZE_M": GM}, num_warps=w, num_stages=s)
+    for BM in [int(BM) for BM in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_M_LIST", "64,128").replace(" ","").split(",")]
+    for BN in [int(BN) for BN in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_N_LIST", "64,128").replace(" ","").split(",")]
+    for BK in [int(BK) for BK in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_K_LIST", "64,128").replace(" ","").split(",")]
+    for GM in [int(GM) for GM in os.environ.get("SDNQ_TRITON_MM_GROUP_SIZE_M_LIST", "4,8").replace(" ","").split(",")]
+    for w in [int(w) for w in os.environ.get("SDNQ_TRITON_MM_NUM_WARPS_LIST", "4,8,16").replace(" ","").split(",")]
+    for s in [int(s) for s in os.environ.get("SDNQ_TRITON_MM_NUM_STAGES_LIST", "2").replace(" ","").split(",")]
+]
+
+
+@triton.autotune(configs=matmul_configs, key=["b_is_contiguous", "bias_ndim", "M_AT", "N_AT", "K_AT", "a_dtype", "out_dtype"], cache_results=True)
+@triton.jit
+def sdnq_scaled_mm_kernel(
+    a_ptr, b_ptr, c_ptr, bias_ptr,
+    scale_a_ptr, scale_b_ptr,
+    M: tl.constexpr,
+    N: tl.constexpr,
+    K: tl.constexpr,
+    b_is_contiguous: tl.constexpr,
+    bias_ndim: tl.constexpr,
+    M_AT: tl.constexpr, # pylint: disable=unused-argument
+    N_AT: tl.constexpr, # pylint: disable=unused-argument
+    K_AT: tl.constexpr, # pylint: disable=unused-argument
+    a_dtype: tl.constexpr, # pylint: disable=unused-argument
+    out_dtype: tl.constexpr, # pylint: disable=unused-argument
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+) -> None:
+    pid = tl.program_id(axis=0)
+    num_pid_m: tl.constexpr = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n: tl.constexpr = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group: tl.constexpr = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+    off_m = pid_m * BLOCK_SIZE_M
+    off_n = pid_n * BLOCK_SIZE_N
+
+    tl.assume(M > 0)
+    tl.assume(N > 0)
+    tl.assume(K > 0)
+    tl.assume(pid_m >= 0)
+    tl.assume(pid_n >= 0)
+    tl.assume(off_m >= 0)
+    tl.assume(off_n >= 0)
+    tl.assume(BLOCK_SIZE_M > 0)
+    tl.assume(BLOCK_SIZE_N > 0)
+    tl.assume(BLOCK_SIZE_K > 0)
+    tl.assume(GROUP_SIZE_M > 0)
+    tl.assume(b_is_contiguous == 0 or b_is_contiguous == 1) # pylint: disable=consider-using-in
+    tl.assume(bias_ndim >= 0 and bias_ndim <= 2) # pylint: disable=consider-using-in
+
+    a_desc = tl.make_tensor_descriptor(base=a_ptr, shape=(M, K), strides=(K, 1), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_K))
+    if b_is_contiguous:
+        b_desc = tl.make_tensor_descriptor(base=b_ptr, shape=(K, N), strides=(N, 1), block_shape=(BLOCK_SIZE_K, BLOCK_SIZE_N))
+    else:
+        offs_k = tl.arange(0, BLOCK_SIZE_K)
+        offs_bn = (off_n + tl.arange(0, BLOCK_SIZE_N)) % N
+        b_ptrs = b_ptr + (offs_k[:, None] + offs_bn[None, :] * K)
+
+    off_k = 0
+    accumulator_dtype: tl.constexpr = tl.int32 if a_ptr.type.element_ty == tl.int8 else tl.float32
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=accumulator_dtype)
+    for _ in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a = a_desc.load([off_m, off_k])
+        if b_is_contiguous:
+            b = b_desc.load([off_k, off_n])
+        else:
+            b = tl.load(b_ptrs, mask=offs_k[:, None] < K - off_k, other=0.0)
+            b_ptrs += BLOCK_SIZE_K
+        accumulator = tl.dot(a, b, accumulator, out_dtype=accumulator_dtype)
+        off_k += BLOCK_SIZE_K
+
+    scale_a_desc = tl.make_tensor_descriptor(base=scale_a_ptr, shape=(M,), strides=(1,), block_shape=(BLOCK_SIZE_M,))
+    scale_b_desc = tl.make_tensor_descriptor(base=scale_b_ptr, shape=(N,), strides=(1,), block_shape=(BLOCK_SIZE_N,))
+    scale_a = scale_a_desc.load([off_m])[:, None].to(tl.float32)
+    scale_b = scale_b_desc.load([off_n])[None, :].to(tl.float32)
+
+    if bias_ndim == 1:
+        accumulator = accumulator.to(tl.float32) * scale_a
+        bias_desc = tl.make_tensor_descriptor(base=bias_ptr, shape=(N,), strides=(1,), block_shape=(BLOCK_SIZE_N,))
+        bias = bias_desc.load([off_n])[None, :].to(tl.float32)
+        accumulator = tl.fma(accumulator, scale_b, bias)
+    elif bias_ndim == 2:
+        accumulator = accumulator.to(tl.float32) * scale_a
+        bias_desc = tl.make_tensor_descriptor(base=bias_ptr, shape=(M, N), strides=(N, 1), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N))
+        bias = bias_desc.load([off_m, off_n]).to(tl.float32)
+        accumulator = tl.fma(accumulator, scale_b, bias)
+    else:
+        accumulator = accumulator.to(tl.float32) * scale_a * scale_b
+
+    accumulator = accumulator.to(c_ptr.type.element_ty)
+    c_desc = tl.make_tensor_descriptor(base=c_ptr, shape=(M, N), strides=(N, 1), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N))
+    c_desc.store([off_m, off_n], accumulator)
+
+
+def sdnq_scaled_mm(
+    a: torch.Tensor,
+    b: torch.Tensor,
+    scale_a: torch.Tensor,
+    scale_b: torch.Tensor,
+    bias: torch.FloatTensor | None = None,
+    out_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    assert a.shape[1] == b.shape[0], "Incompatible dimensions"
+    assert a.is_contiguous(), "Matrix A must be contiguous"
+    assert scale_a.is_contiguous(), "Matrix A scale must be contiguous"
+    assert scale_b.is_contiguous(), "Matrix B scale must be contiguous"
+    if bias is not None:
+        assert bias.is_contiguous(), "Bias must be contiguous"
+        assert bias.ndim in {1, 2}, "Bias must be 1D or 2D"
+    M, K = a.shape
+    K, N = b.shape
+    c = torch.empty((M, N), device=a.device, dtype=out_dtype)
+    def grid(META):
+        return (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), )
+    sdnq_scaled_mm_kernel[grid](
+        a, b, c, bias,
+        scale_a, scale_b,
+        M, N, K,
+        (1 if b.is_contiguous() else 0),
+        (0 if bias is None else bias.ndim),
+        math.ceil(M / min_block_size),
+        math.ceil(N / min_block_size),
+        math.ceil(K / min_block_size),
+        str(a.dtype), str(c.dtype),
+    )
+    return c

modules/sdnq/layers/conv/conv_fp16.py

@@ -2,7 +2,8 @@
 
 import torch
 
-from ...common import compile_func, fp_mm_func
+from ...common import compile_func
+from ...kernel_wrappers import fp_mm_func, fp_scaled_mm_func
 from ...dequantizer import dequantize_symmetric, dequantize_asymmetric
 from ...quant_utils import rotate_hadamard, get_hadamard
 from ...packed_float import unpack_float
@@ -50,7 +51,7 @@ def conv_fp16_matmul(
     input, weight = check_mats(input, weight)
 
     if groups == 1:
-        result = fp_mm_func(input, weight).to(dtype=input_scale.dtype).mul_(input_scale)
+        result = fp_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(mm_output_shape)
     else:
         weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups)
         input = input.view(input.shape[0], groups, input.shape[1] // groups)
@@ -58,10 +59,10 @@ def conv_fp16_matmul(
         for i in range(groups):
             result.append(fp_mm_func(input[:, i], weight[:, i]))
         result = torch.cat(result, dim=-1).to(dtype=input_scale.dtype).mul_(input_scale)
-    if bias is not None:
-        dequantize_asymmetric(result, scale, bias, dtype=return_dtype, result_shape=mm_output_shape)
-    else:
-        dequantize_symmetric(result, scale, dtype=return_dtype, result_shape=mm_output_shape)
+        if bias is not None:
+            result = dequantize_asymmetric(result, scale, bias, dtype=return_dtype, result_shape=mm_output_shape)
+        else:
+            result = dequantize_symmetric(result, scale, dtype=return_dtype, result_shape=mm_output_shape)
 
     if conv_type == 1:
         result = result.transpose_(1,2)

modules/sdnq/layers/conv/conv_fp8.py

@@ -3,6 +3,7 @@
 import torch
 
 from ...common import compile_func
+from ...kernel_wrappers import fp8_mm_func, fp8_scaled_mm_func
 from ...dequantizer import dequantize_symmetric, dequantize_asymmetric
 from ...quant_utils import rotate_hadamard, get_hadamard
 from ...packed_float import unpack_float
@@ -46,21 +47,20 @@ def conv_fp8_matmul(
 
     input, input_scale = quantize_fp_mm_input(input, dtype=scale.dtype)
     input, weight = check_mats(input, weight)
-    dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
 
     if groups == 1:
-        result = torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=input_scale.dtype).mul_(input_scale)
+        result = fp8_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(mm_output_shape)
     else:
         weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups)
         input = input.view(input.shape[0], groups, input.shape[1] // groups)
         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=input_scale.dtype))
-        result = torch.cat(result, dim=-1).mul_(input_scale)
-    if bias is not None:
-        dequantize_asymmetric(result, scale, bias, dtype=return_dtype, result_shape=mm_output_shape)
-    else:
-        dequantize_symmetric(result, scale, dtype=return_dtype, result_shape=mm_output_shape)
+            result.append(fp8_mm_func(input[:, i], weight[:, i]))
+        result = torch.cat(result, dim=-1).to(dtype=input_scale.dtype).mul_(input_scale)
+        if bias is not None:
+            result = dequantize_asymmetric(result, scale, bias, dtype=return_dtype, result_shape=mm_output_shape)
+        else:
+            result = dequantize_symmetric(result, scale, dtype=return_dtype, result_shape=mm_output_shape)
 
     if conv_type == 1:
         result = result.transpose_(1,2)

modules/sdnq/layers/conv/conv_fp8_scaled.py

@@ -1,108 +0,0 @@
-# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
-
-import torch
-
-from ...common import compile_func
-from ...quant_utils import rotate_hadamard, get_hadamard
-from ...packed_float import unpack_float
-
-from .forward import get_conv_args, process_conv_input
-from ..linear.linear_fp8_scaled import quantize_fp_scaled_mm_input
-from ..linear.forward import check_mats
-
-
-def conv_fp8_scaled_matmul(
-    input: torch.FloatTensor,
-    weight: torch.Tensor,
-    scale: torch.FloatTensor,
-    result_shape: torch.Size,
-    reversed_padding_repeated_twice: list[int],
-    padding_mode: str, conv_type: int,
-    groups: int, stride: list[int],
-    padding: list[int], dilation: list[int],
-    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:
-    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)
-
-    if quantized_weight_shape is not None:
-        weight = unpack_float(weight, weights_dtype, quantized_weight_shape).to(dtype=torch.float8_e4m3fn).t_()
-        scale = scale.t()
-
-    if hadamard is not None:
-        input = rotate_hadamard(input, hadamard=hadamard)
-    if svd_up is not None:
-        input = input.flatten(0,-2)
-        svd_bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
-    input, input_scale = quantize_fp_scaled_mm_input(input)
-    input, weight = check_mats(input, weight)
-
-    if groups == 1:
-        if bias is not None and bias.dtype != torch.bfloat16:
-            bias = bias.to(dtype=torch.bfloat16)
-        result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=torch.bfloat16)
-    else:
-        scale = scale.view(groups, 1, scale.shape[1] // groups)
-        input_scale = input_scale.view(groups, input_scale.shape[0] // groups, 1)
-        weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups)
-        input = input.view(input.shape[0], groups, input.shape[1] // groups)
-        result = []
-        if bias is not None:
-            bias = bias.view(groups, bias.shape[0] // groups)
-            if bias.dtype != torch.bfloat16:
-                bias = bias.to(dtype=torch.bfloat16)
-            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=torch.bfloat16))
-        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=torch.bfloat16))
-        result = torch.cat(result, dim=-1)
-    if svd_up is not None:
-        result.add_(svd_bias)
-    result = result.view(mm_output_shape).to(return_dtype)
-
-    if conv_type == 1:
-        result = result.transpose_(1,2)
-    elif conv_type == 2:
-        result = result.permute(0,3,1,2)
-    elif conv_type == 3:
-        result = result.permute(0,4,1,2,3)
-    return result
-
-
-def quantized_conv_forward_fp8_scaled_matmul(self, input) -> torch.FloatTensor:
-    if torch.numel(input) / input.shape[2] < 32:
-        return self._conv_forward(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)
-    if self.sdnq_dequantizer.re_quantize_for_matmul:
-        weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, zero_point=self.zero_point)
-        quantized_weight_shape = None
-    else:
-        weight, scale = self.weight, self.scale
-        quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
-    if self.sdnq_dequantizer.use_hadamard:
-        hadamard = get_hadamard(self.sdnq_dequantizer.hadamard_group_size, dtype=input.dtype, device=input.device)
-    else:
-        hadamard = None
-
-    conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
-    return conv_fp8_scaled_matmul(
-        input, weight, scale,
-        self.sdnq_dequantizer.result_shape,
-        self._reversed_padding_repeated_twice,
-        self.padding_mode, conv_type,
-        self.groups, stride, padding, dilation,
-        bias=self.bias,
-        svd_up=self.svd_up,
-        svd_down=self.svd_down,
-        hadamard=hadamard,
-        quantized_weight_shape=quantized_weight_shape,
-        weights_dtype=self.sdnq_dequantizer.weights_dtype,
-    )
-
-
-conv_fp8_scaled_matmul = compile_func(conv_fp8_scaled_matmul)

modules/sdnq/layers/conv/conv_int8.py

@@ -2,7 +2,8 @@
 
 import torch
 
-from ...common import compile_func, int_mm_func
+from ...common import compile_func
+from ...kernel_wrappers import int_mm_func, int_scaled_mm_func
 from ...dequantizer import dequantize_symmetric, dequantize_asymmetric
 from ...quant_utils import rotate_hadamard, get_hadamard
 from ...packed_int import unpack_int
@@ -64,7 +65,7 @@ def conv_int8_matmul(
     input, weight = check_mats(input, weight)
 
     if groups == 1:
-        result = int_mm_func(input, weight).to(dtype=input_scale.dtype).mul_(input_scale)
+        result = int_scaled_mm_func(input, weight, input_scale, scale, bias=bias, out_dtype=return_dtype).view(mm_output_shape)
     else:
         weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups)
         input = input.view(input.shape[0], groups, input.shape[1] // groups)
@@ -72,10 +73,10 @@ def conv_int8_matmul(
         for i in range(groups):
             result.append(int_mm_func(input[:, i], weight[:, i]))
         result = torch.cat(result, dim=-1).to(dtype=input_scale.dtype).mul_(input_scale)
-    if bias is not None:
-        result = dequantize_asymmetric(result, scale, bias, dtype=return_dtype, result_shape=mm_output_shape)
-    else:
-        result = dequantize_symmetric(result, scale, dtype=return_dtype, result_shape=mm_output_shape)
+        if bias is not None:
+            result = dequantize_asymmetric(result, scale, bias, dtype=return_dtype, result_shape=mm_output_shape)
+        else:
+            result = dequantize_symmetric(result, scale, dtype=return_dtype, result_shape=mm_output_shape)
 
     if conv_type == 1:
         result = result.transpose_(1,2)

modules/sdnq/layers/conv/conv_uint8.py

@@ -2,7 +2,8 @@
 
 import torch
 
-from ...common import compile_func, int_mm_func
+from ...common import compile_func
+from ...kernel_wrappers import int_mm_func, int_scaled_mm_func
 from ...dequantizer import dequantize_asymmetric
 from ...quant_utils import rotate_hadamard, get_hadamard
 from ...packed_int import unpack_int
@@ -67,7 +68,7 @@ def conv_uint8_matmul(
     input, weight = check_mats(input, weight)
 
     if groups == 1:
-        result = int_mm_func(input, weight).to(dtype=input_scale.dtype).mul_(input_scale)
+        result = int_scaled_mm_func(input, weight, input_scale, scale, bias=zero_bias, out_dtype=return_dtype).view(mm_output_shape)
     else:
         weight = weight.view(weight.shape[0], groups, weight.shape[1] // groups)
         input = input.view(input.shape[0], groups, input.shape[1] // groups)
@@ -75,7 +76,7 @@ def conv_uint8_matmul(
         for i in range(groups):
             result.append(int_mm_func(input[:, i], weight[:, i]))
         result = torch.cat(result, dim=-1).to(dtype=input_scale.dtype).mul_(input_scale)
-    result = dequantize_asymmetric(result, scale, zero_bias, dtype=return_dtype, result_shape=mm_output_shape)
+        result = dequantize_asymmetric(result, scale, zero_bias, dtype=return_dtype, result_shape=mm_output_shape)
 
     if conv_type == 1:
         result = result.transpose_(1,2)

modules/sdnq/layers/linear/forward.py

@@ -2,7 +2,7 @@
 
 import torch
 
-from ...common import use_contiguous_int8_mm, use_contiguous_fp16_mm
+from ...kernel_wrappers import use_contiguous_int8_mm, use_contiguous_fp16_mm
 
 
 def check_mats(input: torch.Tensor, weight: torch.Tensor, allow_contiguous_mm: bool = True) -> tuple[torch.Tensor, torch.Tensor]:

modules/sdnq/layers/linear/linear_fp16.py

@@ -2,8 +2,8 @@
 
 import torch
 
-from ...common import compile_func, fp_mm_func
-from ...dequantizer import dequantize_symmetric, dequantize_asymmetric
+from ...common import compile_func
+from ...kernel_wrappers import fp_scaled_mm_func
 from ...quant_utils import rotate_hadamard, get_hadamard
 from ...packed_float import unpack_float
 
@@ -29,6 +29,7 @@ def fp16_matmul(
         weight = weight.to(dtype=torch.float16) # fp8 weights
     return_dtype = input.dtype
     output_shape = (*input.shape[:-1], weight.shape[-1])
+
     if hadamard is not None:
         input = rotate_hadamard(input, hadamard=hadamard)
     if svd_up is not None:
@@ -37,12 +38,10 @@ def fp16_matmul(
             bias = torch.addmm(bias.to(dtype=svd_down.dtype), torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
         else:
             bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
+
     input, input_scale = quantize_fp_mm_input(input, dtype=scale.dtype, matmul_dtype="float16")
     input, weight = check_mats(input, weight)
-    if bias is not None:
-        return dequantize_asymmetric(fp_mm_func(input, weight).to(dtype=input_scale.dtype).mul_(input_scale), scale, bias, dtype=return_dtype, result_shape=output_shape)
-    else:
-        return dequantize_symmetric(fp_mm_func(input, weight).to(dtype=input_scale.dtype).mul_(input_scale), scale, dtype=return_dtype, result_shape=output_shape)
+    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:

modules/sdnq/layers/linear/linear_fp8.py

@@ -3,7 +3,7 @@
 import torch
 
 from ...common import compile_func
-from ...dequantizer import dequantize_symmetric, dequantize_asymmetric
+from ...kernel_wrappers import fp8_scaled_mm_func
 from ...quant_utils import quantize_fp_mm, rotate_hadamard, get_hadamard
 from ...packed_float import unpack_float
 
@@ -36,6 +36,7 @@ def fp8_matmul(
         scale = scale.t()
     return_dtype = input.dtype
     output_shape = (*input.shape[:-1], weight.shape[-1])
+
     if hadamard is not None:
         input = rotate_hadamard(input, hadamard=hadamard)
     if svd_up is not None:
@@ -44,13 +45,10 @@ def fp8_matmul(
             bias = torch.addmm(bias.to(dtype=svd_down.dtype), torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
         else:
             bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
-    dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
+
     input, input_scale = quantize_fp_mm_input(input, dtype=scale.dtype)
     input, weight = check_mats(input, weight, allow_contiguous_mm=False)
-    if bias is not None:
-        return dequantize_asymmetric(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=input_scale.dtype).mul_(input_scale), scale, bias, dtype=return_dtype, result_shape=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=input_scale.dtype).mul_(input_scale), scale, dtype=return_dtype, result_shape=output_shape)
+    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:

modules/sdnq/layers/linear/linear_fp8_scaled.py

@@ -1,75 +0,0 @@
-# pylint: disable=relative-beyond-top-level,redefined-builtin,protected-access
-
-import torch
-
-from ...common import compile_func
-from ...quant_utils import quantize_fp_mm, rotate_hadamard, get_hadamard
-from ...packed_float import unpack_float
-
-from .forward import check_mats
-
-
-def quantize_fp_scaled_mm_input(input: torch.FloatTensor, matmul_dtype: str = "float8_e4m3fn") -> tuple[torch.Tensor, torch.FloatTensor]:
-    input = input.flatten(0,-2).to(dtype=torch.float32)
-    input, input_scale = quantize_fp_mm(input, dim=-1, matmul_dtype=matmul_dtype)
-    return input, input_scale
-
-
-def fp8_scaled_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:
-    if quantized_weight_shape is not None:
-        weight = unpack_float(weight, weights_dtype, quantized_weight_shape).to(dtype=torch.float8_e4m3fn).t_()
-        scale = scale.t()
-    return_dtype = input.dtype
-    output_shape = (*input.shape[:-1], weight.shape[-1])
-    if hadamard is not None:
-        input = rotate_hadamard(input, hadamard=hadamard)
-    if svd_up is not None:
-        input = input.flatten(0,-2)
-        svd_bias = torch.mm(torch.mm(input.to(dtype=svd_down.dtype), svd_down), svd_up)
-    input, input_scale = quantize_fp_scaled_mm_input(input)
-    input, weight = check_mats(input, weight, allow_contiguous_mm=False)
-    if bias is not None and bias.dtype != torch.bfloat16:
-        bias = bias.to(dtype=torch.bfloat16)
-    result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=torch.bfloat16)
-    if svd_up is not None:
-        result.add_(svd_bias)
-    result = result.view(output_shape).to(return_dtype)
-    return result
-
-
-def quantized_linear_forward_fp8_scaled_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)
-    if self.sdnq_dequantizer.re_quantize_for_matmul:
-        weight, scale = self.sdnq_dequantizer.re_quantize_matmul(self.weight, self.scale, zero_point=self.zero_point)
-        quantized_weight_shape = None
-    else:
-        weight, scale = self.weight, self.scale
-        quantized_weight_shape = self.sdnq_dequantizer.quantized_weight_shape if self.sdnq_dequantizer.is_packed else None
-    if self.sdnq_dequantizer.use_hadamard:
-        hadamard = get_hadamard(self.sdnq_dequantizer.hadamard_group_size, dtype=input.dtype, device=input.device)
-    else:
-        hadamard = None
-
-    return fp8_scaled_matmul(
-        input, weight, scale,
-        bias=self.bias,
-        svd_up=self.svd_up,
-        svd_down=self.svd_down,
-        hadamard=hadamard,
-        quantized_weight_shape=quantized_weight_shape,
-        weights_dtype=self.sdnq_dequantizer.weights_dtype,
-    )
-
-
-fp8_scaled_matmul = compile_func(fp8_scaled_matmul)

modules/sdnq/layers/linear/linear_int8.py

@@ -2,8 +2,8 @@
 
 import torch
 
-from ...common import compile_func, int_mm_func
-from ...dequantizer import dequantize_symmetric, dequantize_asymmetric
+from ...common import compile_func
+from ...kernel_wrappers import int_scaled_mm_func
 from ...quant_utils import quantize_int_mm, rotate_hadamard, get_hadamard
 from ...packed_int import unpack_int
 
@@ -64,12 +64,9 @@ def int8_matmul(
         if bias is not None:
             zero_bias.add_(bias)
         bias = zero_bias
-    input, weight = check_mats(input, weight)
 
-    if bias is not None:
-        return dequantize_asymmetric(int_mm_func(input, weight).to(dtype=input_scale.dtype).mul_(input_scale), scale, bias, dtype=return_dtype, result_shape=output_shape)
-    else:
-        return dequantize_symmetric(int_mm_func(input, weight).to(dtype=input_scale.dtype).mul_(input_scale), scale, dtype=return_dtype, result_shape=output_shape)
+    input, weight = check_mats(input, weight)
+    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:

modules/sdnq/layers/linear/linear_uint8.py

@@ -2,8 +2,8 @@
 
 import torch
 
-from ...common import compile_func, int_mm_func
-from ...dequantizer import dequantize_asymmetric
+from ...common import compile_func
+from ...kernel_wrappers import int_scaled_mm_func
 from ...quant_utils import quantize_uint_mm, rotate_hadamard, get_hadamard
 from ...packed_int import unpack_int
 
@@ -61,16 +61,16 @@ def uint8_matmul(
 
     input, input_scale, input_zero_point = quantize_uint_mm_input(input, dtype=scale.dtype)
     if zero_point is not None:
-        zero_bias = torch.sum(input, dim=-1, keepdim=True, dtype=torch.int32).to(input_scale.dtype).mul_(input_scale).mul(zero_point)
-        zero_bias.add_(torch.sum(weight, dim=0, keepdim=True, dtype=torch.int32).to(scale.dtype).mul_(scale).mul(input_zero_point))
+        zero_bias = torch.sum(input, dim=-1, keepdim=True, dtype=torch.int32).to(dtype=input_scale.dtype).mul_(input_scale).mul(zero_point)
+        zero_bias.add_(torch.sum(weight, dim=0, keepdim=True, dtype=torch.int32).to(dtype=scale.dtype).mul_(scale).mul(input_zero_point))
         zero_bias.add_(torch.mul(input_zero_point.mul_(input.shape[-1]), zero_point))
     else:
-        zero_bias = torch.sum(weight, dim=0, keepdim=True, dtype=torch.int32).to(scale.dtype).mul_(scale).mul(input_zero_point)
+        zero_bias = torch.sum(weight, dim=0, keepdim=True, dtype=torch.int32).to(dtype=scale.dtype).mul_(scale).mul(input_zero_point)
     if bias is not None:
         zero_bias.add_(bias)
 
     input, weight = check_mats(input, weight)
-    return dequantize_asymmetric(int_mm_func(input, weight).to(dtype=input_scale.dtype).mul_(input_scale), scale, zero_bias, dtype=return_dtype, result_shape=output_shape)
+    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:

modules/sdnq/loader.py

@@ -3,7 +3,8 @@ import json
 import torch
 
 from modules import shared
-from .common import dtype_dict, is_fp8_mm_supported, use_tensorwise_fp8_matmul, check_torch_compile, linear_types
+from .common import dtype_dict, check_torch_compile, linear_types
+from .kernel_wrappers import is_fp8_mm_supported, use_tensorwise_fp8_matmul
 from .quantizer import QuantizationMethod, SDNQConfig, SDNQQuantizer, sdnq_post_load_quant
 from .quant_utils import prepare_weight_for_matmul, prepare_svd_for_matmul
 from .utils import get_quant_args_from_config, check_param_name_in

modules/sdnq/quant_utils.py

@@ -3,7 +3,8 @@
 import torch
 
 from modules import devices
-from .common import dtype_dict, use_contiguous_int8_mm, use_contiguous_fp16_mm, conv_types, conv_transpose_types
+from .common import dtype_dict, conv_types, conv_transpose_types
+from .kernel_wrappers import use_contiguous_int8_mm, use_contiguous_fp16_mm
 from .utils import is_pow2, is_pow4, next_power_of_2
 
 

modules/sdnq/quantizer.py

@@ -23,8 +23,6 @@ from .common import (
     conv_types,
     conv_transpose_types,
     weights_dtype_order,
-    is_fp8_mm_supported,
-    use_tensorwise_fp8_matmul,
     check_torch_compile,
     compile_func,
 )
@@ -46,6 +44,7 @@ from .utils import (
     add_module_skip_keys,
 )
 
+from .kernel_wrappers import is_fp8_mm_supported, use_tensorwise_fp8_matmul
 from .dequantizer import SDNQDequantizer, dequantize_sdnq_model
 from .packed_int import pack_int
 from .packed_float import pack_float

pipelines/native_transformer.py

@@ -825,7 +825,8 @@ def build_component_prequantized(
     from accelerate import init_empty_weights
     from accelerate.utils import set_module_tensor_to_device
     from diffusers.utils import get_module_from_name
-    from modules.sdnq.common import dtype_dict, check_torch_compile, is_fp8_compile_supported
+    from modules.sdnq.common import dtype_dict, check_torch_compile
+    from modules.sdnq.kernel_wrappers import is_fp8_compile_supported
     from modules.sdnq.quantizer import SDNQConfig, SDNQQuantizer
     from modules.sdnq.dequantizer import SDNQDequantizer
     from modules.sdnq.layers import get_sdnq_wrapper_class