From 32650c7e508c117de391726362edf3a4875f7b26 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Mon, 13 Jul 2026 18:01:01 +0300 Subject: [PATCH] SDNQ add fused Triton kernels --- modules/sdnq/common.py | 99 +---------- modules/sdnq/dequantizer.py | 3 +- modules/sdnq/forward.py | 18 +- modules/sdnq/kernel_wrappers.py | 155 ++++++++++++++++++ modules/sdnq/kernels/openvino_mm.py | 24 +-- modules/sdnq/kernels/triton_mm.py | 32 +++- modules/sdnq/kernels/triton_scaled_mm.py | 147 +++++++++++++++++ modules/sdnq/layers/conv/conv_fp16.py | 13 +- modules/sdnq/layers/conv/conv_fp8.py | 16 +- modules/sdnq/layers/conv/conv_fp8_scaled.py | 108 ------------ modules/sdnq/layers/conv/conv_int8.py | 13 +- modules/sdnq/layers/conv/conv_uint8.py | 7 +- modules/sdnq/layers/linear/forward.py | 2 +- modules/sdnq/layers/linear/linear_fp16.py | 11 +- modules/sdnq/layers/linear/linear_fp8.py | 10 +- .../sdnq/layers/linear/linear_fp8_scaled.py | 75 --------- modules/sdnq/layers/linear/linear_int8.py | 11 +- modules/sdnq/layers/linear/linear_uint8.py | 12 +- modules/sdnq/loader.py | 3 +- modules/sdnq/quant_utils.py | 3 +- modules/sdnq/quantizer.py | 3 +- pipelines/native_transformer.py | 3 +- 22 files changed, 403 insertions(+), 365 deletions(-) create mode 100644 modules/sdnq/kernel_wrappers.py create mode 100644 modules/sdnq/kernels/triton_scaled_mm.py delete mode 100644 modules/sdnq/layers/conv/conv_fp8_scaled.py delete mode 100644 modules/sdnq/layers/linear/linear_fp8_scaled.py diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py index 7aaac74e0..6d4031d37 100644 --- a/modules/sdnq/common.py +++ b/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)) diff --git a/modules/sdnq/dequantizer.py b/modules/sdnq/dequantizer.py index b63d06514..1344f60a2 100644 --- a/modules/sdnq/dequantizer.py +++ b/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 diff --git a/modules/sdnq/forward.py b/modules/sdnq/forward.py index f66275b74..513be32bb 100644 --- a/modules/sdnq/forward.py +++ b/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 diff --git a/modules/sdnq/kernel_wrappers.py b/modules/sdnq/kernel_wrappers.py new file mode 100644 index 000000000..fefcd0177 --- /dev/null +++ b/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 diff --git a/modules/sdnq/kernels/openvino_mm.py b/modules/sdnq/kernels/openvino_mm.py index a6cf9255e..3e850810c 100644 --- a/modules/sdnq/kernels/openvino_mm.py +++ b/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) diff --git a/modules/sdnq/kernels/triton_mm.py b/modules/sdnq/kernels/triton_mm.py index ab0bc2bcd..9b64c3ff8 100644 --- a/modules/sdnq/kernels/triton_mm.py +++ b/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), diff --git a/modules/sdnq/kernels/triton_scaled_mm.py b/modules/sdnq/kernels/triton_scaled_mm.py new file mode 100644 index 000000000..7fd6c40f5 --- /dev/null +++ b/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 diff --git a/modules/sdnq/layers/conv/conv_fp16.py b/modules/sdnq/layers/conv/conv_fp16.py index 83783e63b..19a0d57bb 100644 --- a/modules/sdnq/layers/conv/conv_fp16.py +++ b/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) diff --git a/modules/sdnq/layers/conv/conv_fp8.py b/modules/sdnq/layers/conv/conv_fp8.py index 23afc9cf3..8f0147f52 100644 --- a/modules/sdnq/layers/conv/conv_fp8.py +++ b/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) diff --git a/modules/sdnq/layers/conv/conv_fp8_scaled.py b/modules/sdnq/layers/conv/conv_fp8_scaled.py deleted file mode 100644 index 8ed336e68..000000000 --- a/modules/sdnq/layers/conv/conv_fp8_scaled.py +++ /dev/null @@ -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) diff --git a/modules/sdnq/layers/conv/conv_int8.py b/modules/sdnq/layers/conv/conv_int8.py index e95ba9990..a35c5a9e4 100644 --- a/modules/sdnq/layers/conv/conv_int8.py +++ b/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) diff --git a/modules/sdnq/layers/conv/conv_uint8.py b/modules/sdnq/layers/conv/conv_uint8.py index 615430e55..a87beb0eb 100644 --- a/modules/sdnq/layers/conv/conv_uint8.py +++ b/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) diff --git a/modules/sdnq/layers/linear/forward.py b/modules/sdnq/layers/linear/forward.py index dbdee7521..1a323f0a1 100644 --- a/modules/sdnq/layers/linear/forward.py +++ b/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]: diff --git a/modules/sdnq/layers/linear/linear_fp16.py b/modules/sdnq/layers/linear/linear_fp16.py index fcfec43f9..e4835970c 100644 --- a/modules/sdnq/layers/linear/linear_fp16.py +++ b/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: diff --git a/modules/sdnq/layers/linear/linear_fp8.py b/modules/sdnq/layers/linear/linear_fp8.py index 4b5610e87..8c6de8312 100644 --- a/modules/sdnq/layers/linear/linear_fp8.py +++ b/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: diff --git a/modules/sdnq/layers/linear/linear_fp8_scaled.py b/modules/sdnq/layers/linear/linear_fp8_scaled.py deleted file mode 100644 index 949c58586..000000000 --- a/modules/sdnq/layers/linear/linear_fp8_scaled.py +++ /dev/null @@ -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) diff --git a/modules/sdnq/layers/linear/linear_int8.py b/modules/sdnq/layers/linear/linear_int8.py index c7a0faa29..9f39d624b 100644 --- a/modules/sdnq/layers/linear/linear_int8.py +++ b/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: diff --git a/modules/sdnq/layers/linear/linear_uint8.py b/modules/sdnq/layers/linear/linear_uint8.py index 7ae8a1379..7812a8759 100644 --- a/modules/sdnq/layers/linear/linear_uint8.py +++ b/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: diff --git a/modules/sdnq/loader.py b/modules/sdnq/loader.py index eb80fe5b1..276c609cd 100644 --- a/modules/sdnq/loader.py +++ b/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 diff --git a/modules/sdnq/quant_utils.py b/modules/sdnq/quant_utils.py index 41ea09670..b215275f9 100644 --- a/modules/sdnq/quant_utils.py +++ b/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 diff --git a/modules/sdnq/quantizer.py b/modules/sdnq/quantizer.py index 530c2d66e..66a4131fc 100644 --- a/modules/sdnq/quantizer.py +++ b/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 diff --git a/pipelines/native_transformer.py b/pipelines/native_transformer.py index 398eb67cb..441c5f1c0 100644 --- a/pipelines/native_transformer.py +++ b/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