[kernel slimming] Clean many useless sgl-kernel deprecated kernels (#20277)

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

@@ -1,7 +1,31 @@
+from dataclasses import dataclass
 from typing import Optional, Tuple, Union
 
 import torch
-from sgl_kernel import FusedSetKVBufferArg, apply_rope_with_cos_sin_cache_inplace
+
+from sglang.jit_kernel.rope import FusedSetKVBufferArg as _JitFusedSetKVBufferArg
+from sglang.jit_kernel.rope import (
+    apply_rope_with_cos_sin_cache_inplace as _jit_apply_rope_with_cos_sin_cache_inplace,
+)
+
+
+@dataclass
+class FusedSetKVBufferArg:
+    value: torch.Tensor
+    k_buffer: torch.Tensor
+    v_buffer: torch.Tensor
+    cache_loc: torch.Tensor
+    # Kept for backward compatibility with old sgl_kernel test/bench callsites.
+    k_scale: Optional[float] = None
+    v_scale: Optional[float] = None
+
+    def to_jit(self) -> _JitFusedSetKVBufferArg:
+        return _JitFusedSetKVBufferArg(
+            value=self.value,
+            k_buffer=self.k_buffer,
+            v_buffer=self.v_buffer,
+            cache_loc=self.cache_loc,
+        )
 
 
 # vLLM torch native
@@ -129,14 +153,19 @@ class FlashInferRotaryEmbedding(RotaryEmbedding):
         fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
 
-        apply_rope_with_cos_sin_cache_inplace(
-            positions=positions,
-            query=query,
-            key=key,
-            fused_set_kv_buffer_arg=fused_set_kv_buffer_arg,
-            head_size=self.head_size,
+        query_view = query.view(query.shape[0], -1, self.head_size)
+        key_view = key.view(key.shape[0], -1, self.head_size)
+        _jit_apply_rope_with_cos_sin_cache_inplace(
+            q=query_view,
+            k=key_view,
             cos_sin_cache=self.cos_sin_cache,
+            positions=positions,
             is_neox=self.is_neox_style,
+            fused_args=(
+                fused_set_kv_buffer_arg.to_jit()
+                if fused_set_kv_buffer_arg is not None
+                else None
+            ),
         )
 
         return query, key