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

2024-08-05 23:59:23 +08:00 · 2024-08-05 23:59:23 +08:00 · 3f1b4d38e7
commit 3f1b4d38e7
parent 3f6ff1c1c5
1 changed files with 1 additions and 1 deletions
--- a/flash_attn/modules/mha.py
+++ b/flash_attn/modules/mha.py
@ -196,7 +196,7 @@ class FlashCrossAttention(nn.Module):
            assert cu_seqlens_k is not None
            assert cu_seqlens_k.dtype == torch.int32
            assert max_seqlen_k is not None
-            assert isinstance(max_seqlen, int)
+            assert isinstance(max_seqlen_k, int)
            return flash_attn_varlen_kvpacked_func(
                q,
                kv,