From 13c1bcff3f7c9427bcb6fe5110867061222fa16b Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sat, 27 Jun 2026 01:40:11 +0300 Subject: [PATCH] cleanup --- modules/sdnq/kernels/triton_atten.py | 80 ++++++++++++++-------------- 1 file changed, 40 insertions(+), 40 deletions(-) diff --git a/modules/sdnq/kernels/triton_atten.py b/modules/sdnq/kernels/triton_atten.py index 0198615a2..a1757bb0b 100644 --- a/modules/sdnq/kernels/triton_atten.py +++ b/modules/sdnq/kernels/triton_atten.py @@ -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(