Cast tl.program_id to int64 in llama4_rope and qwen2vl_mrope

src/liger_kernel/ops/backends/_ascend/ops/qwen2vl_mrope.py

@@ -34,7 +34,7 @@ def _triton_qwen2vl_mrope_npu(
     actual_rows = tl.minimum(rows_per_program, total_rows - start_row)
 
     for row_offset in tl.range(0, actual_rows):
-        pid = start_row + row_offset
+        pid = (start_row + row_offset).to(tl.int64)
 
         t_end = mrope_section_t
         h_end = t_end + mrope_section_h

src/liger_kernel/ops/llama4_rope.py

@@ -39,7 +39,7 @@ def _llama4_rope_kernel(
     Grid: (batch*seq, head)
     """
     # 2D grid
-    pid_bs = tl.program_id(0)  # over batch*seq
+    pid_bs = tl.program_id(0).to(tl.int64)  # over batch*seq
     pid_h = tl.program_id(1)  # over heads
 
     batch_idx = pid_bs // seq_len

src/liger_kernel/ops/qwen2vl_mrope.py

@@ -22,7 +22,7 @@ def _triton_qwen2vl_mrope(
     BLOCK_SIZE: tl.constexpr,
     BACKWARD_PASS: tl.constexpr = False,
 ):
-    pid = tl.program_id(0)
+    pid = tl.program_id(0).to(tl.int64)
 
     # locate start address
     q_ptr = q_ptr + pid * (n_qh * hd)