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, 05:30 UTC
Revert "[ROCm][AITER] Support AITER Rope ops in RotaryEmbedding Module." (#22956)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com> Co-authored-by: vllmellm <vllm.ellm@embeddedllm.com>
TJian committed
Aug 14, 2025 at 23:39 UTC
b2f6c247a9b84556a8ea0e75bb4a2db765ff3315
4 files changed
+10
-204
vllm/model_executor/layers/rotary_embedding/base.py
-71
@@ -8,7 +8,6 @@ import torch
8
from vllm.model_executor.custom_op import CustomOp
9
10
from .common import apply_rotary_emb_dispatch, apply_rotary_emb_torch
11
-from .rocm_aiter_rope_ops import is_rocm_rotary_embedding_enabled
11
12
13
@CustomOp.register("rotary_embedding")
@@ -36,7 +35,6 @@ class RotaryEmbedding(CustomOp):
35
cache = cache.to(dtype)
36
self.cos_sin_cache: torch.Tensor
37
self.register_buffer("cos_sin_cache", cache, persistent=False)
39
- self.is_rocm_aiter_enabled = is_rocm_rotary_embedding_enabled()
38
39
def _compute_inv_freq(self, base: float) -> torch.Tensor:
40
"""Compute the inverse frequency."""
@@ -121,75 +119,6 @@ class RotaryEmbedding(CustomOp):
119
self.cos_sin_cache, self.is_neox_style)
120
return query, key
121
124
- def forward_hip(
125
- self,
126
- positions: torch.Tensor,
127
- query: torch.Tensor,
128
- key: Optional[torch.Tensor] = None,
129
- offsets: Optional[torch.Tensor] = None,
130
- is_nope_first=False,
131
- ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
132
- # currently only rotary embedding ops from AITER package are
133
- # supported for HiP forward.
134
- if self.is_rocm_aiter_enabled:
135
- return self.forward_hip_rocm_aiter(positions, query, key, offsets,
136
- is_nope_first)
137
- return self.forward_native(positions, query, key, offsets)
138
-
139
- def forward_hip_rocm_aiter(
140
- self,
141
- positions: torch.Tensor,
142
- # if is_nope_first
143
- # [[batch_size, seq_len, num_heads, nope_size+rope_size]
144
- # if NOT is_nope_first
145
- # [[batch_size, seq_len, num_heads, rope_size+nope_size],
146
- query: torch.Tensor,
147
- key: Optional[torch.Tensor] = None,
148
- offsets: Optional[torch.Tensor] = None,
149
- is_nope_first: bool = False,
150
- ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
151
- if self.cos_sin_cache.device != query.device or \
152
- self.cos_sin_cache.dtype != query.dtype:
153
- self.cos_sin_cache = self.cos_sin_cache.to(query.device,
154
- dtype=query.dtype)
155
- cos, sin = self.cos_sin_cache.chunk(2, dim=-1)
156
-
157
- cos = cos.unsqueeze(-2).unsqueeze(-2)
158
- sin = sin.unsqueeze(-2).unsqueeze(-2)
159
-
160
- rotate_style = 0 if self.is_neox_style else 1
161
-
162
- num_tokens = positions.numel()
163
-
164
- query_shape = query.shape
165
- query = query.view(1, num_tokens, -1, self.head_size)
166
- if key is not None:
167
- key_shape = key.shape
168
- key = key.view(1, num_tokens, -1, self.head_size)
169
-
170
- positions = positions.view(*query.shape[:2])
171
- if offsets is not None:
172
- offsets = offsets.view(*query.shape[:2])
173
-
174
- if not is_nope_first:
175
- query_ = query[..., :self.rotary_dim]
176
- key_ = key[..., :self.rotary_dim] if key is not None else None
177
- else:
178
- query_ = query[..., -self.rotary_dim:]
179
- key_ = key[..., -self.rotary_dim:] if key is not None else None
180
-
181
- if key_ is None:
182
- torch.ops.vllm.rocm_aiter_rotary_emb_without_key_forward_hip(
183
- positions, sin, cos, query_, offsets, rotate_style,
184
- is_nope_first)
185
- return query.view(query_shape), None
186
-
187
- torch.ops.vllm.rocm_aiter_rotary_emb_with_key_forward_hip(
188
- positions, sin, cos, query_, key_, offsets, rotate_style,
189
- is_nope_first)
190
-
191
- return query.view(query_shape), key.view(key_shape)
192
-
122
def forward_xpu(
123
self,
124
positions: torch.Tensor,
vllm/model_executor/layers/rotary_embedding/common.py
+2
-2
@@ -99,7 +99,7 @@ def yarn_linear_ramp_mask(low: float, high: float, dim: int,
99
return ramp_func
100
101
102
-def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
102
+def yarn_get_mscale(scale: float = 1) -> float:
103
if scale <= 1:
104
return 1.0
105
- return 0.1 * mscale * math.log(scale) + 1.0
105
+ return 0.1 * math.log(scale) + 1.0
vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py
+8
-4
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import math
5
from typing import Optional
6
7
import torch
@@ -9,7 +10,13 @@ from vllm.platforms import current_platform
10
11
from .base import RotaryEmbedding
12
from .common import (rotate_gptj, rotate_neox, yarn_find_correction_range,
12
- yarn_get_mscale, yarn_linear_ramp_mask)
13
+ yarn_linear_ramp_mask)
14
+
15
+
16
+def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float:
17
+ if scale <= 1:
18
+ return 1.0
19
+ return 0.1 * mscale * math.log(scale) + 1.0
20
21
22
class DeepseekScalingRotaryEmbedding(RotaryEmbedding):
@@ -89,9 +96,6 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding):
96
offsets: Optional[torch.Tensor] = None,
97
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
98
"""PyTorch-native implementation equivalent to forward()."""
92
- if self.is_rocm_aiter_enabled:
93
- return self.forward_hip_rocm_aiter(positions, query, key, offsets)
94
-
99
assert key is not None
100
query_rot = query[..., :self.rotary_dim]
101
key_rot = key[..., :self.rotary_dim]
vllm/model_executor/layers/rotary_embedding/rocm_aiter_rope_ops.py
deleted
-127
@@ -1,127 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-from typing import Optional
5
-
6
-import torch
7
-
8
-import vllm.envs as envs
9
-from vllm.platforms import current_platform
10
-from vllm.utils import direct_register_custom_op
11
-
12
-
13
-def is_rocm_rotary_embedding_enabled() -> bool:
14
- return (current_platform.is_rocm() and envs.VLLM_ROCM_USE_AITER)
15
-
16
-
17
-def rocm_aiter_rotary_emb_without_key_forward_hip_impl(
18
- positions: torch.Tensor,
19
- sin: torch.Tensor,
20
- cos: torch.Tensor,
21
- query: torch.Tensor,
22
- offsets: Optional[torch.Tensor] = None,
23
- rotate_style: int = 0,
24
- is_nope_first: bool = False,
25
-) -> None:
26
- import aiter as ops
27
- if offsets is None:
28
- ops.rope_cached_positions_fwd_inplace(
29
- query,
30
- cos,
31
- sin,
32
- positions,
33
- rotate_style,
34
- reuse_freqs_front_part=True,
35
- nope_first=is_nope_first,
36
- )
37
- else:
38
- ops.rope_cached_positions_offsets_fwd_inplace(
39
- query,
40
- cos,
41
- sin,
42
- positions,
43
- offsets,
44
- rotate_style,
45
- reuse_freqs_front_part=True,
46
- nope_first=is_nope_first,
47
- )
48
-
49
-
50
-def rocm_aiter_rotary_emb_with_key_forward_hip_impl(
51
- positions: torch.Tensor,
52
- sin: torch.Tensor,
53
- cos: torch.Tensor,
54
- query: torch.Tensor,
55
- key: torch.Tensor,
56
- offsets: Optional[torch.Tensor] = None,
57
- rotate_style: int = 0,
58
- is_nope_first: bool = False,
59
-) -> None:
60
- import aiter as ops
61
- if offsets is None:
62
- ops.rope_cached_positions_2c_fwd_inplace(
63
- query,
64
- key,
65
- cos,
66
- sin,
67
- positions,
68
- rotate_style,
69
- reuse_freqs_front_part=True,
70
- nope_first=is_nope_first,
71
- )
72
- else:
73
- ops.rope_cached_positions_offsets_2c_fwd_inplace(
74
- query,
75
- key,
76
- cos,
77
- sin,
78
- positions,
79
- offsets,
80
- rotate_style,
81
- reuse_freqs_front_part=True,
82
- nope_first=is_nope_first,
83
- )
84
-
85
-
86
-def rocm_aiter_rotary_emb_with_key_forward_hip_fake(
87
- positions: torch.Tensor,
88
- sin: torch.Tensor,
89
- cos: torch.Tensor,
90
- query: torch.Tensor,
91
- key: torch.Tensor,
92
- offsets: Optional[torch.Tensor] = None,
93
- rotate_style: int = 0,
94
- is_nope_first: bool = False,
95
-) -> None:
96
- pass
97
-
98
-
99
-def rocm_aiter_rotary_emb_without_key_forward_hip_fake(
100
- positions: torch.Tensor,
101
- sin: torch.Tensor,
102
- cos: torch.Tensor,
103
- query: torch.Tensor,
104
- offsets: Optional[torch.Tensor] = None,
105
- rotate_style: int = 0,
106
- is_nope_first: bool = False,
107
-) -> None:
108
- pass
109
-
110
-
111
-if is_rocm_rotary_embedding_enabled():
112
-
113
- direct_register_custom_op(
114
- op_name="rocm_aiter_rotary_emb_with_key_forward_hip",
115
- op_func=rocm_aiter_rotary_emb_with_key_forward_hip_impl,
116
- mutates_args=["key", "query"],
117
- fake_impl=rocm_aiter_rotary_emb_with_key_forward_hip_fake,
118
- dispatch_key=current_platform.dispatch_key,
119
- )
120
-
121
- direct_register_custom_op(
122
- op_name="rocm_aiter_rotary_emb_without_key_forward_hip",
123
- op_func=rocm_aiter_rotary_emb_without_key_forward_hip_impl,
124
- mutates_args=["query"],
125
- fake_impl=rocm_aiter_rotary_emb_without_key_forward_hip_fake,
126
- dispatch_key=current_platform.dispatch_key,
127
- )
\ No newline at end of file