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

[2/2] Add rope kernel in sgl-kernel (#14452)

Browse Source
This commit is contained in:
Qiaolin Yu
2025-12-07 21:37:29 -08:00
committed by GitHub
parent 2970f22917
commit 661e9775d0
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/sglang/srt/layers/rotary_embedding.py
View File

@@ -118,7 +118,10 @@ class RotaryEmbedding(CustomOp):
and not (_is_cpu and _is_cpu_amx_available)
and not (_is_xpu)
):
from vllm._custom_ops import rotary_embedding
if _is_cuda:
from sgl_kernel import rotary_embedding
else:
from vllm._custom_ops import rotary_embedding
self.use_fallback_kernel = True
self.fallback_rotary_embedding = rotary_embedding
@@ -340,7 +343,7 @@ class RotaryEmbedding(CustomOp):
else:
assert (
fused_set_kv_buffer_arg is None
), "save kv cache is not supported for vllm_rotary_embedding."
), "save kv cache is not supported for fallback_rotary_embedding."
self.cos_sin_cache = self.cos_sin_cache.to(query.device, dtype=query.dtype)
self.fallback_rotary_embedding(
positions,