Cleanup SDNQ compile

modules/sdnq/common.py

@@ -2,6 +2,8 @@
 
 import os
 import torch
+from functools import partial
+
 from modules import shared
 
 torch_version = float(torch.__version__[:3])
@@ -42,3 +44,7 @@ allowed_types = linear_types + conv_types + conv_transpose_types
 if use_torch_compile:
     torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit)
     torch._dynamo.config.accumulated_recompile_limit = max(8192, torch._dynamo.config.accumulated_recompile_limit)
+    compile_func = partial(torch.compile, fullgraph=True, dynamic=False)
+else:
+    def compile_func(fn, **kwargs):
+        return fn

modules/sdnq/dequantizer.py

@@ -4,7 +4,7 @@ from typing import Tuple
 
 import torch
 
-from .common import dtype_dict, use_torch_compile
+from .common import dtype_dict, compile_func
 from .packed_int import pack_int_symetric, unpack_int_symetric, pack_int_asymetric, unpack_int_asymetric
 
 
@@ -226,21 +226,11 @@ dequantizer_dict = {
 }
 
 
-if use_torch_compile:
-    dequantize_asymmetric_compiled = torch.compile(dequantize_asymmetric, fullgraph=True, dynamic=False)
-    dequantize_symmetric_compiled = torch.compile(dequantize_symmetric, fullgraph=True, dynamic=False)
-    dequantize_packed_int_asymmetric_compiled = torch.compile(dequantize_packed_int_asymmetric, fullgraph=True, dynamic=False)
-    dequantize_packed_int_symmetric_compiled = torch.compile(dequantize_packed_int_symmetric, fullgraph=True, dynamic=False)
-    re_quantize_matmul_asymmetric_compiled = torch.compile(re_quantize_matmul_asymmetric, fullgraph=True, dynamic=False)
-    re_quantize_matmul_symmetric_compiled = torch.compile(re_quantize_matmul_symmetric, fullgraph=True, dynamic=False)
-    re_quantize_matmul_packed_int_asymmetric_compiled = torch.compile(re_quantize_matmul_packed_int_asymmetric, fullgraph=True, dynamic=False)
-    re_quantize_matmul_packed_int_symmetric_compiled = torch.compile(re_quantize_matmul_packed_int_symmetric, fullgraph=True, dynamic=False)
-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
-    re_quantize_matmul_asymmetric_compiled = re_quantize_matmul_asymmetric
-    re_quantize_matmul_symmetric_compiled = re_quantize_matmul_symmetric
-    re_quantize_matmul_packed_int_asymmetric_compiled = re_quantize_matmul_packed_int_asymmetric
-    re_quantize_matmul_packed_int_symmetric_compiled = re_quantize_matmul_packed_int_symmetric
+dequantize_asymmetric_compiled = compile_func(dequantize_asymmetric)
+dequantize_symmetric_compiled = compile_func(dequantize_symmetric)
+dequantize_packed_int_asymmetric_compiled = compile_func(dequantize_packed_int_asymmetric)
+dequantize_packed_int_symmetric_compiled = compile_func(dequantize_packed_int_symmetric)
+re_quantize_matmul_asymmetric_compiled = compile_func(re_quantize_matmul_asymmetric)
+re_quantize_matmul_symmetric_compiled = compile_func(re_quantize_matmul_symmetric)
+re_quantize_matmul_packed_int_asymmetric_compiled = compile_func(re_quantize_matmul_packed_int_asymmetric)
+re_quantize_matmul_packed_int_symmetric_compiled = compile_func(re_quantize_matmul_packed_int_symmetric)

modules/sdnq/layers/conv/conv_fp8.py

@@ -4,7 +4,7 @@ from typing import List
 
 import torch
 
-from ...common import use_torch_compile # noqa: TID252
+from ...common import compile_func # noqa: TID252
 from ..linear.linear_fp8 import quantize_fp8_matmul_input # noqa: TID252
 from .forward import get_conv_args, process_conv_input
 
@@ -68,5 +68,4 @@ def quantized_conv_forward_fp8_matmul(self, input) -> torch.FloatTensor:
     )
 
 
