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 10, 17:45 UTC
[Bugfix][CPU] Fix RotaryEmbedding fallback causing gibberish with --enforce-eager (#31643)
Signed-off-by: rickychen-infinirc <ricky.chen@infinirc.com>
RickyChen / 陳昭儒 committed
Jan 6, 2026 at 01:25 UTC
c455b771fdfd9185ca5c20d315c07dd5390b5459
2 files changed
+25
-2
vllm/model_executor/custom_op.py
+3
-2
@@ -67,8 +67,9 @@ class CustomOp(nn.Module):
67
return self.forward_native(*args, **kwargs)
68
69
def forward_cpu(self, *args, **kwargs):
70
- # By default, we assume that CPU ops are compatible with CUDA ops.
71
- return self.forward_cuda(*args, **kwargs)
70
+ # By default, we assume that CPU ops are compatible with the
71
+ # PyTorch-native implementation.
72
+ return self.forward_native(*args, **kwargs)
73
74
def forward_tpu(self, *args, **kwargs):
75
# By default, we assume that TPU ops are compatible with the
vllm/model_executor/layers/rotary_embedding/base.py
+22
@@ -250,6 +250,28 @@ class RotaryEmbedding(RotaryEmbeddingBase):
250
)
251
return query, key
252
253
+ def forward_cpu(
254
+ self,
255
+ positions: torch.Tensor,
256
+ query: torch.Tensor,
257
+ key: torch.Tensor | None = None,
258
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
259
+ from vllm import _custom_ops as ops
260
+
261
+ self._match_cos_sin_cache_dtype(query)
262
+
263
+ # ops.rotary_embedding() is an in-place operation
264
+ # that updates the query and key tensors.
265
+ ops.rotary_embedding(
266
+ positions,
267
+ query,
268
+ key,
269
+ self.head_size,
270
+ self.cos_sin_cache,
271
+ self.is_neox_style,
272
+ )
273
+ return query, key
274
+
275
def extra_repr(self) -> str:
276
s = f"head_size={self.head_size}, rotary_dim={self.rotary_dim}"
277
s += f", max_position_embeddings={self.max_position_embeddings}"