From eeb34eb6c93e5be937e591c2ada6267ce856bb57 Mon Sep 17 00:00:00 2001 From: Will Feng Date: Tue, 11 Nov 2025 14:13:06 -0800 Subject: [PATCH 1/2] Remove @helion.jit and advise use of @helion.kernel --- examples/all_gather_matmul.py | 2 +- examples/all_reduce.py | 2 +- helion/__init__.py | 2 -- test/test_associative_scan.py | 12 ++++++------ test/test_reduce.py | 6 +++--- 5 files changed, 11 insertions(+), 13 deletions(-) diff --git a/examples/all_gather_matmul.py b/examples/all_gather_matmul.py index 001d67a3c..1f4a0445c 100644 --- a/examples/all_gather_matmul.py +++ b/examples/all_gather_matmul.py @@ -79,7 +79,7 @@ def copy_engine_all_gather_w_progress( # %% -@helion.jit( +@helion.kernel( config=helion.Config( block_sizes=[128, 256, 64], num_warps=8, diff --git a/examples/all_reduce.py b/examples/all_reduce.py index 17616c275..7396e9b9c 100644 --- a/examples/all_reduce.py +++ b/examples/all_reduce.py @@ -81,7 +81,7 @@ def dev_array_to_tensor_short( # %% -@helion.jit( +@helion.kernel( config=helion.Config( block_sizes=[8192], num_warps=32, diff --git a/helion/__init__.py b/helion/__init__.py index b02e42e47..b1f7c8341 100644 --- a/helion/__init__.py +++ b/helion/__init__.py @@ -11,7 +11,6 @@ from .runtime import Config from .runtime import Kernel from .runtime import kernel -from .runtime import kernel as jit # alias from .runtime.settings import RefMode from .runtime.settings import Settings @@ -22,7 +21,6 @@ "Settings", "cdiv", "exc", - "jit", "kernel", "language", "next_power_of_2", diff --git a/test/test_associative_scan.py b/test/test_associative_scan.py index f5c070d79..030def297 100644 --- a/test/test_associative_scan.py +++ b/test/test_associative_scan.py @@ -92,9 +92,9 @@ def cumsum_helper(x: torch.Tensor) -> torch.Tensor: return hl.associative_scan(add_combine_fn, x, dim=0) -@helion.jit +@helion.kernel def jit_add_combine_fn(x, y): - """Addition combine function with @helion.jit decorator (should be ignored).""" + """Addition combine function with @helion.kernel decorator (should be ignored).""" return x + y @@ -496,10 +496,10 @@ def test_codegen_kernel(x: torch.Tensor) -> torch.Tensor: self.assertNotIn("placeholder", code) @skipIfRefEager( - "torch._higher_order_ops.associative_scan with nested @helion.jit is not supported by ref eager mode yet" + "torch._higher_order_ops.associative_scan with nested @helion.kernel is not supported by ref eager mode yet" ) def test_associative_scan_jit_decorator_ignored(self): - """Test that @helion.jit decorator on combine functions is ignored.""" + """Test that @helion.kernel decorator on combine functions is ignored.""" @helion.kernel(autotune_effort="none") def test_jit_kernel(x: torch.Tensor) -> torch.Tensor: @@ -521,8 +521,8 @@ def test_jit_kernel(x: torch.Tensor) -> torch.Tensor: self.assertIn("def jit_add_combine_fn_", code) self.assertIn("tl.associative_scan", code) self.assertIn("param_0 + param_1", code) - # Verify @helion.jit decorator doesn't appear in generated code - self.assertNotIn("@helion.jit", code) + # Verify @helion.kernel decorator doesn't appear in generated code + self.assertNotIn("@helion.kernel", code) @skipIfRefEager( "torch._higher_order_ops.associative_scan with tuple arg is not supported by ref eager mode yet" diff --git a/test/test_reduce.py b/test/test_reduce.py index 5895d442c..224d1cf6f 100644 --- a/test/test_reduce.py +++ b/test/test_reduce.py @@ -73,9 +73,9 @@ def argmax_combine_unpacked_fn(left_value, left_index, right_value, right_index) return max_value, max_index -@helion.jit +@helion.kernel def jit_add_combine_fn(x, y): - """Addition combine function with @helion.jit decorator (should be ignored).""" + """Addition combine function with @helion.kernel decorator (should be ignored).""" return x + y @@ -221,7 +221,7 @@ def test_reduce_product_kernel(x: torch.Tensor) -> torch.Tensor: torch.testing.assert_close(result, expected) def test_reduce_jit_combine_fn(self): - """Test reduce with @helion.jit decorated combine function.""" + """Test reduce with @helion.kernel decorated combine function.""" @helion.kernel(autotune_effort="none") def test_reduce_jit_kernel(x: torch.Tensor) -> torch.Tensor: From e8f11ef5ec6106ef0df7e569353de0d6e73534a3 Mon Sep 17 00:00:00 2001 From: Will Feng Date: Tue, 11 Nov 2025 23:04:52 -0800 Subject: [PATCH 2/2] keep helion.jit for BC --- helion/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/helion/__init__.py b/helion/__init__.py index b1f7c8341..b02e42e47 100644 --- a/helion/__init__.py +++ b/helion/__init__.py @@ -11,6 +11,7 @@ from .runtime import Config from .runtime import Kernel from .runtime import kernel +from .runtime import kernel as jit # alias from .runtime.settings import RefMode from .runtime.settings import Settings @@ -21,6 +22,7 @@ "Settings", "cdiv", "exc", + "jit", "kernel", "language", "next_power_of_2",