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

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