Update attention.py

attention.py

@@ -10,6 +10,23 @@ from torch import nn
 from .fc import FC_CLASS_REGISTRY
 from .norm import NORM_CLASS_REGISTRY
 
+
+def is_flash_v2_installed(v2_version: str='2.0.0'):
+    assert version.parse(v2_version) >= version.parse('2.0.0')
+    try:
+        import flash_attn as flash_attn
+    except:
+        return False
+    return version.parse(flash_attn.__version__) >= version.parse(v2_version)
+
+def is_flash_v1_installed():
+    try:
+        import flash_attn as flash_attn
+    except:
+        return False
+    return version.parse(flash_attn.__version__) < version.parse('2.0.0')
+
+
 def _reset_is_causal(num_query_tokens: int, num_key_tokens: int, original_is_causal: bool) -> bool:
     if original_is_causal and num_query_tokens != num_key_tokens:
         if num_query_tokens != 1: