Inductor breaks down with math.log

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:

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
 

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: