linkedin / linkedin/Liger-Kernel

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

Open Beginner friendly
#1,335 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
6.6k
Forks
603
Avg merge
1d 20h
Merged PRs (30d)
47

Description

### 🐛 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.

Contributor guide

Open the contributing guide

Research direction

Start by comparing src/liger_kernel/ops/llama4_rope.py and src/liger_kernel/ops/qwen2vl_mrope.py with the existing fix in rope.py and rms_norm.py, plus the Ascend llama4_rope implementation. Verify both program_id sites use the established int64 form, then check the relevant kernel changes against the cited long-context offset scenario; done means both affected kernels avoid 32-bit program-index overflow.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
70/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.