From c2401a05487532666b698bd6952fd87401deea2f Mon Sep 17 00:00:00 2001 From: Aditya Singh Date: Wed, 5 Aug 2026 05:00:51 -0700 Subject: [PATCH 1/3] Cast tl.program_id to int64 in llama4_rope and qwen2vl_mrope Both kernels index the flattened batch*seq token dimension with a raw tl.program_id(0), which Triton materializes as int32. The resulting pointer arithmetic (base_offset * q_row_stride in llama4_rope, pid * (n_qh * hd) in qwen2vl_mrope) overflows once the q tensor exceeds 2**31 elements, which is the same out-of-bounds class that #804 fixed in rope.py and rms_norm.py. Applies the identical .to(tl.int64) widening. llama4_rope.py was added after #804 as a fresh kernel body rather than reusing _triton_rope, so it never inherited the fix, and qwen2vl_mrope.py predates #804 entirely. The Ascend port of qwen2vl_mrope has the same gap, widened at the pid used for addressing so loop control stays int32. Fixes #1335 --- src/liger_kernel/ops/backends/_ascend/ops/qwen2vl_mrope.py | 2 +- src/liger_kernel/ops/llama4_rope.py | 2 +- src/liger_kernel/ops/qwen2vl_mrope.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/liger_kernel/ops/backends/_ascend/ops/qwen2vl_mrope.py b/src/liger_kernel/ops/backends/_ascend/ops/qwen2vl_mrope.py index d273b7ec9..8c83e4428 100644 --- a/src/liger_kernel/ops/backends/_ascend/ops/qwen2vl_mrope.py +++ b/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 diff --git a/src/liger_kernel/ops/llama4_rope.py b/src/liger_kernel/ops/llama4_rope.py index 9167d69f3..02a8647be 100644 --- a/src/liger_kernel/ops/llama4_rope.py +++ b/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 diff --git a/src/liger_kernel/ops/qwen2vl_mrope.py b/src/liger_kernel/ops/qwen2vl_mrope.py index fbd120f96..b431c8976 100644 --- a/src/liger_kernel/ops/qwen2vl_mrope.py +++ b/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) From 42bc7ed70336e42af8d1e0e2d7c358b970693703 Mon Sep 17 00:00:00 2001 From: Aditya Singh Date: Sat, 8 Aug 2026 14:09:34 -0700 Subject: [PATCH 2/3] Add int32 row-offset regression test for llama4 rope --- test/transformers/test_llama4_rope.py | 66 +++++++++++++++++++++++++++ 1 file changed, 66 insertions(+) diff --git a/test/transformers/test_llama4_rope.py b/test/transformers/test_llama4_rope.py index 91f3d6dea..8132c4916 100644 --- a/test/transformers/test_llama4_rope.py +++ b/test/transformers/test_llama4_rope.py @@ -147,3 +147,69 @@ 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) + + +def _free_device_memory_bytes(): + if device != "cuda" or not torch.cuda.is_available(): + return 0 + free, _total = torch.cuda.mem_get_info() + return free + + +# 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 +# q, its `.contiguous()` copy inside the op, and the output. +_NEEDED_BYTES = 3 * _SEQ_LEN * _ROW * 2 + + +@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.skipif( + _free_device_memory_bytes() < _NEEDED_BYTES, + reason=f"needs about {_NEEDED_BYTES / 1e9:.0f} GB free to build a q tensor past 2**31 elements", +) +def test_row_offset_does_not_wrap_int32(): + """Regression for #1335: the last token must still be rotated, not read from a wrapped address.""" + dtype = torch.bfloat16 + n_kv_heads = 1 + + q = torch.zeros((1, _SEQ_LEN, _N_Q_HEADS, _HEAD_DIM), device=device, dtype=dtype) + k = torch.zeros((1, _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=_N_Q_HEADS * _HEAD_DIM, + num_attention_heads=_N_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) + + 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[:, -1:].float(), k[:, -1:].float(), freqs_cis[-1:].unsqueeze(0) + ) + 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) From 1b120b97beb2b1e57c2e9dc5a5df15d89d2000ff Mon Sep 17 00:00:00 2001 From: Aditya Singh Date: Sat, 8 Aug 2026 23:55:35 -0700 Subject: [PATCH 3/3] Gate the overflow test with get_total_gpu_memory via pytest.param --- test/transformers/test_llama4_rope.py | 46 +++++++++++++++------------ 1 file changed, 25 insertions(+), 21 deletions(-) diff --git a/test/transformers/test_llama4_rope.py b/test/transformers/test_llama4_rope.py index 8132c4916..2a91afe0a 100644 --- a/test/transformers/test_llama4_rope.py +++ b/test/transformers/test_llama4_rope.py @@ -5,6 +5,7 @@ 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: @@ -149,13 +150,6 @@ def test_functional_correctness(bsz, seq_len, num_q_heads, num_kv_heads, head_di assert torch.allclose(k1_grad, k2_grad, atol=atol, rtol=rtol) -def _free_device_memory_bytes(): - if device != "cuda" or not torch.cuda.is_available(): - return 0 - free, _total = torch.cuda.mem_get_info() - return free - - # 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 @@ -165,39 +159,49 @@ def _free_device_memory_bytes(): _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 -# q, its `.contiguous()` copy inside the op, and the output. -_NEEDED_BYTES = 3 * _SEQ_LEN * _ROW * 2 @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.skipif( - _free_device_memory_bytes() < _NEEDED_BYTES, - reason=f"needs about {_NEEDED_BYTES / 1e9:.0f} GB free to build a q tensor past 2**31 elements", +@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(): +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 = 1 + n_kv_heads = num_kv_heads - q = torch.zeros((1, _SEQ_LEN, _N_Q_HEADS, _HEAD_DIM), device=device, dtype=dtype) - k = torch.zeros((1, _SEQ_LEN, n_kv_heads, _HEAD_DIM), device=device, dtype=dtype) + 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=_N_Q_HEADS * _HEAD_DIM, - num_attention_heads=_N_Q_HEADS, + 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, + 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) + pos_ids = torch.arange(seq_len, device=device).unsqueeze(0) freqs_cis = rotary_emb(q, pos_ids) q_out, k_out = LigerLlama4RopeFunction.apply(q, k, freqs_cis)