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 15, 12:15 UTC
Make the _apply_rotary_emb compatible with dynamo (#17435)
Lu Fang committed
Apr 30, 2025 at 00:52 UTC
ece5a8b0b6657b8c46b20daba2869762d1765008
1 file changed
+3
-1
vllm/model_executor/layers/rotary_embedding.py
+3
-1
@@ -32,6 +32,9 @@ from transformers import PretrainedConfig
32
from vllm.model_executor.custom_op import CustomOp
33
from vllm.platforms import current_platform
34
35
+if current_platform.is_cuda_alike():
36
+ from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
37
+
38
39
def _rotate_neox(x: torch.Tensor) -> torch.Tensor:
40
x1 = x[..., :x.shape[-1] // 2]
@@ -78,7 +81,6 @@ def _apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
81
positional embeddings.
82
"""
83
if current_platform.is_cuda_alike():
81
- from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
84
return apply_rotary_emb(x.unsqueeze(0), cos, sin,
85
not is_neox_style).squeeze(0)
86
else: