llama4_rope / qwen2vl_mrope kernels missing the int64 program_id cast that #804 added to rope.py/rms_norm.py

[AmirF194]

🐛 Describe the bug

llama4_rope.py and qwen2vl_mrope.py are missing the same tl.program_id(0) int32-overflow cast that was already found and fixed in this repo for rope.py and rms_norm.py (#803 / PR #804: "Triton's default 32-bit tl.program_id(0) can overflow, leading to out-of-bounds memory accesses").

Fault lines at current HEAD 91ae44ae659ebaf40d4a851777f4f02515e9ba65:

• src/liger_kernel/ops/llama4_rope.py:43-52: pid_bs = tl.program_id(0) (no cast), then base_offset = batch_idx * seq_len + seq_idx and q_base = q_ptr + base_offset * q_row_stride, where q_row_stride = n_heads * head_dim.
• src/liger_kernel/ops/qwen2vl_mrope.py:25-29: pid = tl.program_id(0) (no cast), then q_ptr = q_ptr + pid * (n_qh * hd).

Neither file was touched by PR #804 (git log --oneline --follow on both files: llama4_rope.py has 2 commits, #843 introduce + #1053 NPU-only; qwen2vl_mrope.py has 7, none mention int64/overflow). llama4_rope.py was added after #804 as a fresh, non-shared kernel body rather than reusing _triton_rope, so it never inherited the fix; qwen2vl_mrope.py predates #804 entirely.

In-repo evidence the pattern is already known here: the Ascend NPU port of llama4_rope (PR #1053, src/liger_kernel/ops/backends/_ascend/ops/llama4_rope.py) does have pid = tl.program_id(0).to(tl.int64), while the CUDA/Triton version and the NPU port of qwen2vl_mrope do not.

I don't have GPU access to trigger the actual illegal-memory-access crash (same limitation noted on #803 for very long sequences), so this is reported as an arithmetic finding, not an observed crash. The overflow condition, using the numbers from #803's own long-context motivation:

batch=4, seq_len=420_000, n_heads=40, head_dim=128
max_base_offset = (batch-1)*seq_len + (seq_len-1) = 1_679_999
q_row_stride     = n_heads*head_dim = 5_120
true ptr offset  = 1_679_999 * 5_120 = 8_601_594_880   (> int32 max = 2_147_483_647)

so at this scale the element offset already exceeds what a 32-bit computation can represent, the same shape of overflow #804 fixed for rope.py/rms_norm.py.

Reproduce

Not applicable as a runtime repro (needs a GPU + long-context tensors, same limitation as #803). The arithmetic above was computed this session in plain Python to confirm the magnitude; happy to hand this to someone with GPU access to confirm the actual illegal-memory-access, or to open a PR applying the precedented one-line fix (tl.program_id(0).to(tl.int64), matching PR #804's and #1053's exact pattern) if that's useful without a runtime repro first.

Versions

Reported against main HEAD 91ae44ae659ebaf40d4a851777f4f02515e9ba65 (2026-07-31). No GPU available in this environment to run python -m liger_kernel.env_report against real hardware.

[akshansh47]

Opened #1343 — ports the #804 .to(tl.int64) cast to llama4_rope.py and qwen2vl_mrope.py.