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

sgl-kernel/python/sgl_kernel/__init__.py

@@ -30,6 +30,7 @@ from sgl_kernel.elementwise import (
     gemma_fused_add_rmsnorm,
     gemma_rmsnorm,
     rmsnorm,
+    rotary_embedding,
     silu_and_mul,
 )
 from sgl_kernel.expert_specialization import (

sgl-kernel/python/sgl_kernel/elementwise.py

@@ -353,6 +353,19 @@ def apply_rope_with_cos_sin_cache_inplace(
     )
 
 
+def rotary_embedding(
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    head_size: int,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool = True,
+):
+    torch.ops.sgl_kernel.rotary_embedding.default(
+        positions, query, key, head_size, cos_sin_cache, is_neox
+    )
+
+
 def downcast_fp8(
     k: torch.Tensor,
     v: torch.Tensor,

sgl-kernel/python/sgl_kernel/testing/rotary_embedding.py

@@ -97,10 +97,6 @@ class RotaryEmbedding(torch.nn.Module):
         num_tokens = positions.shape[0]
         cos_sin = self.cos_sin_cache.index_select(0, positions)
 
-        # Modification: float32 is required for the rotary embedding to work correctly
-        query = query.to(torch.float32)
-        key = key.to(torch.float32)
-
         cos, sin = cos_sin.chunk(2, dim=-1)
 
         query_shape = query.shape
@@ -146,6 +142,31 @@ class FlashInferRotaryEmbedding(RotaryEmbedding):
         return query, key
 
 
+class SglKernelRotaryEmbedding(RotaryEmbedding):
+    def forward_cuda(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        offsets: Optional[torch.Tensor] = None,
+        fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        assert (
+            fused_set_kv_buffer_arg is None
+        ), "fused_set_kv_buffer_arg is not supported for sgl-kernel implementation"
+        if self.cos_sin_cache.dtype != query.dtype:
+            self.cos_sin_cache = self.cos_sin_cache.to(query.dtype)
+        torch.ops.sgl_kernel.rotary_embedding(
+            positions,
+            query,
+            key,
+            self.head_size,
+            self.cos_sin_cache,
+            self.is_neox_style,
+        )
+        return query, key
+
+
 class MHATokenToKVPool:
     KV_POOL_SIZE = 16384