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, 05:00 UTC
[Perf] enable flashinfer rotary_embedding custom ops in DeepSeek rotary (#30729)
Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com>
jiahanc committed
Dec 18, 2025 at 11:31 UTC
53ad423f2638a3cbb95149928f127e07564581b7
2 files changed
+24
-2
vllm/model_executor/layers/rotary_embedding/base.py
+4
-1
@@ -38,7 +38,10 @@ class RotaryEmbeddingBase(CustomOp):
38
# and current_platform.is_cuda()
39
# and has_flashinfer()
40
# and self.head_size in [64, 128, 256, 512])
41
- self.use_flashinfer = False
41
+
42
+ # Check if use_flashinfer is already set
43
+ if not hasattr(self, "use_flashinfer"):
44
+ self.use_flashinfer = False
45
46
cache = self._compute_cos_sin_cache()
47
if not self.use_flashinfer:
vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py
+20
-1
@@ -6,6 +6,7 @@ import math
6
import torch
7
8
from vllm.platforms import current_platform
9
+from vllm.utils.flashinfer import has_flashinfer
10
11
from .base import RotaryEmbeddingBase
12
from .common import (
@@ -56,6 +57,13 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
57
/ yarn_get_mscale(self.scaling_factor, float(mscale_all_dim))
58
* attn_factor
59
)
60
+ self.use_flashinfer = (
61
+ self.enabled()
62
+ and dtype in (torch.float16, torch.bfloat16)
63
+ and current_platform.is_cuda()
64
+ and has_flashinfer()
65
+ and head_size in [64, 128, 256, 512]
66
+ )
67
super().__init__(
68
head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
69
)
@@ -162,4 +170,15 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
170
key: torch.Tensor | None = None,
171
offsets: torch.Tensor | None = None,
172
) -> tuple[torch.Tensor, torch.Tensor | None]:
165
- return self.forward_native(positions, query, key, offsets)
173
+ if self.use_flashinfer:
174
+ torch.ops.vllm.flashinfer_rotary_embedding(
175
+ torch.add(positions, offsets) if offsets is not None else positions,
176
+ query,
177
+ key,
178
+ self.head_size,
179
+ self.cos_sin_cache,
180
+ self.is_neox_style,
181
+ )
182
+ return query, key
183
+ else:
184
+ return self.forward_native(positions, query, key, offsets)