refactor to avoid incorrect rope dispatch

Signed-off-by: Yan Ma <yan.ma@intel.com>

Author: yma11

vllm/model_executor/layers/rotary_embedding/base.py

@@ -14,7 +14,7 @@
 
 
 @CustomOp.register("rotary_embedding")
-class RotaryEmbedding(CustomOp):
+class RotaryEmbeddingBase(CustomOp):
     """Original rotary positional embedding."""
 
     def __init__(
@@ -86,6 +86,21 @@ def _match_cos_sin_cache_dtype(self, query: torch.Tensor) -> None:
         ):
             self.cos_sin_cache = self.cos_sin_cache.to(query.device, dtype=query.dtype)
 
+
+class RotaryEmbedding(RotaryEmbeddingBase):
+    def __init__(
+        self,
+        head_size: int,
+        rotary_dim: int,
+        max_position_embeddings: int,
+        base: float,
+        is_neox_style: bool,
+        dtype: torch.dtype,
+    ) -> None:
+        super().__init__(
+            head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
+        )
+
     def forward_native(
         self,
         positions: torch.Tensor,

vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py

@@ -7,7 +7,7 @@
 
 from vllm.platforms import current_platform
 
-from .base import RotaryEmbedding
+from .base import RotaryEmbeddingBase
 from .common import (
     rotate_gptj,
     rotate_neox,
@@ -22,7 +22,7 @@ def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
     return 0.1 * mscale * math.log(scale) + 1.0
 
 
-class DeepseekScalingRotaryEmbedding(RotaryEmbedding):
+class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
     """RotaryEmbedding extended with YaRN method.
 
     Credits to Peng et al. github.com/jquesnelle/yarn
@@ -146,5 +146,11 @@ def forward_native(
             key = key_rot
         return query, key
 
-    forward_cuda = forward_native
-    forward_xpu = forward_native
+    def forward_cuda(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor | None = None,
+        offsets: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        return self.forward_native(positions, query, key, offsets)

vllm/model_executor/layers/rotary_embedding/llama4_vision_rope.py

@@ -5,10 +5,10 @@
 
 import torch
 
-from .base import RotaryEmbedding
+from .base import RotaryEmbeddingBase
 
 
-class Llama4VisionRotaryEmbedding(RotaryEmbedding):
+class Llama4VisionRotaryEmbedding(RotaryEmbeddingBase):
     def __init__(
         self,
         head_size: int,
@@ -72,6 +72,9 @@ def forward_native(  # type: ignore[override]
         key_out = torch.view_as_real(key_ * freqs_ci).flatten(3)
         return query_out.type_as(query), key_out.type_as(key)
 
-    forward_cuda = forward_native
-    forward_hip = forward_native
-    forward_xpu = forward_native
+    def forward_cuda(  # type: ignore[override]
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        return self.forward_native(query, key)

vllm/model_executor/layers/rotary_embedding/mrope.py

@@ -7,7 +7,7 @@
 
 from vllm.triton_utils import tl, triton
 
-from .base import RotaryEmbedding
+from .base import RotaryEmbeddingBase
 from .common import apply_rotary_emb_dispatch
 from .yarn_scaling_rope import YaRNScalingRotaryEmbedding, yarn_get_mscale
 
@@ -199,7 +199,7 @@ def apply_interleaved_rope(x: torch.Tensor, mrope_section: list[int]) -> torch.T
     return x_t
 
 
-class MRotaryEmbedding(RotaryEmbedding):
+class MRotaryEmbedding(RotaryEmbeddingBase):
     """Rotary Embedding with Multimodal Sections."""
 
     def __init__(
@@ -357,24 +357,6 @@ def forward_cuda(
         key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
         return query, key
 
-    def forward_xpu(
-        self,
-        positions: torch.Tensor,
-        query: torch.Tensor,
-        key: torch.Tensor | None = None,
-        offsets: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        return self.forward_native(positions, query, key, offsets)
-
-    def forward_cpu(
-        self,
-        positions: torch.Tensor,
-        query: torch.Tensor,
-        key: torch.Tensor | None = None,
-        offsets: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        return self.forward_native(positions, query, key, offsets)
-
     @staticmethod
     def get_next_input_positions(
         mrope_position_delta: int,