[2/n jit_kernel restruct] unify rotary embedding entrypoints under rope.py (#20247)

python/sglang/jit_kernel/rope.py

@@ -17,6 +17,15 @@ if TYPE_CHECKING:
     from tvm_ffi.module import Module
 
 
+@cache_once
+def _jit_rotary_embedding_module() -> Module:
+    return load_jit(
+        "rotary_embedding",
+        cuda_files=["elementwise/pos_enc.cuh"],
+        cuda_wrappers=[("rotary_embedding", "RotaryEmbeddingKernel::run")],
+    )
+
+
 @cache_once
 def _jit_fused_rope_module(is_neox: bool, rope_dim: int, dtype: torch.dtype) -> Module:
     args = make_cpp_args(is_neox, rope_dim, is_arch_support_pdl(), dtype)
@@ -31,6 +40,56 @@ def _jit_fused_rope_module(is_neox: bool, rope_dim: int, dtype: torch.dtype) ->
     )
 
 
+@register_custom_op(
+    op_name="rotary_embedding_with_key",
+    mutates_args=["query", "key"],
+)
+def rotary_embedding_with_key(
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    head_size: int,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool = True,
+) -> None:
+    module = _jit_rotary_embedding_module()
+    module.rotary_embedding(positions, query, key, head_size, cos_sin_cache, is_neox)
+
+
+@register_custom_op(
+    op_name="rotary_embedding_without_key",
+    mutates_args=["query"],
+)
+def rotary_embedding_without_key(
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    head_size: int,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool = True,
+) -> None:
+    module = _jit_rotary_embedding_module()
+    module.rotary_embedding(positions, query, None, head_size, cos_sin_cache, is_neox)
+
+
+def rotary_embedding(
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    key: Optional[torch.Tensor],
+    head_size: int,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool = True,
+):
+    if key is None:
+        rotary_embedding_without_key(
+            positions, query, head_size, cos_sin_cache, is_neox
+        )
+    else:
+        rotary_embedding_with_key(
+            positions, query, key, head_size, cos_sin_cache, is_neox
+        )
+    return query, key
+
+
 @dataclass
 class FusedSetKVBufferArg:
     """