EngineX/xc-llm-ascend
3
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

adopt rope in vllm-ascend (#530)

Browse Source 
### What this PR does / why we need it?
Adopt custom kernel rotary embedding in actual model inference,
customized rotary_embedding will generate contiguous query and key in
the cpp side to reduce the overhead of two contiguous and index_select
compared with rotary_embedding in torch_npu. For now, rotary_embedding
can only support the scenario of `is_neox = true`, non-neox version rope
will be updated soon in the future.
---------

Signed-off-by: ganyi <pleaplusone.gy@gmail.com>
This commit is contained in:
Pleaplusone
2025-04-18 08:56:05 +08:00
committed by GitHub
parent 23f85e3f74
commit 66a0837963
5 changed files with 37 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
tests/ops/test_rotary_embedding.py
View File

@@ -184,7 +184,7 @@ def test_rotary_embedding_quant_with_leading_dim(
)
ref_query, ref_key = rope.forward_native(positions, query, key)
torch.ops._C.rotary_embedding(
query, key = torch.ops._C.rotary_embedding(
positions,
query,
key,
@@ -194,11 +194,11 @@ def test_rotary_embedding_quant_with_leading_dim(
)
# Compare the results.
torch.testing.assert_close(query,
torch.testing.assert_close(query.view(ref_query.size()),
ref_query,
atol=DEFAULT_ATOL,
rtol=DEFAULT_RTOL)
torch.testing.assert_close(key,
torch.testing.assert_close(key.view(ref_key.size()),
ref_key,
atol=DEFAULT_ATOL,
rtol=DEFAULT_RTOL)