Update modeling_llama.py

modeling_llama.py

@@ -60,14 +60,10 @@ def is_flash_attn_available():
         return False
 
     # Let's add an extra check to see if cuda is available
-    import torch
 
     return _is_package_available("flash_attn") and torch.cuda.is_available()
 
-if is_flash_attn_available():
-    from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func, flash_attn_qkvpacked_func, flash_attn_with_kvcache
-    # from flash_attn.flash_attn_interface import flash_attn_unpadded_qkvpacked_func
-    from flash_attn.bert_padding import unpad_input, pad_input
+
 
 
 
@@ -232,7 +228,10 @@ class LlamaAttention(nn.Module):
 
         attention_mask: [bsz, q_len]
         """
-
+        if is_flash_attn_available():
+            from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func, flash_attn_qkvpacked_func, flash_attn_with_kvcache
+            # from flash_attn.flash_attn_interface import flash_attn_unpadded_qkvpacked_func
+            from flash_attn.bert_padding import unpad_input, pad_input
         bsz, q_len, *_ = qkv.size()
 
         if key_padding_mask is None:
@@ -342,6 +341,7 @@ class LlamaAttention(nn.Module):
             return attn_output, attn_weights, past_key_value
         # use flash attention
         elif past_key_value is not None:
+            from flash_attn.flash_attn_interface import  flash_attn_with_kvcache
             output = flash_attn_with_kvcache(
                         query_states.transpose(1, 2),
                         key_states.transpose(1, 2),