[Bugfix] Fix int32 index overflow in LoRA punica kernels at long context (#53034)
Signed-off-by: Shuai Shao <ss@rungalileo.io> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
Shuai Shao committed
Aug 21, 2026 at 03:37 UTC
a60c66e3dcd563a293a9a3f54541e7405e95c82e
1 file changed
+6
-2
vllm/lora/ops/triton_ops/kernel_utils.py
+6
-2
@@ -223,7 +223,8 @@ def do_expand_kernel(
223
offset_cm = tl.arange(0, BLOCK_M)
224
c_ptr = (
225
out_ptr
226
- + ram[:, None] * output_d0_stride
226
+ # int64 to keep the row offset from overflowing at long context lengths
227
+ + ram[:, None].to(tl.int64) * output_d0_stride
228
+ offset_cn[None, :] * output_d1_stride
229
)
230
c_mask = (offset_cm[:, None] < M_LEN) & (offset_cn[None, :] < (cur_slice_start + N))
@@ -291,7 +292,10 @@ def do_shrink_kernel(
292
# Identify A and B block pointers
293
offset_k = pid_sk * BLOCK_K + tl.arange(0, BLOCK_K)
294
a_ptr = (
294
- input_ptr + ram[:, None] * input_d0_stride + offset_k[None, :] * input_d1_stride
295
+ input_ptr
296
+ # int64 to keep the row offset from overflowing at long context lengths
297
+ + ram[:, None].to(tl.int64) * input_d0_stride
298
+ + offset_k[None, :] * input_d1_stride
299
)
300
b_ptr = (
301
cur_lora_ptr