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 11, 03:45 UTC
[CustomOp][Refactor] Extract common methods for ApplyRotaryEmb CustomOp (#31021)
Signed-off-by: shen-shanshan <467638484@qq.com>
Shanshan Shen committed
Dec 19, 2025 at 22:16 UTC
23a1946e3b1b2dd45a98c3dfd86556563543c91b
1 file changed
+35
-28
vllm/model_executor/layers/rotary_embedding/common.py
+35
-28
@@ -178,6 +178,37 @@ class ApplyRotaryEmb(CustomOp):
178
output = output.to(origin_dtype)
179
return output
180
181
+ def _pre_process(
182
+ self,
183
+ x: torch.Tensor,
184
+ cos: torch.Tensor,
185
+ sin: torch.Tensor,
186
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Size, torch.dtype]:
187
+ origin_shape = x.shape
188
+ if len(origin_shape) == 3:
189
+ # x: [seq_len, num_heads, head_size]
190
+ x = x.unsqueeze(0)
191
+
192
+ origin_dtype = x.dtype
193
+ if self.enable_fp32_compute:
194
+ x = x.float()
195
+ cos = cos.float()
196
+ sin = sin.float()
197
+
198
+ return x, cos, sin, origin_shape, origin_dtype
199
+
200
+ def _post_process(
201
+ self,
202
+ output: torch.Tensor,
203
+ origin_shape: torch.Size,
204
+ origin_dtype: torch.dtype,
205
+ ) -> torch.Tensor:
206
+ if len(origin_shape) == 3:
207
+ output = output.squeeze(0)
208
+ if self.enable_fp32_compute:
209
+ output = output.to(origin_dtype)
210
+ return output
211
+
212
def forward_native(
213
self,
214
x: torch.Tensor,
@@ -197,16 +228,7 @@ class ApplyRotaryEmb(CustomOp):
228
) -> torch.Tensor:
229
from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
230
200
- origin_dtype = x.dtype
201
- if self.enable_fp32_compute:
202
- x = x.float()
203
- cos = cos.float()
204
- sin = sin.float()
205
-
206
- origin_shape = x.shape
207
- if len(origin_shape) == 3:
208
- # x: [seq_len, num_heads, head_size]
209
- x = x.unsqueeze(0)
231
+ x, cos, sin, origin_shape, origin_dtype = self._pre_process(x, cos, sin)
232
233
"""
234
Arguments of apply_rotary_emb() in vllm_flash_attn:
@@ -218,10 +240,7 @@ class ApplyRotaryEmb(CustomOp):
240
interleaved = not self.is_neox_style
241
output = apply_rotary_emb(x, cos, sin, interleaved)
242
221
- if len(origin_shape) == 3:
222
- output = output.squeeze(0)
223
- if self.enable_fp32_compute:
224
- output = output.to(origin_dtype)
243
+ output = self._post_process(output, origin_shape, origin_dtype)
244
return output
245
246
def forward_hip(
@@ -231,16 +250,7 @@ class ApplyRotaryEmb(CustomOp):
250
sin: torch.Tensor,
251
) -> torch.Tensor:
252
if self.apply_rotary_emb_flash_attn is not None:
234
- origin_dtype = x.dtype
235
- if self.enable_fp32_compute:
236
- x = x.float()
237
- cos = cos.float()
238
- sin = sin.float()
239
-
240
- origin_shape = x.shape
241
- if len(origin_shape) == 3:
242
- # x: [seq_len, num_heads, head_size]
243
- x = x.unsqueeze(0)
253
+ x, cos, sin, origin_shape, origin_dtype = self._pre_process(x, cos, sin)
254
255
"""
256
Arguments of apply_rotary() in flash_attn:
@@ -254,10 +264,7 @@ class ApplyRotaryEmb(CustomOp):
264
x, cos, sin, interleaved=interleaved
265
).type_as(x)
266
257
- if len(origin_shape) == 3:
258
- output = output.squeeze(0)
259
- if self.enable_fp32_compute:
260
- output = output.to(origin_dtype)
267
+ output = self._post_process(output, origin_shape, origin_dtype)
268
else:
269
# Falling back to PyTorch native implementation.
270
output = self.forward_native(x, cos, sin)