diff --git a/modules/sdnq/kernel_wrappers.py b/modules/sdnq/kernel_wrappers.py index b713bba6f..a3cdccd2a 100644 --- a/modules/sdnq/kernel_wrappers.py +++ b/modules/sdnq/kernel_wrappers.py @@ -38,16 +38,10 @@ if devices.backend in {"ipex", "xpu"}: else: is_alchemist_or_igpu = 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(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 = 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 @@ -56,65 +50,53 @@ else: use_tensorwise_fp8_matmul = bool(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", "xpu", "mps", "openvino", "zluda"}) + use_contiguous_int8_mm = bool(is_rdna2_and_older or devices.backend in {"ipex", "xpu", "cpu", "mps", "openvino", "zluda"}) use_contiguous_fp16_mm = bool(use_contiguous_int8_mm or devices.backend == "rocm") - use_contiguous_fp8_mm = use_contiguous_fp16_mm + use_contiguous_fp8_mm = use_contiguous_fp16_mm and (is_fp8_mm_supported or use_triton_mm) 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 - use_contiguous_fp8_mm = use_contiguous_fp16_mm + use_contiguous_fp8_mm = use_contiguous_fp16_mm and (is_fp8_mm_supported or use_triton_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 +use_openvino_mm = bool(os.environ.get("SDNQ_USE_OPENVINO_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 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 as e: use_openvino_mm = False - shared.log.warning(f"SDNQ: OpenVINO kernels are not available! Falling back to PyTorch Eager kernels. Error message: {e}") -elif use_triton_mm: + openvino_int_mm = None + openvino_fp_mm = None + shared.log.warning(f"SDNQ: OpenVINO MM kernels are not available! Falling back to PyTorch Eager kernels for CPU device. Error message: {e}") +else: + openvino_int_mm = None + openvino_fp_mm = None + + +if use_triton_mm: try: from .kernels.triton_mm import sdnq_triton_mm - int_mm_func = sdnq_triton_mm - fp_mm_func = sdnq_triton_mm if is_fp8_mm_supported: - fp8_mm_func = sdnq_triton_mm use_tensorwise_fp8_matmul = True - if use_triton_scaled_mm: - from .kernels.triton_scaled_mm import sdnq_scaled_mm - int_scaled_mm_func = sdnq_scaled_mm - fp_scaled_mm_func = sdnq_scaled_mm - if is_fp8_mm_supported: - fp8_scaled_mm_func = sdnq_scaled_mm except Exception as e: use_triton_mm = False - use_triton_scaled_mm = False - shared.log.warning(f"SDNQ: Triton kernels are not available! Falling back to PyTorch Eager kernels. Error message: {e}") + sdnq_triton_mm = None + shared.log.warning(f"SDNQ: Triton MM kernels are not available! Falling back to PyTorch Eager kernels. Error message: {e}") +else: + sdnq_triton_mm = None -if ( - fp_mm_func is None and not is_alchemist_or_igpu - and devices.backend in {"cuda", "rocm", "ipex", "xpu", "zluda"} - and os.environ.get("SDNQ_USE_TRITON_MM", "1").lower() not in {"0", "false", "no"} -): +if use_triton_scaled_mm: try: - from .kernels.triton_mm import sdnq_triton_mm - fp_mm_func = sdnq_triton_mm - if use_triton_scaled_mm: - from .kernels.triton_scaled_mm import sdnq_scaled_mm - fp_scaled_mm_func = sdnq_scaled_mm - except Exception: - use_triton_mm = False + from .kernels.triton_scaled_mm import sdnq_scaled_mm + except Exception as e: use_triton_scaled_mm = False + sdnq_scaled_mm = None + shared.log.warning(f"SDNQ: Triton Scaled MM kernels are not available! Falling back to PyTorch Eager kernels. Error message: {e}") +else: + sdnq_scaled_mm = None if os.environ.get("SDNQ_INCLUDE_MM_KERNEL_IN_COMPILE", None) is None: @@ -123,8 +105,14 @@ 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 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) + + +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) + def fp_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor: if b.dtype == torch.float8_e4m3fn: @@ -137,59 +125,91 @@ def fp_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch 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) - int_mm_func = int_mm_torch - - -if fp_mm_func is None: - if devices.backend == "cuda": - fp_mm_func = fp_mm_torch_cuda +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, out_dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype) else: - fp_mm_func = fp_mm_torch + return torch.addcmul(bias, int_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype) -if fp8_mm_func is None: +if use_tensorwise_fp8_matmul or not is_fp8_mm_supported: + 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) + + +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, out_dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype) + else: + return torch.addcmul(bias, fp_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype) + + +def int_mm_func(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.int32) -> torch.FloatTensor: + if sdnq_triton_mm is not None and a.device.type in {"cuda", "xpu"}: + return sdnq_triton_mm(a, b, out_dtype=out_dtype) + elif openvino_int_mm is not None and a.device.type == "cpu": + return openvino_int_mm(a, b, out_dtype=out_dtype) + else: + return int_mm_torch(a, b, out_dtype=out_dtype) + + +def fp8_mm_func(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor: if is_fp8_mm_supported: - 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 - use_contiguous_fp8_mm = False - else: - fp8_mm_func = fp_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, out_dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype) + if sdnq_triton_mm is not None and a.device.type in {"cuda", "xpu"}: + return sdnq_triton_mm(a, b, out_dtype=out_dtype) + elif openvino_fp_mm is not None and a.device.type == "cpu": + return openvino_fp_mm(a, b, out_dtype=out_dtype) else: - return torch.addcmul(bias, int_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype) - int_scaled_mm_func = compile_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, out_dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype) - else: - return torch.addcmul(bias, fp_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype) - fp_scaled_mm_func = compile_func(fp_scaled_mm_torch) - - -if fp8_scaled_mm_func is None: - if use_tensorwise_fp8_matmul or not is_fp8_mm_supported: - 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) + return fp8_mm_torch(a, b, out_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 = compile_func(fp8_scaled_mm_torch) + if openvino_fp_mm is not None and a.device.type == "cpu": + return openvino_fp_mm(a, b, out_dtype=out_dtype) + else: + return fp_mm_torch(a, b, out_dtype=out_dtype) + + +def fp_mm_func(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor: + if sdnq_triton_mm is not None and a.device.type in {"cuda", "xpu"}: + return sdnq_triton_mm(a, b, out_dtype=out_dtype) + elif openvino_fp_mm is not None and a.device.type == "cpu": + return openvino_fp_mm(a, b, out_dtype=out_dtype) + else: + return fp_mm_torch(a, b, out_dtype=out_dtype) + + +def int_scaled_mm_func(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 sdnq_scaled_mm is not None and a.device.type in {"cuda", "xpu"}: + return sdnq_scaled_mm(a, b, scale_a, scale_b, bias=bias, out_dtype=out_dtype) + else: + return int_scaled_mm_torch(a, b, scale_a, scale_b, bias=bias, out_dtype=out_dtype) + + +def fp8_scaled_mm_func(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 is_fp8_mm_supported and sdnq_scaled_mm is not None and a.device.type in {"cuda", "xpu"}: + return sdnq_scaled_mm(a, b, scale_a, scale_b, bias=bias, out_dtype=out_dtype) + else: + return fp8_scaled_mm_torch(a, b, scale_a, scale_b, bias=bias, out_dtype=out_dtype) + + +def fp_scaled_mm_func(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 sdnq_scaled_mm is not None and a.device.type in {"cuda", "xpu"}: + return sdnq_scaled_mm(a, b, scale_a, scale_b, bias=bias, out_dtype=out_dtype) + else: + return fp_scaled_mm_torch(a, b, scale_a, scale_b, bias=bias, out_dtype=out_dtype) + + +int_mm_torch = compile_func(int_mm_torch) +fp8_mm_torch = compile_func(fp8_mm_torch) +fp_mm_torch = compile_func(fp_mm_torch) +int_scaled_mm_torch = compile_func(int_scaled_mm_torch) +fp8_scaled_mm_torch = compile_func(fp8_scaled_mm_torch) +fp_scaled_mm_torch = compile_func(fp_scaled_mm_torch) diff --git a/modules/sdnq/kernels/openvino_mm.py b/modules/sdnq/kernels/openvino_mm.py index a5ab57411..56ee3c4f9 100644 --- a/modules/sdnq/kernels/openvino_mm.py +++ b/modules/sdnq/kernels/openvino_mm.py @@ -18,8 +18,9 @@ for ov_device in core.get_available_devices(): def ov_mm(infer_request: ov.InferRequest, out_name: str, A: torch.Tensor, B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor: 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)) + A, B = A.contiguous(), B.contiguous() + infer_request.set_tensor("A", ov.Tensor(A.detach().to("cpu").numpy(), shared_memory=True)) + infer_request.set_tensor("B", ov.Tensor(B.detach().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, dtype=out_dtype) diff --git a/modules/sdnq/kernels/triton_mm.py b/modules/sdnq/kernels/triton_mm.py index be49a94fe..7479c14b5 100644 --- a/modules/sdnq/kernels/triton_mm.py +++ b/modules/sdnq/kernels/triton_mm.py @@ -14,7 +14,7 @@ matmul_configs = [ for BK in [int(BK) for BK in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_K_LIST", "32,64,128").replace(" ","").split(",")] for GM in [int(GM) for GM in os.environ.get("SDNQ_TRITON_MM_GROUP_SIZE_M_LIST", "8").replace(" ","").split(",")] for w in [int(w) for w in os.environ.get("SDNQ_TRITON_MM_NUM_WARPS_LIST", "16" if torch.xpu.is_available() else "4").replace(" ","").split(",")] - for s in [int(s) for s in os.environ.get("SDNQ_TRITON_MM_NUM_STAGES_LIST", "1" if (torch.cuda.is_available() and torch.version.hip) else "2").replace(" ","").split(",")] + for s in [int(s) for s in os.environ.get("SDNQ_TRITON_MM_NUM_STAGES_LIST", "1,2" if (torch.cuda.is_available() and torch.version.hip) else "2").replace(" ","").split(",")] ] diff --git a/modules/sdnq/kernels/triton_scaled_mm.py b/modules/sdnq/kernels/triton_scaled_mm.py index c2c1d8689..c546b4ecf 100644 --- a/modules/sdnq/kernels/triton_scaled_mm.py +++ b/modules/sdnq/kernels/triton_scaled_mm.py @@ -14,7 +14,7 @@ matmul_configs = [ for BK in [int(BK) for BK in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_K_LIST", "32,64,128").replace(" ","").split(",")] for GM in [int(GM) for GM in os.environ.get("SDNQ_TRITON_MM_GROUP_SIZE_M_LIST", "8").replace(" ","").split(",")] for w in [int(w) for w in os.environ.get("SDNQ_TRITON_MM_NUM_WARPS_LIST", "16" if torch.xpu.is_available() else "4").replace(" ","").split(",")] - for s in [int(s) for s in os.environ.get("SDNQ_TRITON_MM_NUM_STAGES_LIST", "1" if (torch.cuda.is_available() and torch.version.hip) else "2").replace(" ","").split(",")] + for s in [int(s) for s in os.environ.get("SDNQ_TRITON_MM_NUM_STAGES_LIST", "1,2" if (torch.cuda.is_available() and torch.version.hip) else "2").replace(" ","").split(",")] ]