Add support for flash-attention2

flash-attention2 rename the `flash_attn_unpadded_func`

modeling_qwen.py

@@ -66,11 +66,15 @@ _CONFIG_FOR_DOC = "QWenConfig"
 QWen_PRETRAINED_MODEL_ARCHIVE_LIST = ["qwen-7b"]
 
 try:
-    from flash_attn.flash_attn_interface import flash_attn_unpadded_func
+    # from flash_attn.flash_attn_interface import flash_attn_unpadded_func
+    import flash_attn
+    if int(flash_attn.__version__.split(".")[0]) == 1:
+        from flash_attn.flash_attn_interface import flash_attn_unpadded_func
+    if int(flash_attn.__version__.split(".")[0]) == 2:
+        from flash_attn.flash_attn_interface import flash_attn_varlen_func as flash_attn_unpadded_func
 except ImportError:
     flash_attn_unpadded_func = None
-    print("Warning: import flash_attn fail, please install FlashAttention "
-          "https://github.com/Dao-AILab/flash-attention")
+    print("import flash_attn qkv fail")
 
 
 class FlashSelfAttention(torch.nn.Module):