Cast tl.program_id to int64 in llama4_rope and qwen2vl_mrope (#1342)

## Summary

`llama4_rope.py` and `qwen2vl_mrope.py` index the flattened `batch *
seq` token dimension with a raw `tl.program_id(0)`. Triton materializes
that as int32, so the pointer arithmetic built on top of it wraps once
the offset passes `2**31`. This is the same out-of-bounds class that
#803 reported and #804 fixed in `rope.py` and `rms_norm.py`, and neither
of these two files was touched by that PR. `llama4_rope.py` was added
afterwards as a fresh kernel body rather than reusing `_triton_rope`, so
it never inherited the fix, and `qwen2vl_mrope.py` predates #804.

Fixes #1335

## Details

Where it overflows.

`llama4_rope.py` launches on a `(batch_size * seq_len, n_heads_max)`
grid and passes `q.stride(1)` as `q_row_stride`, which for the
contiguous `(B, S, H, D)` layout is `H * D`. The kernel then computes:

```
base_offset = batch_idx * seq_len + seq_idx      # max B*S - 1
q_base      = q_ptr + base_offset * q_row_stride
```

so the largest element offset is `(B*S - 1) * H * D`, one row short of
`numel(q)`. The int32 product wraps once that passes int32 max, which as
a token budget for a single call means `total_tokens > int32_max / (H *
D)`:

| q width (`H * D`) | overflows at total tokens |
| --- | --- |
| 4096 (32 heads, head_dim 128) | 524,289 |
| 5120 (40 heads, head_dim 128) | 419,432 |
| 8192 (64 heads, head_dim 128) | 262,145 |

Using #803's own long-context motivation, `B=4, S=420000, H=40, D=128`
gives `1679999 * 5120 = 8601594880`, four times past int32 max. Nothing
that extreme is required though. A 64-head model at `B=8, S=65536`
reaches `4294959104`, already twice over.

`qwen2vl_mrope.py` has the same shape of arithmetic, `q_ptr + pid *
(n_qh * hd)` with `pid` up to `B*S - 1`, so the same token thresholds
apply.

The per-head term is safe in both and is left alone. `pid_h *
q_head_stride` is bounded by `n_heads * head_dim`, and widening the
batch term promotes the sum anyway.

The fix is the exact pattern used by #804 and by the Ascend port of
`llama4_rope` from #1053, which already carries
`tl.program_id(0).to(tl.int64)`.

Ascend backend. `backends/_ascend/ops/qwen2vl_mrope.py` has the same
gap, which #1335 also flags. That kernel derives its token index inside
a loop rather than straight from the program id, so I widened at `pid`,
the value actually used for addressing, and left `start_row`,
`rows_per_program` and `actual_rows` as int32 so the loop bound and
`tl.range` are unchanged.

Cost is one `.to(tl.int64)` on a single scalar per program, the same
tradeoff #804 accepted for the two hottest kernels in the repo.

Related, but deliberately not included. `_block_rms_norm_forward_kernel`
in `rms_norm.py`, the blocked forward added after #804, computes
`row_idx = tl.program_id(0) * BLOCK_ROW + tl.arange(0, BLOCK_ROW)` with
no cast, while `_block_rms_norm_backward_kernel` in the same file does
carry `tl.program_id(0).cast(tl.int64)`. That path is gated to `n_cols
<= 256`, so it needs on the order of 8.4M rows to wrap, far beyond what
the rope kernels need. I kept it out to hold this PR to the scope of
#1335. Glad to fold it in here or send it separately, whichever you
prefer.

## Testing Done

I do not have GPU access, so I have not run the Triton suites, and I
would rather say that than tick boxes I cannot honestly tick.

What I did verify:

- `ruff check` and `ruff format --check` pass on all three files.
- The overflow arithmetic above was computed in plain Python against the
real strides read from each kernel's launch site, not from the issue
text.
- Audited every `tl.program_id` under `src/liger_kernel/ops/` to confirm
these are the token-indexing kernels still missing the cast.

The change is a pure type widening. It is the identity for every offset
below `2**31`, so existing shapes stay bit-identical, and reproducing
the actual illegal memory access needs a q tensor above 4 GiB, which is
not CI-runnable. #804 landed under the same constraint and added no test
for the same reason.

- Hardware Type: none, no GPU access
- [ ] run `make test` to ensure correctness
- [x] run `make checkstyle` to ensure code style
- [ ] run `make test-convergence` to ensure convergence

---------

Co-authored-by: Vaibhav Jindal <[email protected]>
Co-authored-by: Copilot <[email protected]>

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)

test/transformers/test_llama4_rope.py

@@ -5,6 +5,7 @@ from test.utils import supports_bfloat16
 
 from liger_kernel.ops import LigerLlama4RopeFunction
 from liger_kernel.transformers.llama4_rope import liger_llama4_text_rotary_pos_emb
+from liger_kernel.utils import get_total_gpu_memory
 from liger_kernel.utils import infer_device
 
 try:
@@ -147,3 +148,74 @@ def test_functional_correctness(bsz, seq_len, num_q_heads, num_kv_heads, head_di
 
     assert torch.allclose(q1_grad, q2_grad, atol=atol, rtol=rtol)
     assert torch.allclose(k1_grad, k2_grad, atol=atol, rtol=rtol)
+
+
+# The kernel indexes the flattened batch*seq dimension as
+# `base_offset * q_row_stride`, whose largest value is
+# `(bsz * seq_len - 1) * n_q_heads * head_dim`. That product is one row short of
+# `q.numel()`, so it only exceeds int32 once q itself holds more than 2**31
+# elements. There is no cheaper shape: the bound is the element count.
+_HEAD_DIM = 128
+_N_Q_HEADS = 64
+_ROW = _N_Q_HEADS * _HEAD_DIM
+_SEQ_LEN = (2**31 - 1) // _ROW + 2  # smallest seq_len whose last offset wraps int32
+
+
+@pytest.mark.skipif(not IS_LLAMA4_AVAILABLE, reason="Llama4 is not available in transformers.")
+@pytest.mark.skipif(not supports_bfloat16(), reason="bfloat16 not supported on this GPU")
+@pytest.mark.parametrize(
+    "bsz, seq_len, num_q_heads, num_kv_heads, head_dim",
+    [
+        pytest.param(
+            1,
+            _SEQ_LEN,
+            _N_Q_HEADS,
+            1,
+            _HEAD_DIM,
+            marks=pytest.mark.skipif(
+                infer_device() == "cpu" or get_total_gpu_memory() < 20,
+                reason="This test requires a GPU with at least 20GB of memory",
+            ),
+        ),
+    ],
+)
+def test_row_offset_does_not_wrap_int32(bsz, seq_len, num_q_heads, num_kv_heads, head_dim):
+    """Regression for #1335: the last token must still be rotated, not read from a wrapped address."""
+    dtype = torch.bfloat16
+    n_kv_heads = num_kv_heads
+
+    q = torch.zeros((bsz, seq_len, num_q_heads, head_dim), device=device, dtype=dtype)
+    k = torch.zeros((bsz, seq_len, n_kv_heads, head_dim), device=device, dtype=dtype)
+
+    # Only the final token carries a signal. It sits at the offset that wraps.
+    q[0, -1].fill_(1.0)
+    k[0, -1].fill_(1.0)
+
+    config = Llama4TextConfig(
+        hidden_size=num_q_heads * head_dim,
+        num_attention_heads=num_q_heads,
+        num_key_value_heads=n_kv_heads,
+        head_dim=head_dim,
+        max_position_embeddings=seq_len,
+        rope_theta=10000.0,
+        rope_scaling=None,
+    )
+    rotary_emb = Llama4TextRotaryEmbedding(config=config, device=device)
+    pos_ids = torch.arange(seq_len, device=device).unsqueeze(0)
+    freqs_cis = rotary_emb(q, pos_ids)
+
+    # The kernel rotates q/k in place, so keep an unrotated copy of the last token for the reference.
+    q_last = q[:, -1:].float().clone()
+    k_last = k[:, -1:].float().clone()
+
+    q_out, k_out = LigerLlama4RopeFunction.apply(q, k, freqs_cis)
+
+    # A wrapped index leaves the last row untouched or fills it from elsewhere.
+    assert torch.isfinite(q_out[0, -1]).all()
+    assert torch.isfinite(k_out[0, -1]).all()
+    assert q_out[0, -1].abs().sum() > 0
+    assert k_out[0, -1].abs().sum() > 0
+
+    ref_q, ref_k = apply_rotary_emb(q_last, k_last, freqs_cis[:, -1:])
+    assert torch.allclose(q_out[:, -1:].float(), ref_q, atol=1e-1, rtol=1e-5)
+    assert torch.allclose(k_out[:, -1:].float(), ref_k, atol=1e-1, rtol=1e-5)