chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[2/n jit_kernel restruct] unify rotary embedding entrypoints under rope.py (#20247)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-03-10 17:49:57 +08:00
committed by GitHub
parent 6407891b4f
commit 51d9d34977
5 changed files with 63 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/layers/rotary_embedding/base.py
View File

@@ -71,10 +71,10 @@ class RotaryEmbedding(MultiPlatformOp):
and not (_is_npu)
and not (_is_musa)
):
# rotary_embedding from sglang.jit_kernel.pos_enc and vllm._custom_ops has the same implementation.
# rotary_embedding from sglang.jit_kernel.rope and vllm._custom_ops has the same implementation.
# TODO: Test on different devices and remove this conditional.
if _is_cuda:
from sglang.jit_kernel.pos_enc import rotary_embedding
from sglang.jit_kernel.rope import rotary_embedding
elif _is_hip:
from sgl_kernel import rotary_embedding
else: