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
[Bugfix] Add contiguous call inside rope kernel wrapper (#17091)
Signed-off-by: 苏政渊 <suzhengyuan@moonshot.cn> Co-authored-by: 苏政渊 <suzhengyuan@moonshot.cn>
Zhengyuan Su (苏政渊) committed
Apr 29, 2025 at 10:24 UTC
17eb306fcc7018f01e37d6bb548c7cdfc65602a6
2 files changed
+17
-7
vllm/_custom_ops.py
+14
-3
@@ -158,8 +158,13 @@ def rotary_embedding(
158
cos_sin_cache: torch.Tensor,
159
is_neox: bool,
160
) -> None:
161
- torch.ops._C.rotary_embedding(positions, query, key, head_size,
162
- cos_sin_cache, is_neox)
161
+ # TODO: Remove this contiguous call when the kernel is updated to support tensor slices
162
+ query_contiguous = query.contiguous()
163
+ key_contiguous = key.contiguous()
164
+ torch.ops._C.rotary_embedding(positions, query_contiguous, key_contiguous,
165
+ head_size, cos_sin_cache, is_neox)
166
+ query.copy_(query_contiguous)
167
+ key.copy_(key_contiguous)
168
169
170
def batched_rotary_embedding(positions: torch.Tensor, query: torch.Tensor,
@@ -167,9 +172,15 @@ def batched_rotary_embedding(positions: torch.Tensor, query: torch.Tensor,
172
cos_sin_cache: torch.Tensor, is_neox: bool,
173
rot_dim: int,
174
cos_sin_cache_offsets: torch.Tensor) -> None:
170
- torch.ops._C.batched_rotary_embedding(positions, query, key, head_size,
175
+ # TODO: Remove this contiguous call when the kernel is updated to support tensor slices
176
+ query_contiguous = query.contiguous()
177
+ key_contiguous = key.contiguous()
178
+ torch.ops._C.batched_rotary_embedding(positions, query_contiguous,
179
+ key_contiguous, head_size,
180
cos_sin_cache, is_neox, rot_dim,
181
cos_sin_cache_offsets)
182
+ query.copy_(query_contiguous)
183
+ key.copy_(key_contiguous)
184
185
186
# layer norm ops
vllm/v1/attention/backends/mla/common.py
+3
-4
@@ -938,8 +938,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
938
decode_ql_nope, decode_q_pe = \
939
self._q_proj_and_k_up_proj(decode_hs_or_q_c)
940
decode_q_pe[...], decode_k_pe[...] = self.rotary_emb(
941
- attn_metadata.decode.input_positions, decode_q_pe.contiguous(),
942
- decode_k_pe)
941
+ attn_metadata.decode.input_positions, decode_q_pe, decode_k_pe)
942
943
if has_prefill:
944
assert attn_metadata.prefill is not None
@@ -948,8 +947,8 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
947
prefill_q_pe = prefill_q[..., self.qk_nope_head_dim:]
948
949
prefill_q_pe[...], prefill_k_pe[...] = self.rotary_emb(
951
- attn_metadata.prefill.input_positions,
952
- prefill_q_pe.contiguous(), prefill_k_pe)
950
+ attn_metadata.prefill.input_positions, prefill_q_pe,
951
+ prefill_k_pe)
952
953
# write the latent and rope to kv cache
954
if kv_cache.numel() > 0: