From 661e9775d0c2b0b5bc479951ed3493be7a5e6bdf Mon Sep 17 00:00:00 2001 From: Qiaolin Yu Date: Sun, 7 Dec 2025 21:37:29 -0800 Subject: [PATCH] [2/2] Add rope kernel in sgl-kernel (#14452) --- python/sglang/srt/layers/rotary_embedding.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index 0cdb7e1ae..0c9dc2f3a 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -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,