squall/flash-attention
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Use int64_t instead of uint32_t in kernel_traits.h

Browse Source
This commit is contained in:
Tri Dao 2024-01-22 22:39:29 -08:00
parent 000b67f5d8
commit 36bc29edf7
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
csrc/flash_attn/src/kernel_traits.h
View File

@ -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<