Skip to content

Commit d4ee86d

Browse files
David9857Yikun
authored andcommitted
feat: replace _npu_rotary_embedding with npu_mrope
Signed-off-by: David9857 <985700846@qq.com>
1 parent ab5d110 commit d4ee86d

File tree

1 file changed

+4
-3
lines changed

1 file changed

+4
-3
lines changed

vllm_ascend/ops/rotary_embedding.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -64,13 +64,14 @@ def rope_forward_oot(
6464
# TODO: Remove the contiguous in the future.
6565
query = query.contiguous().view(query.shape[0], -1)
6666
key = key.contiguous().view(key.shape[0], -1)
67-
torch_npu._npu_rotary_embedding(
67+
query, key = torch_npu.npu_mrope(
6868
positions,
6969
query,
7070
key,
71-
self.head_size,
7271
self.cos_sin_cache,
73-
neox_style,
72+
self.head_size,
73+
mrope_section=[0,0,0],
74+
rotary_mode='half' if neox_style else 'interleave'
7475
)
7576
return query.view(query_shape), key.view(key_shape)
7677

0 commit comments

Comments
 (0)