A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 01:45 UTC
[XPU][bugfix] fix rope for llama4 and deepseek (#25145)
Signed-off-by: Yan Ma <yan.ma@intel.com>
Yan Ma committed
Oct 30, 2025 at 09:43 UTC
b798e39f931ad42354e0223de3d49e24523b79af
4 files changed
+22
-32
vllm/model_executor/layers/rotary_embedding/base.py
+16
-1
@@ -14,7 +14,7 @@ from .rocm_aiter_rope_ops import (
14
15
16
@CustomOp.register("rotary_embedding")
17
-class RotaryEmbedding(CustomOp):
17
+class RotaryEmbeddingBase(CustomOp):
18
"""Original rotary positional embedding."""
19
20
def __init__(
@@ -86,6 +86,21 @@ class RotaryEmbedding(CustomOp):
86
):
87
self.cos_sin_cache = self.cos_sin_cache.to(query.device, dtype=query.dtype)
88
89
+
90
+class RotaryEmbedding(RotaryEmbeddingBase):
91
+ def __init__(
92
+ self,
93
+ head_size: int,
94
+ rotary_dim: int,
95
+ max_position_embeddings: int,
96
+ base: float,
97
+ is_neox_style: bool,
98
+ dtype: torch.dtype,
99
+ ) -> None:
100
+ super().__init__(
101
+ head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
102
+ )
103
+
104
def forward_native(
105
self,
106
positions: torch.Tensor,
vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py
+2
-2
@@ -7,7 +7,7 @@ import torch
7
8
from vllm.platforms import current_platform
9
10
-from .base import RotaryEmbedding
10
+from .base import RotaryEmbeddingBase
11
from .common import (
12
rotate_gptj,
13
rotate_neox,
@@ -22,7 +22,7 @@ def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
22
return 0.1 * mscale * math.log(scale) + 1.0
23
24
25
-class DeepseekScalingRotaryEmbedding(RotaryEmbedding):
25
+class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
26
"""RotaryEmbedding extended with YaRN method.
27
28
Credits to Peng et al. github.com/jquesnelle/yarn
vllm/model_executor/layers/rotary_embedding/llama4_vision_rope.py
+2
-9
@@ -5,10 +5,10 @@ import math
5
6
import torch
7
8
-from .base import RotaryEmbedding
8
+from .base import RotaryEmbeddingBase
9
10
11
-class Llama4VisionRotaryEmbedding(RotaryEmbedding):
11
+class Llama4VisionRotaryEmbedding(RotaryEmbeddingBase):
12
def __init__(
13
self,
14
head_size: int,
@@ -78,10 +78,3 @@ class Llama4VisionRotaryEmbedding(RotaryEmbedding):
78
key: torch.Tensor | None = None,
79
) -> tuple[torch.Tensor, torch.Tensor | None]:
80
return self.forward_native(query, key)
81
-
82
- def forward_hip( # type: ignore[override]
83
- self,
84
- query: torch.Tensor,
85
- key: torch.Tensor | None = None,
86
- ) -> tuple[torch.Tensor, torch.Tensor | None]:
87
- return self.forward_native(query, key)
vllm/model_executor/layers/rotary_embedding/mrope.py
+2
-20
@@ -7,7 +7,7 @@ import torch
7
8
from vllm.triton_utils import tl, triton
9
10
-from .base import RotaryEmbedding
10
+from .base import RotaryEmbeddingBase
11
from .common import apply_rotary_emb_dispatch
12
from .yarn_scaling_rope import YaRNScalingRotaryEmbedding, yarn_get_mscale
13
@@ -199,7 +199,7 @@ def apply_interleaved_rope(x: torch.Tensor, mrope_section: list[int]) -> torch.T
199
return x_t
200
201
202
-class MRotaryEmbedding(RotaryEmbedding):
202
+class MRotaryEmbedding(RotaryEmbeddingBase):
203
"""Rotary Embedding with Multimodal Sections."""
204
205
def __init__(
@@ -357,24 +357,6 @@ class MRotaryEmbedding(RotaryEmbedding):
357
key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
358
return query, key
359
360
- def forward_xpu(
361
- self,
362
- positions: torch.Tensor,
363
- query: torch.Tensor,
364
- key: torch.Tensor | None = None,
365
- offsets: torch.Tensor | None = None,
366
- ) -> tuple[torch.Tensor, torch.Tensor | None]:
367
- return self.forward_native(positions, query, key, offsets)
368
-
369
- def forward_cpu(
370
- self,
371
- positions: torch.Tensor,
372
- query: torch.Tensor,
373
- key: torch.Tensor | None = None,
374
- offsets: torch.Tensor | None = None,
375
- ) -> tuple[torch.Tensor, torch.Tensor | None]:
376
- return self.forward_native(positions, query, key, offsets)
377
-
360
@staticmethod
361
def get_next_input_positions(
362
mrope_position_delta: int,