[update] slow attn

.gitignore

@@ -0,0 +1 @@
+*.pyc

grn/models/basic.py

@@ -201,24 +201,7 @@ class SelfAttention(nn.Module):
         if self.use_flex_attn and attn_fn is not None:
             attn_output = attn_fn(query_states.to(value_states.dtype), key_states.to(value_states.dtype), value_states, scale=scale).transpose(1, 2).reshape(B, L, C)
         else:
-            if attn_bias_or_two_vector is None:
-                # fa2, flash_attn_func input/output should be (batch_size, seqlen, nheads, headdim)
-                from flash_attn import flash_attn_func, flash_attn_varlen_func
-                attn_output = flash_attn_varlen_func(
-                    q = query_states.permute([0,2,1,3]).to(torch.bfloat16).squeeze(0),
-                    k = key_states.permute([0,2,1,3]).to(torch.bfloat16).squeeze(0),
-                    v = value_states.permute([0,2,1,3]).to(torch.bfloat16).squeeze(0),
-                    cu_seqlens_q = torch.tensor([0] + split_cond_uncond, device=query_states.device).cumsum(-1).to(torch.int32),
-                    cu_seqlens_k = torch.tensor([0] + cu_seqlens_k, device=query_states.device).cumsum(-1).to(torch.int32),
-                    max_seqlen_q = max(split_cond_uncond),
-                    max_seqlen_k = max(cu_seqlens_k),
-                    softmax_scale=scale,
-                )
-                attn_output = attn_output.reshape(B, L, C)
-                # attn_output = flash_attn_func(query_states.permute([0,2,1,3]).to(torch.bfloat16), key_states.permute([0,2,1,3]).to(torch.bfloat16), value_states.permute([0,2,1,3]).to(torch.bfloat16), softmax_scale=scale)
-            else:
-                # slow attn
-                attn_output = slow_attn(query=query_states, key=key_states, value=value_states, scale=scale, attn_mask=attn_bias_or_two_vector, dropout_p=0).transpose(1, 2).reshape(B, L, C)
+            attn_output = slow_attn(query=query_states, key=key_states, value=value_states, scale=scale, attn_mask=attn_bias_or_two_vector, dropout_p=0).transpose(1, 2).reshape(B, L, C)
 
             # fa3, flash_attn_func input/output should be (batch_size, seqlen, nheads, headdim)
             # from flash_attn_interface import flash_attn_qkvpacked_func, flash_attn_func

requirements.txt

@@ -15,4 +15,3 @@ ftfy>=6.1.1
 transformers>=4.35.0
 regex>=2023.10.3
 pyyaml>=6.0
-flash-attn