Fix overflow indexing in causal_conv1d kernel (#20938)

Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
This commit is contained in:
Thomas Parnell 2025-07-14 23:43:07 +02:00 committed by GitHub
parent 149f2435a5
commit 86f3ac21ce
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -92,7 +92,8 @@ def _causal_conv1d_fwd_kernel( # continuous batching
if IS_CONTINUOUS_BATCHING:
# cache_idx
conv_state_batch_coord = tl.load(conv_state_indices_ptr + idx_seq)
conv_state_batch_coord = tl.load(conv_state_indices_ptr + idx_seq).to(
tl.int64)
else:
# cache_idx
conv_state_batch_coord = idx_seq