From a60c66e3dcd563a293a9a3f54541e7405e95c82e Mon Sep 17 00:00:00 2001 From: Shuai Shao Date: Fri, 21 Aug 2026 03:37:06 -0700 Subject: [PATCH] [Bugfix] Fix int32 index overflow in LoRA punica kernels at long context (#53034) Signed-off-by: Shuai Shao Co-authored-by: Jee Jee Li --- vllm/lora/ops/triton_ops/kernel_utils.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/vllm/lora/ops/triton_ops/kernel_utils.py b/vllm/lora/ops/triton_ops/kernel_utils.py index c6c2a02fde..34a9939ffc 100644 --- a/vllm/lora/ops/triton_ops/kernel_utils.py +++ b/vllm/lora/ops/triton_ops/kernel_utils.py @@ -223,7 +223,8 @@ def do_expand_kernel( offset_cm = tl.arange(0, BLOCK_M) c_ptr = ( out_ptr - + ram[:, None] * output_d0_stride + # int64 to keep the row offset from overflowing at long context lengths + + ram[:, None].to(tl.int64) * output_d0_stride + offset_cn[None, :] * output_d1_stride ) c_mask = (offset_cm[:, None] < M_LEN) & (offset_cn[None, :] < (cur_slice_start + N)) @@ -291,7 +292,10 @@ def do_shrink_kernel( # Identify A and B block pointers offset_k = pid_sk * BLOCK_K + tl.arange(0, BLOCK_K) a_ptr = ( - input_ptr + ram[:, None] * input_d0_stride + offset_k[None, :] * input_d1_stride + input_ptr + # int64 to keep the row offset from overflowing at long context lengths + + ram[:, None].to(tl.int64) * input_d0_stride + + offset_k[None, :] * input_d1_stride ) b_ptr = ( cur_lora_ptr