Update llama_xformers_attention.py

llama_xformers_attention.py

@@ -3,7 +3,11 @@ import torch.nn as nn
 
 from typing import Optional, Tuple
 
-from transformers.models.llama.modeling_llama import LlamaAttention
+from transformers.models.llama.modeling_llama import LlamaAttention, apply_rotary_pos_emb
+
+from xformers.ops.fmha import (
+    memory_efficient_attention,
+)
 
 class LlamaXFormersAttention(LlamaAttention):
     def forward(