FIX: tl.program_id() does indeed not have a cast method in triton2.3.1

wizyoung · Sep 26, 2024 · cd15004 · cd15004
1 parent 5c38d09
commit cd15004
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/liger_kernel/ops/geglu.py b/src/liger_kernel/ops/geglu.py
@@ -25,7 +25,7 @@
 def _geglu_tanh_forward_kernel(
     a, b, c, stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr
 ):
-    program_id = tl.program_id(0).cast(tl.int64)
+    program_id = tl.program_id(0).to(tl.int64)
 
     # locate start index
     a += program_id * stride
@@ -52,7 +52,7 @@ def _geglu_tanh_forward_kernel(
 def _geglu_tanh_backward_kernel(
     dc, a, b, stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr
 ):
-    program_id = tl.program_id(0).cast(tl.int64)
+    program_id = tl.program_id(0).to(tl.int64)
 
     # locate start index
     dc += program_id * stride

diff --git a/src/liger_kernel/ops/swiglu.py b/src/liger_kernel/ops/swiglu.py
@@ -14,7 +14,7 @@ def silu(x):
 def _swiglu_forward_kernel(
     a_ptr, b_ptr, c_ptr, stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr
 ):
-    program_id = tl.program_id(0).cast(tl.int64)
+    program_id = tl.program_id(0).to(tl.int64)
 
     # locate start index
     a_ptr += program_id * stride
@@ -35,7 +35,7 @@ def _swiglu_forward_kernel(
 def _swiglu_backward_kernel(
     dc_ptr, a_ptr, b_ptr, stride, n_cols: tl.constexpr, BLOCK_SIZE: tl.constexpr
 ):
-    program_id = tl.program_id(0).cast(tl.int64)
+    program_id = tl.program_id(0).to(tl.int64)
 
     # locate start index
     dc_ptr += program_id * stride