From 51d9d34977582dc565c784e4435cf23cf3ea0640 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com> Date: Tue, 10 Mar 2026 17:49:57 +0800 Subject: [PATCH] [2/n jit_kernel restruct] unify rotary embedding entrypoints under rope.py (#20247) --- .../sglang/jit_kernel/benchmark/bench_rope.py | 2 +- python/sglang/jit_kernel/pos_enc.py | 86 ------------------- python/sglang/jit_kernel/rope.py | 59 +++++++++++++ .../sglang/jit_kernel/tests/test_pos_enc.py | 2 +- .../srt/layers/rotary_embedding/base.py | 4 +- 5 files changed, 63 insertions(+), 90 deletions(-) delete mode 100644 python/sglang/jit_kernel/pos_enc.py diff --git a/python/sglang/jit_kernel/benchmark/bench_rope.py b/python/sglang/jit_kernel/benchmark/bench_rope.py index 1e8d324e3..74c4f1003 100644 --- a/python/sglang/jit_kernel/benchmark/bench_rope.py +++ b/python/sglang/jit_kernel/benchmark/bench_rope.py @@ -70,7 +70,7 @@ def sglang_pos_enc_rope( positions: torch.Tensor, is_neox: bool, ) -> None: - from sglang.jit_kernel.pos_enc import rotary_embedding_with_key + from sglang.jit_kernel.rope import rotary_embedding_with_key head_size = q.shape[-1] rotary_embedding_with_key( diff --git a/python/sglang/jit_kernel/pos_enc.py b/python/sglang/jit_kernel/pos_enc.py deleted file mode 100644 index 16f8ac37e..000000000 --- a/python/sglang/jit_kernel/pos_enc.py +++ /dev/null @@ -1,86 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING - -import torch - -from sglang.jit_kernel.utils import cache_once, load_jit -from sglang.srt.utils.custom_op import register_custom_op - -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")], - ) - - -@register_custom_op( - op_name="rotary_embedding_with_key", - mutates_args=["query", "key"], -) -def rotary_embedding_with_key( - positions: torch.Tensor, # [batch_size, seq_len] or [num_tokens] - query: torch.Tensor, # [batch_size, seq_len, num_heads * head_size] or - # [num_tokens, num_heads * head_size] or - # [batch_size, seq_len, num_heads, head_size] or - # [num_tokens, num_heads, head_size] - key: torch.Tensor, # [batch_size, seq_len, num_kv_heads * head_size] or - # [num_tokens, num_kv_heads * head_size] or - # [batch_size, seq_len, num_heads, head_size] or - # [num_tokens, num_heads, head_size] - head_size: int, - cos_sin_cache: torch.Tensor, # [max_position, rot_dim] - is_neox: bool = True, -) -> None: - """ - Apply rotary embedding to query and key tensors. - - Args: - positions: Position indices of shape [num_tokens] or [batch_size, seq_len] - query: Query tensor of shape [num_tokens, num_heads, head_size] or [num_tokens, num_heads * head_size] - key: Key tensor of shape [num_tokens, num_kv_heads, head_size] or [num_tokens, num_kv_heads * head_size] - cos_sin_cache: Cosine and sine cache of shape [max_position, rot_dim] - is_neox: Whether to use GPT-NeoX style rotary embedding (True) or GPT-J style (False) - """ - 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: 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 diff --git a/python/sglang/jit_kernel/rope.py b/python/sglang/jit_kernel/rope.py index ebc316094..d9cbe0a8b 100644 --- a/python/sglang/jit_kernel/rope.py +++ b/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: """ diff --git a/python/sglang/jit_kernel/tests/test_pos_enc.py b/python/sglang/jit_kernel/tests/test_pos_enc.py index 4b809b002..fc973f9b9 100644 --- a/python/sglang/jit_kernel/tests/test_pos_enc.py +++ b/python/sglang/jit_kernel/tests/test_pos_enc.py @@ -6,7 +6,7 @@ import torch import triton import triton.language as tl -from sglang.jit_kernel.pos_enc import rotary_embedding +from sglang.jit_kernel.rope import rotary_embedding @triton.jit diff --git a/python/sglang/srt/layers/rotary_embedding/base.py b/python/sglang/srt/layers/rotary_embedding/base.py index fa3068b3d..e98007325 100644 --- a/python/sglang/srt/layers/rotary_embedding/base.py +++ b/python/sglang/srt/layers/rotary_embedding/base.py @@ -71,10 +71,10 @@ class RotaryEmbedding(MultiPlatformOp): and not (_is_npu) and not (_is_musa) ): - # rotary_embedding from sglang.jit_kernel.pos_enc and vllm._custom_ops has the same implementation. + # rotary_embedding from sglang.jit_kernel.rope and vllm._custom_ops has the same implementation. # TODO: Test on different devices and remove this conditional. if _is_cuda: - from sglang.jit_kernel.pos_enc import rotary_embedding + from sglang.jit_kernel.rope import rotary_embedding elif _is_hip: from sgl_kernel import rotary_embedding else: