Cast tl.program_id to int64 in llama4_rope and qwen2vl_mrope

[adityasingh2400]

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
• run make checkstyle to ensure code style
• run make test-convergence to ensure convergence

[adityasingh2400]

Updated the branch onto main, so this is no longer behind and is mergeable again.

Worth flagging that no CI has run on this yet. All four workflows have been sitting at action_required since the PR opened:

• Checkstyle
• NVIDIA GPU
• AMD GPU
• Intel GPU

Approving the workflow run is all it needs to produce a signal.

[Tcc0403]

Could you help adding a mentioned test case to existing test suite?

[adityasingh2400]

Added in 42bc7ed, test_row_offset_does_not_wrap_int32 in test/transformers/test_llama4_rope.py.

One thing you should know before deciding whether to keep it, because it constrains what a real test can look like. The kernel indexes with base_offset * q_row_stride, whose largest value is (bsz * seq_len - 1) * n_q_heads * head_dim. That is one row short of q.numel(), so the product only passes int32 once q itself holds more than 2**31 elements. There is no cheaper shape, the bound is the element count, so trading heads against sequence length just moves the same 2.1 billion elements around:

H=64  D=128  seq_len=262145  q = 2,147,491,840 elements
H=128 D=128  seq_len=131073  q = 2,147,500,032 elements

At bfloat16 that is 4.3 GB for q, and the op does q.to(compute_dtype).contiguous() and returns a new tensor, so the real ask is roughly 13 GB free. The test is gated on torch.cuda.mem_get_info() and skips below that, so it is inert on a runner that cannot hold it.

The test writes into the last token only, since that is the row whose offset wraps, then checks the output is finite, nonzero, and matches apply_rotary_emb for that position.

Worth noting that #804, which made the same int64 change in rope.py and rms_norm.py, shipped with no test, I assume for this reason. If you would rather not carry a 13 GB test I am happy to drop it and leave the fix as is, or move it behind an explicit marker you run manually.

Separately, CI still has not run on this branch. All four workflows have been at action_required since the PR opened.

[Tcc0403]

can you add the overflow test case as a pytest.param with marks=pytest.mark.skipif total memory less than 20GB?

there's a util func you can reuse for this purpose

Liger-Kernel/src/liger_kernel/utils.py

Line 202 in a3f3654

def get_total_gpu_memory() -> int:

if there's any concern, feel free to chat about it!

[adityasingh2400]

Done in 1b120b9, using get_total_gpu_memory as you suggested.

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",
    ),
),

One thing I added on purpose. The mark is evaluated at collection time, and get_total_gpu_memory raises RuntimeError when infer_device() returns cpu, so calling it unguarded would fail collection for the whole file on a cpu only machine. The or short circuits before that. I checked it locally where infer_device() is cpu and the util is never reached.

The test body now takes the parametrized values rather than the module constants, so the shape lives entirely in the param. ruff passes.