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 8, 22:45 UTC
[Bugfix] Support `RotaryEmbedding` CustomOp for gpt-oss (#33800)
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Simon Danielsson committed
Feb 4, 2026 at 21:17 UTC
4292c90a2a188121ccbfd132def62031283d9d8a
4 files changed
+97
-18
tests/compile/test_rotary_embedding_compile.py
new
+68
@@ -0,0 +1,68 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+import vllm.envs as envs
8
+from vllm.compilation.decorators import support_torch_compile
9
+from vllm.config import (
10
+ CompilationConfig,
11
+ ModelConfig,
12
+ VllmConfig,
13
+ set_current_vllm_config,
14
+)
15
+from vllm.config.compilation import CompilationMode, CUDAGraphMode
16
+from vllm.model_executor.layers.rotary_embedding import get_rope
17
+from vllm.platforms import current_platform
18
+
19
+
20
+@support_torch_compile
21
+class RotaryEmbeddingCompileModule(torch.nn.Module):
22
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
23
+ super().__init__()
24
+ self.rotary_emb = get_rope(
25
+ head_size=32,
26
+ max_position=128,
27
+ dtype=torch.float32,
28
+ rope_parameters={"rope_type": "default", "rope_theta": 10000},
29
+ is_neox_style=True,
30
+ )
31
+
32
+ def forward(
33
+ self, positions: torch.Tensor, query: torch.Tensor, key: torch.Tensor
34
+ ) -> torch.Tensor:
35
+ q_rot, k_rot = self.rotary_emb(positions, query, key)
36
+ return q_rot + k_rot
37
+
38
+
39
+@pytest.mark.skipif(current_platform.is_cpu(), reason="Requires GPU for torch.compile")
40
+def test_rotary_embedding_torch_compile_with_custom_op(monkeypatch):
41
+ # Ensure env toggles take effect for this test only.
42
+ # The bytecode hook is required to detect buffer mutation in compiled code,
43
+ # and AOT compile bypasses that hook entirely.
44
+ envs.disable_envs_cache()
45
+ monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "1")
46
+ monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "0")
47
+
48
+ device = "cuda"
49
+ positions = torch.arange(16, device=device)
50
+ query = torch.randn(16, 32, device=device, dtype=torch.bfloat16)
51
+ key = torch.randn(16, 32, device=device, dtype=torch.bfloat16)
52
+
53
+ vllm_config = VllmConfig(
54
+ model_config=ModelConfig(dtype=torch.bfloat16),
55
+ compilation_config=CompilationConfig(
56
+ mode=CompilationMode.VLLM_COMPILE,
57
+ backend="inductor",
58
+ custom_ops=["+rotary_embedding"],
59
+ cudagraph_mode=CUDAGraphMode.NONE,
60
+ cudagraph_num_of_warmups=0,
61
+ ),
62
+ )
63
+
64
+ with set_current_vllm_config(vllm_config):
65
+ model = RotaryEmbeddingCompileModule(vllm_config=vllm_config)
66
+ model(positions, query, key)
67
+ assert model._compiled_bytecode is not None
68
+ assert "update" not in model._compiled_bytecode.co_names
vllm/model_executor/layers/rotary_embedding/base.py
+23
-12
@@ -86,14 +86,23 @@ class RotaryEmbeddingBase(CustomOp):
86
cache = torch.cat((cos, sin), dim=-1)
87
return cache
88
89
- def _match_cos_sin_cache_dtype(self, query: torch.Tensor) -> None:
89
+ def _match_cos_sin_cache_dtype(self, query: torch.Tensor) -> torch.Tensor:
90
# __setattr__ in nn.Module (called by `self.cos_sin_cache = ...`)
91
# is expensive, so avoid calling it if possible
92
+ cos_sin_cache = self.cos_sin_cache
93
if (
93
- self.cos_sin_cache.device != query.device
94
- or self.cos_sin_cache.dtype != query.dtype
94
+ cos_sin_cache.device == query.device
95
+ and self.cos_sin_cache.dtype == query.dtype
96
):
96
- self.cos_sin_cache = self.cos_sin_cache.to(query.device, dtype=query.dtype)
97
+ return cos_sin_cache
98
+
99
+ cos_sin_cache = cos_sin_cache.to(query.device, dtype=query.dtype)
100
+ # Avoid mutating buffers during torch.compile (cudagraph) tracing.
101
+ if torch.compiler.is_compiling():
102
+ return cos_sin_cache
103
+
104
+ self.cos_sin_cache = cos_sin_cache
105
+ return cos_sin_cache
106
107
def get_cos_sin(self, seqlen: int) -> tuple[torch.Tensor, torch.Tensor]:
108
cos_sin = self.cos_sin_cache[:seqlen]
@@ -172,13 +181,14 @@ class RotaryEmbedding(RotaryEmbeddingBase):
181
key: torch.Tensor | None = None,
182
) -> tuple[torch.Tensor, torch.Tensor | None]:
183
"""A PyTorch-native implementation of forward()."""
184
+ cos_sin_cache = self._match_cos_sin_cache_dtype(query)
185
return self.forward_static(
186
positions,
187
query,
188
key,
189
self.head_size,
190
self.rotary_dim,
181
- self.cos_sin_cache,
191
+ cos_sin_cache,
192
self.is_neox_style,
193
)
194
@@ -201,7 +211,7 @@ class RotaryEmbedding(RotaryEmbeddingBase):
211
212
from vllm import _custom_ops as ops
213
204
- self._match_cos_sin_cache_dtype(query)
214
+ cos_sin_cache = self._match_cos_sin_cache_dtype(query)
215
216
# ops.rotary_embedding() is an in-place operation
217
# that updates the query and key tensors.
@@ -210,7 +220,7 @@ class RotaryEmbedding(RotaryEmbeddingBase):
220
query,
221
key,
222
self.head_size,
213
- self.cos_sin_cache,
223
+ cos_sin_cache,
224
self.is_neox_style,
225
)
226
return query, key
@@ -222,12 +232,12 @@ class RotaryEmbedding(RotaryEmbeddingBase):
232
key: torch.Tensor | None = None,
233
) -> tuple[torch.Tensor, torch.Tensor | None]:
234
if self.is_rocm_triton_rotary_embed_enabled:
225
- self._match_cos_sin_cache_dtype(query)
235
+ cos_sin_cache = self._match_cos_sin_cache_dtype(query)
236
rocm_aiter_ops.triton_rotary_embed(
237
positions,
238
query,
239
key,
230
- self.cos_sin_cache,
240
+ cos_sin_cache,
241
self.head_size,
242
self.rotary_dim,
243
self.is_neox_style,
@@ -249,12 +259,13 @@ class RotaryEmbedding(RotaryEmbeddingBase):
259
else:
260
from vllm import _custom_ops as ops
261
262
+ cos_sin_cache = self._match_cos_sin_cache_dtype(query)
263
ops.rotary_embedding(
264
positions,
265
query,
266
key,
267
self.head_size,
257
- self.cos_sin_cache,
268
+ cos_sin_cache,
269
self.is_neox_style,
270
)
271
return query, key
@@ -267,7 +278,7 @@ class RotaryEmbedding(RotaryEmbeddingBase):
278
) -> tuple[torch.Tensor, torch.Tensor | None]:
279
from vllm import _custom_ops as ops
280
270
- self._match_cos_sin_cache_dtype(query)
281
+ cos_sin_cache = self._match_cos_sin_cache_dtype(query)
282
283
# ops.rotary_embedding() is an in-place operation
284
# that updates the query and key tensors.
@@ -276,7 +287,7 @@ class RotaryEmbedding(RotaryEmbeddingBase):
287
query,
288
key,
289
self.head_size,
279
- self.cos_sin_cache,
290
+ cos_sin_cache,
291
self.is_neox_style,
292
)
293
return query, key
vllm/model_executor/layers/rotary_embedding/deepseek_scaling_rope.py
+2
-2
@@ -120,14 +120,14 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbeddingBase):
120
) -> tuple[torch.Tensor, torch.Tensor | None]:
121
"""PyTorch-native implementation equivalent to forward()."""
122
assert key is not None
123
- self._match_cos_sin_cache_dtype(query)
123
+ cos_sin_cache = self._match_cos_sin_cache_dtype(query)
124
query_rot = query[..., : self.rotary_dim]
125
key_rot = key[..., : self.rotary_dim]
126
if self.rotary_dim < self.head_size:
127
query_pass = query[..., self.rotary_dim :]
128
key_pass = key[..., self.rotary_dim :]
129
130
- cos_sin = self.cos_sin_cache[
130
+ cos_sin = cos_sin_cache[
131
torch.add(positions, offsets) if offsets is not None else positions
132
]
133
cos, sin = cos_sin.chunk(2, dim=-1)
vllm/model_executor/layers/rotary_embedding/mrope.py
+4
-4
@@ -277,9 +277,9 @@ class MRotaryEmbedding(RotaryEmbeddingBase):
277
assert positions.ndim == 1 or positions.ndim == 2
278
assert key is not None
279
280
- self._match_cos_sin_cache_dtype(query)
280
+ cos_sin_cache = self._match_cos_sin_cache_dtype(query)
281
num_tokens = positions.shape[-1]
282
- cos_sin = self.cos_sin_cache[positions]
282
+ cos_sin = cos_sin_cache[positions]
283
cos, sin = cos_sin.chunk(2, dim=-1)
284
if positions.ndim == 2:
285
assert self.mrope_section
@@ -329,9 +329,9 @@ class MRotaryEmbedding(RotaryEmbeddingBase):
329
assert positions.ndim == 1 or positions.ndim == 2
330
assert key is not None
331
332
- self._match_cos_sin_cache_dtype(query)
332
+ cos_sin_cache = self._match_cos_sin_cache_dtype(query)
333
num_tokens = positions.shape[-1]
334
- cos_sin = self.cos_sin_cache[positions]
334
+ cos_sin = cos_sin_cache[positions]
335
cos, sin = cos_sin.chunk(2, dim=-1)
336
query_shape = query.shape
337
key_shape = key.shape