Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 17:24:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

cleanup sdnq

Browse Source
This commit is contained in:
Disty0
2026-07-14 20:55:47 +03:00
parent b1c3578b27
commit 47cb113b75
5 changed files with 21 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/sdnq/kernel_wrappers.py
+10 -4
View File
@@ -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:
modules/sdnq/kernels/triton_atten.py
+6 -6
View File
@@ -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)
modules/sdnq/kernels/triton_scaled_mm.py
+3 -8
View File
@@ -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))
modules/sdnq/layers/linear/linear_int8.py
+1 -1
View File
@@ -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
modules/sdnq/layers/linear/linear_uint8.py
+1 -1
View File
@@ -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: