linkedin / linkedin/Liger-Kernel
llama4_rope / qwen2vl_mrope kernels missing the int64 program_id cast that #804 added to rope.py/rms_norm.py
- 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
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