Update to use the new attribute setting for tf32. (#835)

Author: choijon5

helion/_testing.py

@@ -458,8 +458,8 @@ def run_example(
         atol: Absolute tolerance for correctness check (default: 1e-1)
         bwd: Whether to also test backward pass (default: False)
     """
-    torch.backends.cuda.matmul.fp32_precision = "tf32"
-    torch.backends.cudnn.conv.fp32_precision = "tf32"  # type: ignore[reportAttributeAccessIssue]
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
 
     # Normalize to dict format
     kernels = kernel_fn if isinstance(kernel_fn, dict) else {kernel_name: kernel_fn}

helion/autotuner/base_cache.py

@@ -65,7 +65,9 @@ def torch_key_wrapper() -> str:
 
 @functools.cache
 def triton_key_wrapper() -> str:
-    from torch._inductor.runtime.triton_compat import triton_key
+    from torch._inductor.runtime.triton_compat import (
+        triton_key,  # pyright: ignore[reportAttributeAccessIssue]
+    )
 
     return triton_key()
 

helion/autotuner/local_cache.py

@@ -55,12 +55,12 @@ def _generate_key(self) -> LooseAutotuneCacheKey:
         for arg in self.args:
             if isinstance(arg, torch.Tensor):
                 nms = torch.xpu if torch.xpu.is_available() else torch.cuda
-                device_properties = nms.get_device_properties(arg.device)
+                device_properties = nms.get_device_properties(arg.device)  # pyright: ignore[reportAttributeAccessIssue]
                 if torch.version.cuda is not None:  # pyright: ignore[reportAttributeAccessIssue]
                     hardware = device_properties.name
-                    runtime_name = str(torch.version.cuda)
+                    runtime_name = str(torch.version.cuda)  # pyright: ignore[reportAttributeAccessIssue]
                 elif torch.version.hip is not None:  # pyright: ignore[reportAttributeAccessIssue]
-                    hardware = device_properties.gcnArchName
+                    hardware = device_properties.gcnArchName  # pyright: ignore[reportAttributeAccessIssue]
                     runtime_name = torch.version.hip  # pyright: ignore[reportAttributeAccessIssue]
                 else:
                     hardware = device_properties.name