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
[Perf] Use FlashInfer RoPE for RotaryEmbedding.forward_cuda when available (#21126)
Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Michael Goin committed
Sep 19, 2025 at 16:06 UTC
48ecb4438b2845f757edf228c5455ca6095938af
5 files changed
+78
-14
vllm/model_executor/layers/rotary_embedding/base.py
+28
-10
@@ -6,6 +6,8 @@ from typing import Optional
6
import torch
7
8
from vllm.model_executor.custom_op import CustomOp
9
+from vllm.platforms import current_platform
10
+from vllm.utils.flashinfer import has_flashinfer
11
12
from .common import apply_rotary_emb_torch
13
@@ -30,9 +32,17 @@ class RotaryEmbedding(CustomOp):
32
self.base = base
33
self.is_neox_style = is_neox_style
34
self.dtype = dtype
35
+ # Flashinfer only supports head_size=64, 128, 256, 512.
36
+ # https://github.com/flashinfer-ai/flashinfer/blob/ebfd655efe830048dba5d582aaa61d61d1cf9a87/include/flashinfer/utils.cuh#L174-L202
37
+ self.use_flashinfer = (self.enabled()
38
+ and dtype in (torch.float16, torch.bfloat16)
39
+ and current_platform.is_cuda()
40
+ and has_flashinfer()
41
+ and self.head_size in [64, 128, 256, 512])
42
43
cache = self._compute_cos_sin_cache()
35
- cache = cache.to(dtype)
44
+ if not self.use_flashinfer:
45
+ cache = cache.to(dtype)
46
self.cos_sin_cache: torch.Tensor
47
self.register_buffer("cos_sin_cache", cache, persistent=False)
48
@@ -57,6 +67,14 @@ class RotaryEmbedding(CustomOp):
67
cache = torch.cat((cos, sin), dim=-1)
68
return cache
69
70
+ def _match_cos_sin_cache_dtype(self, query: torch.Tensor) -> None:
71
+ # __setattr__ in nn.Module (called by `self.cos_sin_cache = ...`)
72
+ # is expensive, so avoid calling it if possible
73
+ if self.cos_sin_cache.device != query.device or \
74
+ self.cos_sin_cache.dtype != query.dtype:
75
+ self.cos_sin_cache = self.cos_sin_cache.to(query.device,
76
+ dtype=query.dtype)
77
+
78
def forward_native(
79
self,
80
positions: torch.Tensor,
@@ -94,15 +112,16 @@ class RotaryEmbedding(CustomOp):
112
query: torch.Tensor,
113
key: Optional[torch.Tensor] = None,
114
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
97
- from vllm import _custom_ops as ops
115
+ if self.use_flashinfer:
116
+ torch.ops.vllm.flashinfer_rotary_embedding(positions, query, key,
117
+ self.head_size,
118
+ self.cos_sin_cache,
119
+ self.is_neox_style)
120
+ return query, key
121
99
- # __setattr__ in nn.Module (called by `self.cos_sin_cache = ...`)
100
- # is expensive, so avoid calling it if possible
101
- if self.cos_sin_cache.device != query.device or \
102
- self.cos_sin_cache.dtype != query.dtype:
103
- self.cos_sin_cache = self.cos_sin_cache.to(query.device,
104
- dtype=query.dtype)
122
+ from vllm import _custom_ops as ops
123
124
+ self._match_cos_sin_cache_dtype(query)
125
# ops.rotary_embedding() is an in-place operation
126
# that updates the query and key tensors.
127
ops.rotary_embedding(positions, query, key, self.head_size,
@@ -117,8 +136,7 @@ class RotaryEmbedding(CustomOp):
136
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
137
from vllm._ipex_ops import ipex_ops as ops
138
120
- self.cos_sin_cache = self.cos_sin_cache.to(positions.device,
121
- dtype=query.dtype)
139
+ self._match_cos_sin_cache_dtype(query)
140
# ops.rotary_embedding() is an in-place operation
141
# that updates the query and key tensors.
142
if key is None:
vllm/model_executor/layers/rotary_embedding/common.py
+46
@@ -6,6 +6,7 @@ import math
6
import torch
7
8
from vllm.platforms import current_platform
9
+from vllm.utils import direct_register_custom_op
10
11
if current_platform.is_cuda():
12
from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
@@ -103,3 +104,48 @@ def yarn_get_mscale(scale: float = 1) -> float:
104
if scale <= 1:
105
return 1.0
106
return 0.1 * math.log(scale) + 1.0
107
+
108
+
109
+def _flashinfer_rotary_embedding(
110
+ positions: torch.Tensor,
111
+ query: torch.Tensor,
112
+ key: torch.Tensor,
113
+ head_size: int,
114
+ cos_sin_cache: torch.Tensor,
115
+ is_neox: bool,
116
+) -> None:
117
+ """Custom op wrapper for flashinfer's rotary embedding.
118
+
119
+ This is an in-place operation that modifies query and key tensors directly.
120
+ """
121
+ from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
122
+
123
+ apply_rope_with_cos_sin_cache_inplace(
124
+ positions=positions,
125
+ query=query,
126
+ key=key,
127
+ head_size=head_size,
128
+ cos_sin_cache=cos_sin_cache,
129
+ is_neox=is_neox,
130
+ )
131
+
132
+
133
+def _flashinfer_rotary_embedding_fake(
134
+ positions: torch.Tensor,
135
+ query: torch.Tensor,
136
+ key: torch.Tensor,
137
+ head_size: int,
138
+ cos_sin_cache: torch.Tensor,
139
+ is_neox: bool,
140
+) -> None:
141
+ return
142
+
143
+
144
+# Register flashinfer rotary embedding custom op
145
+direct_register_custom_op(
146
+ op_name="flashinfer_rotary_embedding",
147
+ op_func=_flashinfer_rotary_embedding,
148
+ mutates_args=["query", "key"], # These tensors are modified in-place
149
+ fake_impl=_flashinfer_rotary_embedding_fake,
150
+ dispatch_key=current_platform.dispatch_key,
151
+)
vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py
+1
-3
@@ -97,15 +97,13 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding):
97
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
98
"""PyTorch-native implementation equivalent to forward()."""
99
assert key is not None
100
+ self._match_cos_sin_cache_dtype(query)
101
query_rot = query[..., :self.rotary_dim]
102
key_rot = key[..., :self.rotary_dim]
103
if self.rotary_dim < self.head_size:
104
query_pass = query[..., self.rotary_dim:]
105
key_pass = key[..., self.rotary_dim:]
106
106
- if self.cos_sin_cache.device != positions.device:
107
- self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(
108
- positions.device)
107
cos_sin = self.cos_sin_cache[torch.add(positions, offsets)
108
if offsets is not None else positions]
109
cos, sin = cos_sin.chunk(2, dim=-1)
vllm/model_executor/layers/rotary_embedding/llama4_vision_rope.py
+1
-1
@@ -59,7 +59,7 @@ class Llama4VisionRotaryEmbedding(RotaryEmbedding):
59
key: Optional[torch.Tensor] = None,
60
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
61
assert key is not None
62
- self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(query.device)
62
+ self._match_cos_sin_cache_dtype(query)
63
query_ = torch.view_as_complex(query.float().reshape(
64
*query.shape[:-1], -1, 2))
65
key_ = torch.view_as_complex(key.float().reshape(
vllm/model_executor/layers/rotary_embedding/mrope.py
+2
@@ -245,6 +245,7 @@ class MRotaryEmbedding(RotaryEmbedding):
245
assert positions.ndim == 1 or positions.ndim == 2
246
assert key is not None
247
248
+ self._match_cos_sin_cache_dtype(query)
249
num_tokens = positions.shape[-1]
250
cos_sin = self.cos_sin_cache[positions]
251
cos, sin = cos_sin.chunk(2, dim=-1)
@@ -293,6 +294,7 @@ class MRotaryEmbedding(RotaryEmbedding):
294
assert positions.ndim == 1 or positions.ndim == 2
295
assert key is not None
296
297
+ self._match_cos_sin_cache_dtype(query)
298
num_tokens = positions.shape[-1]
299
cos_sin = self.cos_sin_cache[positions]
300
cos, sin = cos_sin.chunk(2, dim=-1)