Skip to content

Commit e306666

Browse files
committed
remove .clone to optimize the use of AITER RoPE
Signed-off-by: tjtanaavllm <tunjian.tan@amd.com>
1 parent fdbfe6d commit e306666

1 file changed

Lines changed: 4 additions & 6 deletions

File tree

vllm/v1/attention/backends/mla/common.py

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -952,9 +952,8 @@ def forward(
952952
self._q_proj_and_k_up_proj(decode_hs_or_q_c)
953953

954954
if self.use_rocm_aiter:
955-
decode_q_pe[...], decode_k_pe[...] = self.rotary_emb(
956-
attn_metadata.decode.input_positions,
957-
decode_q_pe.clone().contiguous(), decode_k_pe.clone())
955+
self.rotary_emb(attn_metadata.decode.input_positions,
956+
decode_q_pe, decode_k_pe)
958957
else:
959958
decode_q_pe[...], decode_k_pe[...] = self.rotary_emb(
960959
attn_metadata.decode.input_positions,
@@ -967,9 +966,8 @@ def forward(
967966
prefill_q_pe = prefill_q[..., self.qk_nope_head_dim:]
968967

969968
if self.use_rocm_aiter:
970-
prefill_q_pe[...], prefill_k_pe[...] = self.rotary_emb(
971-
attn_metadata.prefill.input_positions,
972-
prefill_q_pe.clone().contiguous(), prefill_k_pe.clone())
969+
self.rotary_emb(attn_metadata.prefill.input_positions,
970+
prefill_q_pe, prefill_k_pe)
973971
else:
974972
prefill_q_pe[...], prefill_k_pe[...] = self.rotary_emb(
975973
attn_metadata.prefill.input_positions,

0 commit comments

Comments
 (0)