From f0ef2018822f990b78712d5da35c757a6a40ed94 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 28 Jun 2026 22:16:53 +0300 Subject: [PATCH] cleanup --- modules/sdnq/kernels/triton_atten.py | 8 ++++---- modules/sdnq/kernels/triton_mm.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/modules/sdnq/kernels/triton_atten.py b/modules/sdnq/kernels/triton_atten.py index 72a3afb59..fa7c283d5 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_SIZE_M': BM, 'BLOCK_SIZE_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").replace(" ","").split(",")] for BN in [int(BN) for BN in os.environ.get("SDNQ_TRITON_ATTEN_BLOCK_SIZE_N_LIST", "32").replace(" ","").split(",")] for w in [int(w) for w in os.environ.get("SDNQ_TRITON_ATTEN_NUM_WARPS_LIST", "2,4,8").replace(" ","").split(",")] @@ -142,11 +142,11 @@ def sdnq_attn_kernel( if is_causal: causal_mask = offs_m[:, None] >= (start_n + offs_n[None, :]) - qk = tl.where(causal_mask, qk, -float('inf')) + qk = tl.where(causal_mask, qk, float("-inf")) if mask_ptr is not None: if mask.dtype == tl.int1: - qk = tl.where(mask, qk, -float('inf')) + qk = tl.where(mask, qk, float("-inf")) else: qk = qk + mask @@ -217,7 +217,7 @@ def sdnq_triton_atten_forward( if attn_mask is not None: attn_mask = attn_mask.expand((QZ, QH, QN, KN)).contiguous() if not math.log(KN, 2).is_integer(): - pad_value = -float('inf') if torch.is_floating_point(attn_mask) else 0 + 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(KN) - KN), value=pad_value) if attn_mask.dtype == torch.bool: attn_mask = attn_mask.to(dtype=torch.int8) diff --git a/modules/sdnq/kernels/triton_mm.py b/modules/sdnq/kernels/triton_mm.py index 086933363..e43073b64 100644 --- a/modules/sdnq/kernels/triton_mm.py +++ b/modules/sdnq/kernels/triton_mm.py @@ -17,7 +17,7 @@ import triton.language as tl min_block_size = int(os.environ.get("SDNQ_TRITON_MM_MIN_BLOCK_SIZE", "64")) matmul_configs = [ - triton.Config({'BLOCK_SIZE_M': BM, 'BLOCK_SIZE_N': BN, "BLOCK_SIZE_K": BK, "GROUP_SIZE_M": GM}, num_warps=w, num_stages=s) + triton.Config({"BLOCK_SIZE_M": BM, "BLOCK_SIZE_N": BN, "BLOCK_SIZE_K": BK, "GROUP_SIZE_M": GM}, num_warps=w, num_stages=s) for BM in [int(BM) for BM in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_M_LIST", "64,128,256").replace(" ","").split(",")] for BN in [int(BN) for BN in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_N_LIST", "64,128,256").replace(" ","").split(",")] for BK in [int(BK) for BK in os.environ.get("SDNQ_TRITON_MM_BLOCK_SIZE_K_LIST", "64,128").replace(" ","").split(",")]