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, 19:45 UTC
[Attention] Remove slow setattr in MLA (#14769)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Mar 13, 2025 at 17:31 UTC
d47807ba0806c5bbd8fd08c19013c327b34dcac5
1 file changed
+7
-2
vllm/model_executor/layers/rotary_embedding.py
+7
-2
@@ -161,8 +161,13 @@ class RotaryEmbedding(CustomOp):
161
) -> Tuple[torch.Tensor, torch.Tensor]:
162
from vllm import _custom_ops as ops
163
164
- self.cos_sin_cache = self.cos_sin_cache.to(query.device,
165
- dtype=query.dtype)
164
+ # __setattr__ in nn.Module (called by `self.cos_sin_cache = ...`)
165
+ # is expensive, so avoid calling it if possible
166
+ if self.cos_sin_cache.device != query.device or \
167
+ self.cos_sin_cache.dtype != query.dtype:
168
+ self.cos_sin_cache = self.cos_sin_cache.to(query.device,
169
+ dtype=query.dtype)
170
+
171
# ops.rotary_embedding()/batched_rotary_embedding()
172
# are in-place operations that update the query and key tensors.
173
if offsets is not None: