[Bugfix] Fix int32 overflow in triton_decode_attention page offsets (#47671)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Yifan Qiao committed
Jul 6, 2026 at 07:36 UTC
095adf1fdc93d5554a3afbf26e97b399309c80db
1 file changed
+2
-2
vllm/v1/attention/ops/triton_decode_attention.py
+2
-2
@@ -133,7 +133,7 @@ def _fwd_kernel_stage1(
133
+ offs_n // PAGE_SIZE,
134
mask=offs_n < split_kv_end,
135
other=0,
136
- )
136
+ ).to(tl.int64) # page_number * page stride overflows int32
137
kv_in_page = offs_n % PAGE_SIZE
138
offs_buf_k = (
139
(kv_page_number * stride_buf_kpbs + kv_in_page * stride_buf_kbs)[
@@ -375,7 +375,7 @@ def _fwd_grouped_kernel_stage1(
375
mask=offs_n < split_kv_end,
376
other=0,
377
cache_modifier=".ca",
378
- )
378
+ ).to(tl.int64) # page_number * page stride overflows int32
379
kv_off_k = (
380
kv_page_number * stride_buf_kpbs + (offs_n % PAGE_SIZE) * stride_buf_kbs
381
)