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

Use int64_t instead of uint32_t for index_t

Browse Source
This commit is contained in:
Tri Dao 2024-01-22 11:25:50 -08:00
parent e43a4ceaab
commit 000b67f5d8
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
csrc/flash_attn/src/flash.h
View File

@ -22,7 +22,7 @@ constexpr int D_DIM = 2;
//////////////////////////////////////////////////////////////////////////////////////////////////// ////////////////////////////////////////////////////////////////////////////////////////////////////
struct Qkv_params { struct Qkv_params {
using index_t = uint32_t; using index_t = int64_t;
// The QKV matrices. // The QKV matrices.
void *__restrict__ q_ptr; void *__restrict__ q_ptr;
void *__restrict__ k_ptr; void *__restrict__ k_ptr;
@ -99,7 +99,7 @@ struct Flash_fwd_params : public Qkv_params {
void * __restrict__ rotary_sin_ptr; void * __restrict__ rotary_sin_ptr;
// The indices to index into the KV cache. // The indices to index into the KV cache.
int *__restrict__ cache_batch_idx; int * __restrict__ cache_batch_idx;
// The dropout probability (probability of keeping an activation). // The dropout probability (probability of keeping an activation).
float p_dropout; float p_dropout;