minke.yu/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
python/sglang/srt/layers/rotary_embedding.py
+5 -2
View File
@@ -118,7 +118,10 @@ class RotaryEmbedding(CustomOp):
and not (_is_cpu and _is_cpu_amx_available) and not (_is_cpu and _is_cpu_amx_available)
and not (_is_xpu) 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.use_fallback_kernel = True
self.fallback_rotary_embedding = rotary_embedding self.fallback_rotary_embedding = rotary_embedding
@@ -340,7 +343,7 @@ class RotaryEmbedding(CustomOp):
else: else:
assert ( assert (
fused_set_kv_buffer_arg is None 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.cos_sin_cache = self.cos_sin_cache.to(query.device, dtype=query.dtype)
self.fallback_rotary_embedding( self.fallback_rotary_embedding(
positions, positions,