A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[AMD][Kernel][Bugfix] Cast offsets tensor bn to tl.int64 to avoid GPU segfault (#23692)
Signed-off-by: Randall Smith <Randall.Smith@amd.com>
rasmith committed
Sep 2, 2025 at 17:13 UTC
457e4719710e44014fb8b2eaf668d0c5cfa145c1
1 file changed
+3
-3
vllm/attention/ops/prefix_prefill.py
+3
-3
@@ -146,7 +146,7 @@ def _fwd_kernel(Q,
146
start_n = tl.multiple_of(start_n, BLOCK_SIZE)
147
# -- compute qk ----
148
bn = tl.load(B_Loc + cur_batch * stride_b_loc_b +
149
- (start_n // BLOCK_SIZE) * stride_b_loc_s)
149
+ (start_n // BLOCK_SIZE) * stride_b_loc_s).to(tl.int64)
150
# [D,BLOCK_SIZE]
151
off_k = (
152
bn[None, :] * stride_k_cache_bs + cur_kv_head * stride_k_cache_h +
@@ -367,7 +367,7 @@ def _fwd_kernel_flash_attn_v2(
367
bn = tl.load(B_Loc + cur_batch * stride_b_loc_b +
368
((start_n + offs_n) // block_size) * stride_b_loc_s,
369
mask=(start_n + offs_n) < cur_batch_ctx_len,
370
- other=0)
370
+ other=0).to(tl.int64)
371
off_k = (
372
bn[None, :] * stride_k_cache_bs + cur_kv_head * stride_k_cache_h +
373
(offs_d[:, None] // x) * stride_k_cache_d +
@@ -575,7 +575,7 @@ def _fwd_kernel_alibi(
575
bn = tl.load(B_Loc + cur_batch * stride_b_loc_b +
576
((start_n + offs_n) // block_size) * stride_b_loc_s,
577
mask=(start_n + offs_n) < cur_batch_ctx_len,
578
- other=0)
578
+ other=0).to(tl.int64)
579
off_k = (
580
bn[None, :] * stride_k_cache_bs + cur_kv_head * stride_k_cache_h +
581
(offs_d[:, None] // x) * stride_k_cache_d +