Fix compatibility with torch 2.7 (#45)

Pull Request resolved: #45
Approved by: https://github.com/oulgen, https://github.com/yf225

Author: jansel

examples/bmm.py

@@ -40,4 +40,6 @@ def check(b: int, m: int, k: int, n: int) -> None:
 
 
 if __name__ == "__main__":
+    # torch.baddbmm support for 16-bit tensors requires torch 2.8+
+    assert torch.__version__.split(".")[:2] >= ["2", "8"], "Requires torch 2.8+"
     check(16, 512, 768, 1024)

helion/runtime/kernel.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable
+import contextlib
 import dataclasses
 import functools
 import inspect
@@ -247,9 +248,7 @@ def __init__(self, kernel: Kernel, args: tuple[object, ...]) -> None:
                     constexpr_args[name] = arg
                 else:
                     self.fake_args.append(self.env.to_fake(arg, ArgumentOrigin(name)))
-            with torch.fx.experimental._config.patch(  # pyre-ignore[16]
-                skip_dtype_check_in_meta_registrations=True
-            ):
+            with _maybe_skip_dtype_check_in_meta_registrations():
                 self.host_fn: HostFunction = HostFunction(
                     self.kernel.fn, self.fake_args, constexpr_args
                 )
@@ -542,3 +541,13 @@ def _find_device(args: tuple[object, ...]) -> torch.device:
                 except exc.NoTensorArgs:
                     pass
     raise exc.NoTensorArgs
+
+
+def _maybe_skip_dtype_check_in_meta_registrations() -> (
+    contextlib.AbstractContextManager[None, None]
+):
+    if hasattr(torch.fx.experimental._config, "skip_dtype_check_in_meta_registrations"):
+        return torch.fx.experimental._config.patch(  # pyre-ignore[16]
+            skip_dtype_check_in_meta_registrations=True
+        )
+    return contextlib.nullcontext()

test/test_examples.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
 from pathlib import Path
+import unittest
 
 from expecttest import TestCase
+from packaging import version
 import torch
 
 from helion._testing import DEVICE
@@ -141,6 +143,10 @@ def _matmul_make_precompiler(x: torch.Tensor, y: torch.Tensor):
     return make_precompiler(_matmul_kernel)(x, y, out, _BLOCK_SIZE_0, _BLOCK_SIZE_1, _BLOCK_SIZE_2, num_warps=4, num_stages=3)""",
         )
 
+    @unittest.skipIf(
+        version.parse(torch.__version__.split("+")[0]) < version.parse("2.8"),
+        "Requires torch 2.8+",
+    )
     def test_bmm(self):
         args = (
             torch.randn([16, 512, 768], device=DEVICE, dtype=torch.float16),