Cleanup SDNQ and add SDNQ_USE_TENSORWISE_FP8_MATMUL env var

modules/sdnq/common.py

@@ -1,5 +1,6 @@
 # pylint: disable=redefined-builtin,no-member,protected-access
 
+import os
 import torch
 from modules import devices, shared
 
@@ -31,7 +32,7 @@ if hasattr(torch, "float8_e5m2fnuz"):
     dtype_dict["float8_e5m2fnuz"] = {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": "fp8", "torch_dtype": torch.float8_e5m2fnuz, "storage_dtype": torch.float8_e5m2fnuz, "is_unsigned": False, "is_integer": False}
 
 use_torch_compile = shared.opts.sdnq_dequantize_compile # this setting requires a full restart of the webui to apply
-use_tensorwise_fp8_matmul = True # row-wise FP8 only exist on H100 hardware, sdnq will use software row-wise with tensorwise hardware with this setting
+use_tensorwise_fp8_matmul = os.environ.get('SDNQ_USE_TENSORWISE_FP8_MATMUL', "1").lower() not in {"0", "false", "no"} # row-wise FP8 only exist on H100 hardware, sdnq will use software row-wise with tensorwise hardware with this setting
 
 quantized_matmul_dtypes = ("int8", "int7", "int6", "int5", "int4", "int3", "int2", "float8_e4m3fn", "float8_e5m2")
 if devices.backend in {"cpu", "openvino"}:
@@ -43,8 +44,5 @@ conv_transpose_types = ("ConvTranspose1d", "ConvTranspose2d", "ConvTranspose3d")
 allowed_types = linear_types + conv_types + conv_transpose_types
 
 if use_torch_compile:
-    try:
-        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)
-    except Exception as e:
-        shared.log.warning(f"Quantization: type=sdnq Failed to increase the cache size for torch.compile: {e}")
+    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)

modules/sdnq/dequantizer.py

@@ -1,7 +1,6 @@
 # pylint: disable=redefined-builtin,no-member,protected-access
 
 import torch
-from modules import shared
 
 from .common import dtype_dict, use_torch_compile
 from .packed_int import pack_int_symetric, unpack_int_symetric, pack_int_asymetric, unpack_int_asymetric
@@ -170,17 +169,10 @@ dequantizer_dict = {
 
 
 if use_torch_compile:
-    try:
-        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)
-    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
+    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)
 else:
     dequantize_asymmetric_compiled = dequantize_asymmetric
     dequantize_symmetric_compiled = dequantize_symmetric

modules/sdnq/layers/conv/conv_fp8.py

@@ -65,7 +65,4 @@ def quantized_conv_forward_fp8_matmul(self, input) -> torch.FloatTensor:
 
 
 if use_torch_compile:
-    try:
-        conv_fp8_matmul = torch.compile(conv_fp8_matmul, fullgraph=True, dynamic=False)
-    except Exception:
-        pass
+    conv_fp8_matmul = torch.compile(conv_fp8_matmul, fullgraph=True, dynamic=False)

modules/sdnq/layers/conv/conv_fp8_tensorwise.py

@@ -64,7 +64,4 @@ def quantized_conv_forward_fp8_matmul_tensorwise(self, input) -> torch.FloatTens
 
 
 if use_torch_compile:
-    try:
-        conv_fp8_matmul_tensorwise = torch.compile(conv_fp8_matmul_tensorwise, fullgraph=True, dynamic=False)
-    except Exception:
-        pass
+    conv_fp8_matmul_tensorwise = torch.compile(conv_fp8_matmul_tensorwise, fullgraph=True, dynamic=False)

modules/sdnq/layers/conv/conv_int8.py

@@ -70,7 +70,4 @@ def quantized_conv_forward_int8_matmul(self, input) -> torch.FloatTensor:
 
 
 if use_torch_compile:
-    try:
-        conv_int8_matmul = torch.compile(conv_int8_matmul, fullgraph=True, dynamic=False)
-    except Exception:
-        pass
+    conv_int8_matmul = torch.compile(conv_int8_matmul, fullgraph=True, dynamic=False)

modules/sdnq/layers/linear/linear_fp8.py

@@ -34,7 +34,4 @@ def quantized_linear_forward_fp8_matmul(self, input: torch.FloatTensor) -> torch
 
 
 if use_torch_compile:
-    try:
-        fp8_matmul = torch.compile(fp8_matmul, fullgraph=True, dynamic=False)
-    except Exception:
-        pass
+    fp8_matmul = torch.compile(fp8_matmul, fullgraph=True, dynamic=False)

modules/sdnq/layers/linear/linear_fp8_tensorwise.py

@@ -42,7 +42,4 @@ def quantized_linear_forward_fp8_matmul_tensorwise(self, input: torch.FloatTenso
 
 
 if use_torch_compile:
-    try:
-        fp8_matmul_tensorwise = torch.compile(fp8_matmul_tensorwise, fullgraph=True, dynamic=False)
-    except Exception:
-        pass
+    fp8_matmul_tensorwise = torch.compile(fp8_matmul_tensorwise, fullgraph=True, dynamic=False)

modules/sdnq/layers/linear/linear_int8.py

@@ -46,7 +46,4 @@ def quantized_linear_forward_int8_matmul(self, input: torch.FloatTensor) -> torc
 
 
 if use_torch_compile:
-    try:
-        int8_matmul = torch.compile(int8_matmul, fullgraph=True, dynamic=False)
-    except Exception:
-        pass
+    int8_matmul = torch.compile(int8_matmul, fullgraph=True, dynamic=False)