mirror of
https://github.com/vladmandic/automatic
synced 2026-09-20 01:31:13 +02:00
Add device checks on kernel_wrappers and support OpenVINO MM for CPU and Triton MM for GPU at the same time
This commit is contained in:
Dity0
+115
-95
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(",")]
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -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(",")]
|
||||
]
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user