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

Fix: check the type of max_seqlen_k instead of checking max_seqlen twice (#1127)

Browse Source
This commit is contained in:
SueJane 2024-08-05 23:59:23 +08:00 committed by GitHub
parent 3f6ff1c1c5
commit 3f1b4d38e7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
flash_attn/modules/mha.py
View File

@ -196,7 +196,7 @@ class FlashCrossAttention(nn.Module):
assert cu_seqlens_k is not None assert cu_seqlens_k is not None
assert cu_seqlens_k.dtype == torch.int32 assert cu_seqlens_k.dtype == torch.int32
assert max_seqlen_k is not None assert max_seqlen_k is not None
assert isinstance(max_seqlen, int) assert isinstance(max_seqlen_k, int)
return flash_attn_varlen_kvpacked_func( return flash_attn_varlen_kvpacked_func(
q, q,
kv, kv,