[Perf]Optimize MRotaryEmbedding implementation to use cuda operator for improved inference performance

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

Author: cynthieye

vllm/model_executor/layers/rotary_embedding.py

@@ -956,6 +956,12 @@ def forward(
         """
         assert positions.ndim == 1 or positions.ndim == 2
 
+        if current_platform.is_cuda_alike():
+            from vllm import _custom_ops as ops
+            ops.rotary_embedding(positions, query, key, self.head_size,
+                                 self.cos_sin_cache, self.is_neox_style)
+            return query, key
+
         num_tokens = positions.shape[-1]
         cos_sin = self.cos_sin_cache[positions]
         cos, sin = cos_sin.chunk(2, dim=-1)