[update] fallback slow_attn

grn/models/basic.py

@@ -180,22 +180,8 @@ class SelfAttention(nn.Module):
             )
             attn_output = attn_output[0].reshape(B, L, C).contiguous()
         else:
-            from kernels import get_kernel
-            _k = get_kernel("kernels-community/vllm-flash-attn3", revision="main")
-            flash_attn_varlen_func = _k.flash_attn_varlen_func
-            attn_output = flash_attn_varlen_func(
-                q = query_states.squeeze(0),
-                k = key_states.squeeze(0),
-                v = value_states.squeeze(0),
-                cu_seqlens_q=cu_seqlens,
-                cu_seqlens_k=cu_seqlens,
-                max_seqlen_q=max_seqlen,
-                max_seqlen_k=max_seqlen,
-                softmax_scale=self.scale,
-            )
-            attn_output = attn_output.reshape(B, L, C).contiguous()
-            # # slow attn
-            # attn_output = slow_attn(query=query_states.transpose(1, 2), key=key_states.transpose(1, 2), value=value_states.transpose(1, 2), scale=self.scale, attn_mask=attn_bias_or_two_vector, dropout_p=0).transpose(1, 2).reshape(B, L, C)
+            # slow attn
+            attn_output = slow_attn(query=query_states.transpose(1, 2), key=key_states.transpose(1, 2), value=value_states.transpose(1, 2), scale=self.scale, attn_mask=attn_bias_or_two_vector, dropout_p=0).transpose(1, 2).reshape(B, L, C)
 
         if sp_manager.sp_on():
             # [B, raw_L, C/sp] --> [B, raw_L/sp, C]

requirements.txt

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