diff --git a/modules/sdnq/kernels/triton_atten.py b/modules/sdnq/kernels/triton_atten.py index 7acb55931..0b80e92dd 100644 --- a/modules/sdnq/kernels/triton_atten.py +++ b/modules/sdnq/kernels/triton_atten.py @@ -6,6 +6,7 @@ import triton.language as tl from ..common import compile_func # pylint: disable=relative-beyond-top-level from ..quant_utils import quantize_int_mm, quantize_fp_mm, get_hadamard, get_hadamard_group_size, apply_hadamard # pylint: disable=relative-beyond-top-level +from ..utils import is_pow2, next_power_of_2 # pylint: disable=relative-beyond-top-level min_block_size = int(os.environ.get("SDNQ_TRITON_ATTEN_MIN_BLOCK_SIZE", "256")) @@ -261,15 +262,15 @@ def get_attn_inputs( out_dtype = query.dtype if scale is None: scale = QHD ** -0.5 - if not math.log2(QHD).is_integer(): - query = torch.nn.functional.pad(query, (0, triton.next_power_of_2(QHD) - QHD)) - key = torch.nn.functional.pad(key, (0, triton.next_power_of_2(KHD) - KHD)) - value = torch.nn.functional.pad(value, (0, triton.next_power_of_2(VHD) - VHD)) + if not is_pow2(QHD): + query = torch.nn.functional.pad(query, (0, next_power_of_2(QHD) - QHD)) + key = torch.nn.functional.pad(key, (0, next_power_of_2(KHD) - KHD)) + value = torch.nn.functional.pad(value, (0, next_power_of_2(VHD) - VHD)) if attn_mask is not None: attn_mask = attn_mask.expand((QZ, QH, QN, KN)) - if not math.log2(KN).is_integer(): + if not is_pow2(KN): 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) + attn_mask = torch.nn.functional.pad(attn_mask, (0, next_power_of_2(KN) - KN), value=pad_value) if attn_mask.dtype == torch.bool: attn_mask = attn_mask.to(dtype=torch.int8) attn_mask = attn_mask.contiguous() @@ -308,7 +309,7 @@ def sdnq_triton_atten( hadamard = None if use_hadamard and do_quantize and matmul_dtype not in {None, "none", "no"}: - hadamard_channel_size = min(triton.next_power_of_2(QHD), triton.next_power_of_2(KHD)) + hadamard_channel_size = next_power_of_2(min(QHD, KHD)) hadamard_group_size = min(hadamard_group_size, hadamard_channel_size) use_hadamard, hadamard_group_size = get_hadamard_group_size(hadamard_channel_size, hadamard_group_size) if use_hadamard: diff --git a/modules/sdnq/quant_utils.py b/modules/sdnq/quant_utils.py index 4765c5f66..2bfacbcd9 100644 --- a/modules/sdnq/quant_utils.py +++ b/modules/sdnq/quant_utils.py @@ -5,6 +5,7 @@ import torch from modules import devices from .common import dtype_dict, use_contiguous_int8_mm, use_contiguous_fp16_mm, conv_types, conv_transpose_types +from .utils import is_pow2, is_pow4, next_power_of_2 @devices.inference_context() @@ -100,9 +101,9 @@ def build_hadamard_n4(n: int, dtype: torch.dtype | None = None, device: torch.de @devices.inference_context() def build_hadamard(n: int, dtype: torch.dtype | None = None, device: torch.device | None = None) -> torch.FloatTensor: - if math.log(n, 4).is_integer(): + if is_pow4(n): return build_hadamard_n4(n, device=device, dtype=dtype) - elif math.log2(n).is_integer(): + elif is_pow2(n): return build_hadamard_n2(n, device=device, dtype=dtype) else: raise RuntimeError(f"Hadamard Group Size must be a power of 2 but got {n}.") @@ -141,12 +142,10 @@ def rotate_hadamard(weight: torch.Tensor, group_size: int = 256, hadamard: torch def get_hadamard_group_size(channel_size: int, group_size: int) -> tuple[bool, int]: - group_size = 2 ** int(math.log2(min(channel_size, group_size))) + group_size = next_power_of_2(min(channel_size, group_size)) if channel_size % group_size != 0: - hadamard_pow2 = int(math.log2(group_size)) while channel_size % group_size != 0: - hadamard_pow2 -= 1 - group_size = 2 ** hadamard_pow2 + group_size = group_size // 2 use_hadamard = group_size >= 4 return use_hadamard, group_size diff --git a/modules/sdnq/utils.py b/modules/sdnq/utils.py index 25bc0f60e..cc423432f 100644 --- a/modules/sdnq/utils.py +++ b/modules/sdnq/utils.py @@ -12,6 +12,20 @@ from .common import ( ) +def is_pow2(n: int) -> bool: + return (n & (n - 1)) == 0 + + +def is_pow4(n: int) -> bool: + return is_pow2(n) and (n.bit_length() & 1 == 1) + + +def next_power_of_2(n: int) -> int: + if is_pow2(n): + return n + return 2 ** n.bit_length() + + def check_param_name_in(param_name: str, param_list: list[str]) -> str: split_param_name = param_name.split(".") for param in param_list: