SDNQ include torch mm in compile graph and add SDNQ_INCLUDE_MM_KERNEL_IN_COMPILE

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]

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)

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)

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)

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)

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)