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 3, 20:15 UTC
[CPU] Fix rotary embedding for CPU without flash-attn ops (#42225)
Signed-off-by: jmamou <jonathan.mamou@intel.com> Signed-off-by: Jonathan Mamou <jonathan.mamou@intel.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Jonathan Mamou committed
May 12, 2026 at 18:05 UTC
c8a6e272e0d3994f299ae522975514fcd98b7c76
1 file changed
+2
-1
vllm/model_executor/layers/rotary_embedding/common.py
+2
-1
@@ -8,6 +8,7 @@ import torch
8
9
from vllm.logger import init_logger
10
from vllm.model_executor.custom_op import CustomOp
11
+from vllm.platforms import current_platform
12
from vllm.utils.torch_utils import direct_register_custom_op
13
14
logger = init_logger(__name__)
@@ -134,7 +135,7 @@ class ApplyRotaryEmb(CustomOp):
135
self.enable_fp32_compute = enable_fp32_compute
136
137
self.apply_rotary_emb_flash_attn = None
137
- if find_spec("flash_attn") is not None:
138
+ if not current_platform.is_cpu() and find_spec("flash_attn") is not None:
139
from flash_attn.ops.triton.rotary import apply_rotary
140
141
self.apply_rotary_emb_flash_attn = apply_rotary