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

cleanup

Browse Source
This commit is contained in:
Disty0
2026-06-27 01:40:11 +03:00
parent b92ccf4adf
commit 13c1bcff3f
1 changed files with 40 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/sdnq/kernels/triton_atten.py
+40 -40
View File
@@ -10,7 +10,7 @@ from ..quant_utils import quantize_int_mm, quantize_fp_mm, get_hadamard, apply_h
min_block_size = int(os.environ.get("SDNQ_TRITON_ATTEN_MIN_BLOCK_SIZE", "32"))
matmul_configs = [
triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_warps=w, num_stages=s)
triton.Config({'BLOCK_SIZE_M': BM, 'BLOCK_SIZE_N': BN}, num_warps=w, num_stages=s)
for BM in [int(BM) for BM in os.environ.get("SDNQ_TRITON_ATTEN_BLOCK_SIZE_M_LIST", "64,128").replace(" ","").split(",")]
for BN in [int(BN) for BN in os.environ.get("SDNQ_TRITON_ATTEN_BLOCK_SIZE_N_LIST", "32,64").replace(" ","").split(",")]
for w in [int(w) for w in os.environ.get("SDNQ_TRITON_ATTEN_NUM_WARPS_LIST", "2,4,8").replace(" ","").split(",")]
@@ -62,15 +62,15 @@ def quantize_attn(
return q_q, q_scale, k_q, k_scale, v_q, v_scale
@triton.autotune(configs=matmul_configs, key=["qn_t", "qhd","vn_t", "vhd", "qk_is_quantized", "pv_is_quantized", "q_dtype", "v_dtype", "out_dtype", "mask_dtype"], cache_results=True)
@triton.autotune(configs=matmul_configs, key=["QN_AT", "QHD", "VN_AT", "VHD", "qk_is_quantized", "pv_is_quantized", "q_dtype", "v_dtype", "out_dtype", "mask_dtype"], cache_results=True)
@triton.jit
def sdnq_attn_kernel(
q_ptr, k_ptr, v_ptr, q_scale_ptr, k_scale_ptr, v_scale_ptr, out_ptr, mask_ptr,
qz: tl.constexpr, qh: tl.constexpr, qn: tl.constexpr, qhd: tl.constexpr,
kz: tl.constexpr, kh: tl.constexpr, kn: tl.constexpr, khd: tl.constexpr,
vz: tl.constexpr, vh: tl.constexpr, vn: tl.constexpr, vhd: tl.constexpr,
oz: tl.constexpr, oh: tl.constexpr, on: tl.constexpr, ohd: tl.constexpr,
mz: tl.constexpr, mh: tl.constexpr, mqn: tl.constexpr, mkn: tl.constexpr,
QZ: tl.constexpr, QH: tl.constexpr, QN: tl.constexpr, QHD: tl.constexpr,
KZ: tl.constexpr, KH: tl.constexpr, KN: tl.constexpr, KHD: tl.constexpr,
VZ: tl.constexpr, VH: tl.constexpr, VN: tl.constexpr, VHD: tl.constexpr,
OZ: tl.constexpr, OH: tl.constexpr, ON: tl.constexpr, OHD: tl.constexpr,
MZ: tl.constexpr, MH: tl.constexpr, MQN: tl.constexpr, MKN: tl.constexpr,
stride_qz: tl.constexpr, stride_qh: tl.constexpr, stride_qn: tl.constexpr, stride_qhd: tl.constexpr,
stride_kz: tl.constexpr, stride_kh: tl.constexpr, stride_kn: tl.constexpr, stride_khd: tl.constexpr,
stride_vz: tl.constexpr, stride_vh: tl.constexpr, stride_vn: tl.constexpr, stride_vhd: tl.constexpr,
@@ -79,43 +79,43 @@ def sdnq_attn_kernel(
stride_skz: tl.constexpr, stride_skh: tl.constexpr, stride_skn: tl.constexpr,
stride_svz: tl.constexpr, stride_svh: tl.constexpr, stride_svn: tl.constexpr,
stride_mz: tl.constexpr, stride_mh: tl.constexpr, stride_mqn: tl.constexpr, stride_mkn: tl.constexpr,
qn_t: tl.constexpr, vn_t: tl.constexpr,
QN_AT: tl.constexpr, VN_AT: tl.constexpr,
qk_is_quantized: tl.constexpr, pv_is_quantized: tl.constexpr,
q_dtype: tl.constexpr, v_dtype: tl.constexpr,
out_dtype: tl.constexpr, mask_dtype: tl.constexpr,
BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
): # pylint: disable=unused-argument
start_m = tl.program_id(0)
off_z = tl.program_id(2)
off_h = tl.program_id(1)
num_kv_groups = qh // vh
num_kv_groups = QH // VH
off_h_kv = off_h // num_kv_groups
m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0
acc = tl.zeros([BLOCK_M, vhd], dtype=tl.float32)
m_i = tl.zeros([BLOCK_SIZE_M], dtype=tl.float32) - float("inf")
l_i = tl.zeros([BLOCK_SIZE_M], dtype=tl.float32) + 1.0
acc = tl.zeros([BLOCK_SIZE_M, VHD], dtype=tl.float32)
Q_desc = tl.make_tensor_descriptor(q_ptr + off_z * stride_qz + off_h * stride_qh, shape=[qn, qhd], strides=[stride_qn, stride_qhd], block_shape=[BLOCK_M, qhd])
K_desc = tl.make_tensor_descriptor(k_ptr + off_z * stride_kz + off_h_kv * stride_kh, shape=[kn, khd], strides=[stride_kn, stride_khd], block_shape=[BLOCK_N, khd])
V_desc = tl.make_tensor_descriptor(v_ptr + off_z * stride_vz + off_h_kv * stride_vh, shape=[vn, vhd], strides=[stride_vn, stride_vhd], block_shape=[BLOCK_N, vhd])
Q_desc = tl.make_tensor_descriptor(q_ptr + off_z * stride_qz + off_h * stride_qh, shape=[QN, QHD], strides=[stride_qn, stride_qhd], block_shape=[BLOCK_SIZE_M, QHD])
K_desc = tl.make_tensor_descriptor(k_ptr + off_z * stride_kz + off_h_kv * stride_kh, shape=[KN, KHD], strides=[stride_kn, stride_khd], block_shape=[BLOCK_SIZE_N, KHD])
V_desc = tl.make_tensor_descriptor(v_ptr + off_z * stride_vz + off_h_kv * stride_vh, shape=[VN, VHD], strides=[stride_vn, stride_vhd], block_shape=[BLOCK_SIZE_N, VHD])
if q_scale_ptr is not None:
q_scale_desc = tl.make_tensor_descriptor(q_scale_ptr + off_z * stride_sqz + off_h * stride_sqh, shape=[qn], strides=[stride_sqn], block_shape=[BLOCK_M])
k_scale_desc = tl.make_tensor_descriptor(k_scale_ptr + off_z * stride_skz + off_h_kv * stride_skh, shape=[kn], strides=[stride_skn], block_shape=[BLOCK_N])
q_scale = q_scale_desc.load([start_m * BLOCK_M])[:, None]
q_scale_desc = tl.make_tensor_descriptor(q_scale_ptr + off_z * stride_sqz + off_h * stride_sqh, shape=[QN], strides=[stride_sqn], block_shape=[BLOCK_SIZE_M])
k_scale_desc = tl.make_tensor_descriptor(k_scale_ptr + off_z * stride_skz + off_h_kv * stride_skh, shape=[KN], strides=[stride_skn], block_shape=[BLOCK_SIZE_N])
q_scale = q_scale_desc.load([start_m * BLOCK_SIZE_M])[:, None]
if v_scale_ptr is not None:
v_scale_desc = tl.make_tensor_descriptor(v_scale_ptr + off_z * stride_svz + off_h_kv * stride_svh, shape=[vn], strides=[stride_svn], block_shape=[BLOCK_N])
v_scale_desc = tl.make_tensor_descriptor(v_scale_ptr + off_z * stride_svz + off_h_kv * stride_svh, shape=[VN], strides=[stride_svn], block_shape=[BLOCK_SIZE_N])
if mask_ptr is not None:
mask_desc = tl.make_tensor_descriptor(mask_ptr + (off_z * stride_mz + off_h * stride_mh), shape=[mqn, mkn], strides=[stride_mqn, stride_mkn], block_shape=[BLOCK_M, BLOCK_N])
mask_desc = tl.make_tensor_descriptor(mask_ptr + (off_z * stride_mz + off_h * stride_mh), shape=[MQN, MKN], strides=[stride_mqn, stride_mkn], block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_N])
q = Q_desc.load([start_m * BLOCK_M, 0])
q = Q_desc.load([start_m * BLOCK_SIZE_M, 0])
lo, hi = 0, kn
for start_n in range(lo, hi, BLOCK_N):
start_n = tl.multiple_of(start_n, BLOCK_N)
lo, hi = 0, KN
for start_n in range(lo, hi, BLOCK_SIZE_N):
start_n = tl.multiple_of(start_n, BLOCK_SIZE_N)
skip = False
if mask_ptr is not None:
mask = mask_desc.load([start_m * BLOCK_M, start_n])
mask = mask_desc.load([start_m * BLOCK_SIZE_M, start_n])
if mask.dtype == tl.int8:
mask = mask.to(tl.int1)
if mask.dtype == tl.int1 and tl.max(mask) == 0:
@@ -169,8 +169,8 @@ def sdnq_attn_kernel(
m_i = m_ij
acc = (acc / l_i[:, None]).to(out_ptr.type.element_ty)
O_desc = tl.make_tensor_descriptor(out_ptr + off_z * stride_oz + off_h * stride_oh, shape=[on, ohd], strides=[stride_on, stride_ohd], block_shape=[BLOCK_M, ohd])
O_desc.store([start_m * BLOCK_M, 0], acc)
O_desc = tl.make_tensor_descriptor(out_ptr + off_z * stride_oz + off_h * stride_oh, shape=[QN, QHD], strides=[stride_on, stride_ohd], block_shape=[BLOCK_SIZE_M, QHD])
O_desc.store([start_m * BLOCK_SIZE_M, 0], acc)
def sdnq_triton_atten_forward(
@@ -193,24 +193,24 @@ def sdnq_triton_atten_forward(
assert not is_causal
if out_dtype is None:
out_dtype = query.dtype
qz, qh, qn, qhd = query.shape
if not math.log(qhd, 2).is_integer():
head_dim_pow2 = triton.next_power_of_2(qhd)
query = torch.nn.functional.pad(query, (0, head_dim_pow2 - qhd))
key = torch.nn.functional.pad(key, (0, head_dim_pow2 - qhd))
value = torch.nn.functional.pad(value, (0, head_dim_pow2 - qhd))
QZ, QH, QN, QHD = query.shape
if not math.log(QHD, 2).is_integer():
head_dim_pow2 = triton.next_power_of_2(QHD)
query = torch.nn.functional.pad(query, (0, head_dim_pow2 - QHD))
key = torch.nn.functional.pad(key, (0, head_dim_pow2 - QHD))
value = torch.nn.functional.pad(value, (0, head_dim_pow2 - QHD))
if scale is None:
scale = qhd ** -0.5
scale = QHD ** -0.5
if attn_mask is not None:
attn_mask = attn_mask.expand((qz, qh, qn, key.shape[-2])).contiguous()
attn_mask = attn_mask.expand((QZ, QH, QN, key.shape[-2])).contiguous()
if not math.log(key.shape[-2], 2).is_integer():
pad_value = -float('inf') if torch.is_floating_point(attn_mask) else 0
attn_mask = torch.nn.functional.pad(attn_mask, (0, triton.next_power_of_2(key.shape[-2]) - key.shape[-2]), value=pad_value)
if attn_mask.dtype == torch.bool:
attn_mask = attn_mask.to(dtype=torch.int8)
def grid(META):
return (triton.cdiv(qn, META["BLOCK_M"]), qh, qz)
out = torch.empty((qz, qh, qn, value.shape[-1]), dtype=out_dtype, device=query.device)
return (triton.cdiv(QN, META["BLOCK_SIZE_M"]), QH, QZ)
out = torch.empty((QZ, QH, QN, value.shape[-1]), dtype=out_dtype, device=query.device)
query, query_scale, key, key_scale, value, value_scale = quantize_attn(
query, key, value,
hadamard=hadamard,
@@ -228,12 +228,12 @@ def sdnq_triton_atten_forward(
*(key_scale.stride() if key_scale is not None else (0, 0, 0)),
*(value_scale.stride() if value_scale is not None else (0, 0, 0)),
*(attn_mask.stride() if attn_mask is not None else (0, 0, 0, 0)),
math.ceil(qn / min_block_size), math.ceil(value.shape[-2] / min_block_size),
math.ceil(QN / min_block_size), math.ceil(value.shape[-2] / min_block_size),
bool(query_scale is not None), bool(value_scale is not None),
str(query.dtype), str(value.dtype), str(out.dtype),
str(attn_mask.dtype if attn_mask is not None else None),
)
return out[..., :qhd]
return out[..., :QHD]
def sdnq_triton_atten(