diff --git a/modules/sdnq/kernel_wrappers.py b/modules/sdnq/kernel_wrappers.py index f123632cb..9388fd0ae 100644 --- a/modules/sdnq/kernel_wrappers.py +++ b/modules/sdnq/kernel_wrappers.py @@ -82,6 +82,7 @@ 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 @@ -104,6 +105,7 @@ elif use_triton_mm: except Exception: use_triton_mm = False + 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 @@ -119,12 +121,14 @@ if int_mm_func is None: 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 else: fp_mm_func = fp_mm_torch + if fp8_mm_func is None: if is_fp8_mm_supported: def fp8_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor: @@ -138,19 +142,21 @@ if fp8_mm_func is None: 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) + return int_mm_func(a,b, out_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) + 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 = 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) + 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).to(dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype) + 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 = 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: diff --git a/modules/sdnq/kernels/triton_atten.py b/modules/sdnq/kernels/triton_atten.py index 5fd7d9fd6..761fb7ce6 100644 --- a/modules/sdnq/kernels/triton_atten.py +++ b/modules/sdnq/kernels/triton_atten.py @@ -139,9 +139,9 @@ def sdnq_attn_kernel( if qk_is_quantized: k_scale = k_scale_desc.load([start_n])[None, :] if q.dtype == tl.int8: - qk = tl.dot(q, k, out_dtype=tl.int32).to(tl.float32) * q_scale * k_scale + qk = tl.mul(tl.mul(tl.dot(q, k, out_dtype=tl.int32).to(tl.float32), q_scale), k_scale) else: - qk = tl.dot(q, k, out_dtype=tl.float32) * q_scale * k_scale + qk = tl.mul(tl.mul(tl.dot(q, k, out_dtype=tl.float32), q_scale), k_scale) else: qk = tl.dot(q, k, out_dtype=tl.float32) @@ -171,14 +171,14 @@ def sdnq_attn_kernel( v_scale = v_scale_desc.load([start_n])[None, :] p *= v_scale if v.dtype == tl.int8: - p_scale = tl.max(p, 1)[:, None] * (1 / 127.0) + p_scale = tl.mul(tl.max(p, 1)[:, None], (1 / 127.0)) p_scale = tl.where(p_scale <= 2e-38, 1.0, p_scale) - p = tl.floor(p * (1 / p_scale) + 0.5).to(tl.int8) + p = tl.floor(tl.fma(p, (1 / p_scale), 0.5)).to(tl.int8) acc = tl.fma(tl.dot(p, v, out_dtype=tl.int32).to(tl.float32), p_scale, acc) else: - p_scale = tl.max(p, 1)[:, None] * (1 / (65504.0 if v.dtype == tl.float16 else 448.0)) + p_scale = tl.mul(tl.max(p, 1)[:, None], (1 / (65504.0 if v.dtype == tl.float16 else 448.0))) p_scale = tl.where(p_scale <= 2e-38, 1.0, p_scale) - p = (p * (1 / p_scale)).to(v.dtype) + p = tl.mul(p, (1 / p_scale)).to(v.dtype) acc = tl.fma(tl.dot(p, v, out_dtype=tl.float32), p_scale, acc) else: p = p.to(v.dtype) diff --git a/modules/sdnq/kernels/triton_scaled_mm.py b/modules/sdnq/kernels/triton_scaled_mm.py index 18779e74d..d222b2fba 100644 --- a/modules/sdnq/kernels/triton_scaled_mm.py +++ b/modules/sdnq/kernels/triton_scaled_mm.py @@ -1,8 +1,3 @@ -""" -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 @@ -96,17 +91,17 @@ def sdnq_scaled_mm_kernel( scale_b = scale_b_desc.load([off_n])[None, :].to(tl.float32) if bias_ndim == 1: - accumulator = accumulator.to(tl.float32) * scale_a + accumulator = tl.mul(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 + accumulator = tl.mul(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 = tl.mul(tl.mul(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)) diff --git a/modules/sdnq/layers/linear/linear_int8.py b/modules/sdnq/layers/linear/linear_int8.py index 9f39d624b..f6c3ee803 100644 --- a/modules/sdnq/layers/linear/linear_int8.py +++ b/modules/sdnq/layers/linear/linear_int8.py @@ -60,7 +60,7 @@ def int8_matmul( input, input_scale = quantize_int_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 = torch.sum(input, dim=-1, keepdim=True, dtype=torch.int32).to(dtype=input_scale.dtype).mul_(input_scale).mul(zero_point) if bias is not None: zero_bias.add_(bias) bias = zero_bias diff --git a/modules/sdnq/layers/linear/linear_uint8.py b/modules/sdnq/layers/linear/linear_uint8.py index 7812a8759..74ebdc58b 100644 --- a/modules/sdnq/layers/linear/linear_uint8.py +++ b/modules/sdnq/layers/linear/linear_uint8.py @@ -63,7 +63,7 @@ def uint8_matmul( if zero_point is not None: 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)) + zero_bias.add_(torch.mul(input_zero_point, zero_point), alpha=input.shape[-1]) else: 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: