update flash-attn interface of apply_rotary_emb

vllm/model_executor/layers/rotary_embedding.py

@@ -40,6 +40,8 @@ from vllm.platforms import current_platform
 
 if current_platform.is_cuda():
     from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
+if current_platform.is_rocm():
+    from flash_attn.layers.rotary import apply_rotary_emb
 
 
 def _rotate_neox(x: torch.Tensor) -> torch.Tensor:

vllm/model_executor/layers/rotary_embedding/common.py

@@ -9,6 +9,8 @@ from vllm.platforms import current_platform
 
 if current_platform.is_cuda():
     from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
+if current_platform.is_rocm():
+    from flash_attn.layers.rotary import apply_rotary_emb
     
 
 # common functions

vllm/model_executor/models/ernie45_vl.py

@@ -115,6 +115,8 @@ def apply_rotary_pos_emb_vision(t: torch.Tensor,
     apply_rotary_emb = apply_rotary_emb_torch
     if current_platform.is_cuda():
         from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
+    if current_platform.is_rocm():
+        from flash_attn.layers.rotary import apply_rotary_emb
     output = apply_rotary_emb(t_, cos, sin).type_as(t)
     return output
 

vllm/model_executor/models/qwen2_vl.py

@@ -244,6 +244,8 @@ def apply_rotary_pos_emb_vision(t: torch.Tensor,
     apply_rotary_emb = apply_rotary_emb_torch
     if current_platform.is_cuda():
         from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
+    if current_platform.is_rocm():
+        from flash_attn.layers.rotary import apply_rotary_emb
     output = apply_rotary_emb(t_, cos, sin).type_as(t)
     return output