From 0d15fd1f62f3537579314e95a8319db87a4fe387 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Tue, 14 Jul 2026 17:16:25 +0300 Subject: [PATCH] upcast fp8 mm to fp16 on unsupported gpus and disable triton mm on alchemist and igpus --- modules/sdnq/kernel_wrappers.py | 52 ++++++++++++++++++++------------- 1 file changed, 32 insertions(+), 20 deletions(-) diff --git a/modules/sdnq/kernel_wrappers.py b/modules/sdnq/kernel_wrappers.py index fbdb85f89..f123632cb 100644 --- a/modules/sdnq/kernel_wrappers.py +++ b/modules/sdnq/kernel_wrappers.py @@ -32,13 +32,18 @@ if devices.backend == "rocm": else: is_rdna2_and_older = False +if devices.backend in {"ipex", "xpu"}: + is_alchemist_or_igpu = bool(not torch.xpu.get_device_capability(devices.device).get("has_subgroup_2d_block_io", False)) +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(is_rdna2_and_older or devices.backend in {"zluda", "ipex", "xpu"}) + use_triton_mm = bool(not is_alchemist_or_igpu and (is_rdna2_and_older or devices.backend in {"zluda", "ipex", "xpu"})) else: use_triton_mm = bool(os.environ.get("SDNQ_USE_TRITON_MM", "0").lower() not in {"0", "false", "no"}) @@ -56,6 +61,20 @@ else: use_contiguous_fp16_mm = use_contiguous_int8_mm +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 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) + + int_mm_func = None fp_mm_func = None fp8_mm_func = None @@ -81,11 +100,11 @@ elif use_triton_mm: if is_fp8_mm_supported: fp8_mm_func = sdnq_triton_mm fp8_scaled_mm_func = sdnq_scaled_mm - use_tensorwise_fp8_matmul = False + use_tensorwise_fp8_matmul = True 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"}: +if fp_mm_func is None and not is_alchemist_or_igpu 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 @@ -102,25 +121,18 @@ if int_mm_func is None: 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) + fp_mm_func = fp_mm_torch_cuda 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 + 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 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 + else: + fp8_mm_func = fp_mm_torch if int_scaled_mm_func is None: @@ -140,7 +152,7 @@ if fp_scaled_mm_func is None: fp_scaled_mm_func = fp_scaled_mm_torch if fp8_scaled_mm_func is None: - if use_tensorwise_fp8_matmul: + 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)