[Perf]: adjust apply_rotary api

Signed-off-by: MagnetoWang <[email protected]>
Co-authored-by: cynthieye <[email protected]>

Author: MagnetoWang

vllm/model_executor/layers/rotary_embedding.py

@@ -34,8 +34,7 @@
 from vllm.platforms import current_platform
 
 if is_flash_attn_2_available():
-    from flash_attn.ops.triton.rotary import apply_rotary
-
+    from flash_attn.layers.rotary import apply_rotary
 
 def _rotate_neox(x: torch.Tensor) -> torch.Tensor:
     x1 = x[..., :x.shape[-1] // 2]