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,17 +14,17 @@
 
 
 @CustomOp.register("rotary_embedding")
-class RotaryEmbedding(CustomOp):
+class RotaryEmbeddingBase(CustomOp):
     """Original rotary positional embedding."""
 
     def __init__(
-        self,
-        head_size: int,
-        rotary_dim: int,
-        max_position_embeddings: int,
-        base: float,
-        is_neox_style: bool,
-        dtype: torch.dtype,
+            self,
+            head_size: int,
+            rotary_dim: int,
+            max_position_embeddings: int,
+            base: float,
+            is_neox_style: bool,
+            dtype: torch.dtype,
     ) -> None:
         super().__init__()
         self.head_size = head_size
@@ -59,10 +59,10 @@ def _compute_inv_freq(self, base: float) -> torch.Tensor:
         # create the cache on GPU for faster initialization. This may cause
         # a slight numerical difference between the HF implementation and ours.
         inv_freq = 1.0 / (
-            base
-            ** (
-                torch.arange(0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim
-            )
+                base
+                ** (
+                        torch.arange(0, self.rotary_dim, 2, dtype=torch.float) / self.rotary_dim
+                )
         )
         return inv_freq
 
@@ -81,11 +81,24 @@ def _match_cos_sin_cache_dtype(self, query: torch.Tensor) -> None:
         # __setattr__ in nn.Module (called by `self.cos_sin_cache = ...`)
         # is expensive, so avoid calling it if possible
         if (
-            self.cos_sin_cache.device != query.device
-            or self.cos_sin_cache.dtype != query.dtype
+                self.cos_sin_cache.device != query.device
+                or self.cos_sin_cache.dtype != query.dtype
         ):
             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(query, key)

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

@@ -8,7 +8,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
 
@@ -200,7 +200,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__(
@@ -358,24 +358,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)
-
     @classmethod
     def get_input_positions(
         cls,