-if use_torch_compile:
-    conv_fp8_matmul = torch.compile(conv_fp8_matmul, fullgraph=True, dynamic=False)
+conv_fp8_matmul = compile_func(conv_fp8_matmul)

modules/sdnq/layers/conv/conv_fp8_tensorwise.py

@@ -4,7 +4,7 @@ from typing import List
 
 import torch
 
-from ...common import use_torch_compile # noqa: TID252
+from ...common import compile_func # noqa: TID252
 from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252
 from ..linear.linear_fp8_tensorwise import quantize_fp8_matmul_input_tensorwise # noqa: TID252
 from .forward import get_conv_args, process_conv_input
@@ -63,5 +63,4 @@ def quantized_conv_forward_fp8_matmul_tensorwise(self, input) -> torch.FloatTens
     )
 
 
-if use_torch_compile:
-    conv_fp8_matmul_tensorwise = torch.compile(conv_fp8_matmul_tensorwise, fullgraph=True, dynamic=False)
+conv_fp8_matmul_tensorwise = compile_func(conv_fp8_matmul_tensorwise)

modules/sdnq/layers/conv/conv_int8.py

@@ -4,7 +4,7 @@ from typing import List
 
 import torch
 
-from ...common import use_torch_compile # noqa: TID252
+from ...common import compile_func # noqa: TID252
 from ...packed_int import unpack_int_symetric # noqa: TID252
 from ...dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252
 from ..linear.linear_int8 import quantize_int8_matmul_input # noqa: TID252
@@ -75,5 +75,4 @@ def quantized_conv_forward_int8_matmul(self, input) -> torch.FloatTensor:
     )
 
 
-if use_torch_compile:
-    conv_int8_matmul = torch.compile(conv_int8_matmul, fullgraph=True, dynamic=False)
+conv_int8_matmul = compile_func(conv_int8_matmul)

modules/sdnq/layers/linear/linear_fp8.py

@@ -4,7 +4,7 @@ from typing import Tuple
 
 import torch
 
-from ...common import use_torch_compile # noqa: TID252
+from ...common import compile_func # noqa: TID252
 from ...dequantizer import quantize_fp8 # noqa: TID252
 
 
@@ -34,5 +34,4 @@ def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch
     return fp8_matmul(input, self.weight, self.bias, self.sdnq_dequantizer.scale)
 
 
-if use_torch_compile:
-    fp8_matmul = torch.compile(fp8_matmul, fullgraph=True, dynamic=False)
+fp8_matmul = compile_func(fp8_matmul)

modules/sdnq/layers/linear/linear_fp8_tensorwise.py

@@ -4,7 +4,7 @@ from typing import Tuple
 
 import torch
 
-from ...common import use_torch_compile # noqa: TID252
+from ...common import compile_func # noqa: TID252
 from ...dequantizer import quantize_fp8, dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252
 
 
@@ -39,5 +39,4 @@ def quantized_linear_forward_fp8_matmul_tensorwise(self, input: torch.FloatTenso
     return fp8_matmul_tensorwise(input, self.weight, self.bias, self.sdnq_dequantizer.scale)
 
 
-if use_torch_compile:
-    fp8_matmul_tensorwise = torch.compile(fp8_matmul_tensorwise, fullgraph=True, dynamic=False)
+fp8_matmul_tensorwise = compile_func(fp8_matmul_tensorwise)

modules/sdnq/layers/linear/linear_int8.py

@@ -4,7 +4,7 @@ from typing import Tuple
 
 import torch
 
-from ...common import use_torch_compile # noqa: TID252
+from ...common import compile_func # noqa: TID252
 from ...packed_int import unpack_int_symetric # noqa: TID252
 from ...dequantizer import quantize_int8, dequantize_symmetric, dequantize_symmetric_with_bias # noqa: TID252
 
@@ -50,5 +50,4 @@ def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torc
     return int8_matmul(input, weight, self.bias, scale, quantized_weight_shape, self.sdnq_dequantizer.weights_dtype)
 
 
-if use_torch_compile:
-    int8_matmul = torch.compile(int8_matmul, fullgraph=True, dynamic=False)
+int8_matmul = compile_func(int8_matmul)