From c11b34d599892dababef84e586722ba96cea06b4 Mon Sep 17 00:00:00 2001 From: Chunyuan WU Date: Tue, 28 Oct 2025 01:02:18 +0800 Subject: [PATCH] rope xpu: fix missing argument 'fused_set_kv_buffer_arg' and replace native with sgl_kernel_xpu impl (#12006) --- python/sglang/srt/layers/rotary_embedding.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index 48f564f83..d48dc1b88 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -312,10 +312,20 @@ class RotaryEmbedding(CustomOp): query: torch.Tensor, key: torch.Tensor, offsets: Optional[torch.Tensor] = None, + fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - # TODO: make a wrapper, and XPU will implement this kernel later. - self.cos_sin_cache = self.cos_sin_cache.to(query.device) - return self.forward_native(positions, query, key, offsets) + assert ( + fused_set_kv_buffer_arg is None + ), "fused_set_kv_buffer_arg is not supported for xpu implementation" + positions = torch.add(positions, offsets) if offsets is not None else positions + return torch.ops.sgl_kernel.rotary_embedding( + positions, + query, + key, + self.head_size, + self.cos_sin_cache, + self.is_neox_style, + ) class LinearScalingRotaryEmbedding(RotaryEmbedding):