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

Fix Windows build (#816)

Browse Source
This commit is contained in:
Grigory Sizov 2024-02-08 02:41:53 +01:00 committed by GitHub
parent 61a7772479
commit 4687936413
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
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/flash_api.cpp
View File

@ -696,8 +696,8 @@ mha_varlen_fwd(at::Tensor &q, // total_q x num_heads x head_size, total_q := \s
}
if (seqlenq_ngroups_swapped) {
long size_before[] = {batch_size, max_seqlen_q, num_heads_k, head_size_og};
long size_after[] = {batch_size, num_heads_k * max_seqlen_q, head_size_og};
int64_t size_before[] = {batch_size, max_seqlen_q, num_heads_k, head_size_og};
int64_t size_after[] = {batch_size, num_heads_k * max_seqlen_q, head_size_og};
out = out.reshape(size_before).transpose(1, 2).reshape(size_after);
out_padded = out_padded.reshape(size_before).transpose(1, 2).reshape(size_after);
q_padded = q_padded.reshape(size_before).transpose(1, 2).reshape(size_after);