Skip to content

Commit 020cf61

Browse files
tdoubleppatrickvonplaten
authored andcommitted
Fix overflow indexing in causal_conv1d kernel (vllm-project#20938)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com> Signed-off-by: Patrick von Platen <patrick.v.platen@gmail.com>
1 parent 33b41d9 commit 020cf61

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

vllm/model_executor/layers/mamba/ops/causal_conv1d.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,7 +92,8 @@ def _causal_conv1d_fwd_kernel( # continuous batching
9292

9393
if IS_CONTINUOUS_BATCHING:
9494
# cache_idx
95-
conv_state_batch_coord = tl.load(conv_state_indices_ptr + idx_seq)
95+
conv_state_batch_coord = tl.load(conv_state_indices_ptr + idx_seq).to(
96+
tl.int64)
9697
else:
9798
# cache_idx
9899
conv_state_batch_coord = idx_seq

0 commit comments

Comments
 (0)