From 36bc29edf71a47015a03125bff1d98e0127ec71a Mon Sep 17 00:00:00 2001 From: Tri Dao Date: Mon, 22 Jan 2024 22:39:29 -0800 Subject: [PATCH] Use int64_t instead of uint32_t in kernel_traits.h --- csrc/flash_attn/src/kernel_traits.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/csrc/flash_attn/src/kernel_traits.h b/csrc/flash_attn/src/kernel_traits.h index cdc7608..be5fd77 100644 --- a/csrc/flash_attn/src/kernel_traits.h +++ b/csrc/flash_attn/src/kernel_traits.h @@ -24,7 +24,7 @@ struct Flash_kernel_traits { #endif using ElementAccum = float; - using index_t = uint32_t; + using index_t = int64_t; #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 using MMA_Atom_Arch = std::conditional_t<