minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

ROCm use rotary_embedding from sgl-kernel (#18920)

Browse Source
This commit is contained in:
HAI
2026-02-17 03:00:37 -08:00
committed by GitHub
parent 5f81ec1ad5
commit 8bb1037796
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/rotary_embedding.py
+3 -1
View File
@@ -131,8 +131,10 @@ class RotaryEmbedding(MultiPlatformOp):
): ):
# rotary_embedding from sglang.jit_kernel.pos_enc and vllm._custom_ops has the same implementation. # rotary_embedding from sglang.jit_kernel.pos_enc and vllm._custom_ops has the same implementation.
# TODO: Test on different devices and remove this conditional. # TODO: Test on different devices and remove this conditional.
if _is_cuda or _is_hip: if _is_cuda:
from sglang.jit_kernel.pos_enc import rotary_embedding from sglang.jit_kernel.pos_enc import rotary_embedding
elif _is_hip:
from sgl_kernel import rotary_embedding
else: else:
from vllm._custom_ops import rotary_embedding from vllm._custom_ops import rotary_embedding