chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Feature] rewrite rope kernel; remove flashinfer dependencies (#18844)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
DarkSharpness
2026-02-21 21:32:40 +08:00
committed by GitHub
parent 6503f94211
commit d8d0208c63
9 changed files with 1155 additions and 1107 deletions
Download Patch File Download Diff File Expand all files Collapse all files

298
python/sglang/jit_kernel/rope.py
View File

@@ -1,13 +1,16 @@
from __future__ import annotations
import pathlib
from dataclasses import dataclass
from typing import TYPE_CHECKING, Optional
import flashinfer
import torch
from sglang.jit_kernel.utils import cache_once, is_arch_support_pdl, load_jit
from sglang.jit_kernel.utils import (
cache_once,
is_arch_support_pdl,
load_jit,
make_cpp_args,
)
from sglang.srt.utils.custom_op import register_custom_op
if TYPE_CHECKING:
@@ -15,115 +18,19 @@ if TYPE_CHECKING:
@cache_once
def _jit_apply_rope_pos_ids_cos_sin_cache_module() -> Module:
flashinfer_dir = pathlib.Path(flashinfer.__file__).parent.resolve()
assert (
flashinfer_dir / "data" / "include"
).exists(), (
f"flashinfer headers are missing {str(flashinfer_dir / 'data' / 'include')}"
)
flashinfer_include_path = (flashinfer_dir / "data" / "include").resolve()
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)
return load_jit(
"apply_rope_pos_ids_cos_sin_cache",
"fused_rope",
*args,
cuda_files=["elementwise/rope.cuh"],
cuda_wrappers=[
(
"apply_rope_pos_ids_cos_sin_cache",
"ApplyRopePosIdsCosSinCacheKernel::run",
)
("run_rope", f"FusedRopeKernel<{args}>::run"),
("run_rope_store", f"FusedRopeKernel<{args}>::run_fused"),
],
extra_include_paths=[str(flashinfer_include_path)],
)
# Split the ops because k_buffer/v_buffer are mutated only when provided,
# and torch.custom_op cannot express optional mutates_args reliably
@register_custom_op(
op_name="apply_rope_pos_ids_cos_sin_cache_with_kv_cache",
mutates_args=["q", "k", "q_rope", "k_rope", "k_buffer", "v_buffer"],
)
def apply_rope_pos_ids_cos_sin_cache_with_kv_cache(
q: torch.Tensor,
k: torch.Tensor,
q_rope: torch.Tensor,
k_rope: torch.Tensor,
cos_sin_cache: torch.Tensor,
pos_ids: torch.Tensor,
v: torch.Tensor,
k_buffer: torch.Tensor,
v_buffer: torch.Tensor,
kv_cache_loc: torch.Tensor,
interleave: bool = False,
enable_pdl: bool = False,
) -> None:
"""
Apply RoPE (Rotary Positional Embedding) with position IDs and cos/sin cache.
Args:
q: Input Q tensor of shape [nnz, num_qo_heads, head_dim]
k: Input K tensor of shape [nnz, num_kv_heads, head_dim]
q_rope: Output Q tensor with RoPE applied, same shape as q
k_rope: Output K tensor with RoPE applied, same shape as k
cos_sin_cache: Cos/sin cache of shape [max_seq_len, rotary_dim]
pos_ids: Position IDs of shape [nnz]
interleave: Whether to use interleaved RoPE
enable_pdl: Enable PDL (Programmable Data Layout)
v: Optional V tensor for KV caching
k_buffer: Optional K buffer for KV caching
v_buffer: Optional V buffer for KV caching
kv_cache_loc: Optional KV cache location tensor
"""
module = _jit_apply_rope_pos_ids_cos_sin_cache_module()
module.apply_rope_pos_ids_cos_sin_cache(
q,
k,
q_rope,
k_rope,
cos_sin_cache,
pos_ids,
interleave,
enable_pdl,
v,
k_buffer,
v_buffer,
kv_cache_loc,
)
@register_custom_op(
op_name="apply_rope_pos_ids_cos_sin_cache_without_kv_cache",
mutates_args=["q", "k", "q_rope", "k_rope"],
)
def apply_rope_pos_ids_cos_sin_cache_without_kv_cache(
q: torch.Tensor,
k: torch.Tensor,
q_rope: torch.Tensor,
k_rope: torch.Tensor,
cos_sin_cache: torch.Tensor,
pos_ids: torch.Tensor,
interleave: bool = False,
enable_pdl: bool = False,
) -> None:
module = _jit_apply_rope_pos_ids_cos_sin_cache_module()
module.apply_rope_pos_ids_cos_sin_cache(
q,
k,
q_rope,
k_rope,
cos_sin_cache,
pos_ids,
interleave,
enable_pdl,
None,
None,
None,
None,
)
# Adepted from
@dataclass
class FusedSetKVBufferArg:
"""
@@ -133,10 +40,6 @@ class FusedSetKVBufferArg:
Buffer for keys, shape: ``(nnz, num_k_heads * head_size)``.
v_buffer : Optional[torch.Tensor]
Buffer for values, shape: ``(nnz, num_v_heads * head_size)``.
k_scale : Optional[float]
Scale factor for keys.
v_scale : Optional[float]
Scale factor for values.
cache_loc : Optional[torch.Tensor]
Cache location tensor, used for indexing kv cache.
"""
@@ -144,93 +47,116 @@ class FusedSetKVBufferArg:
value: torch.Tensor
k_buffer: torch.Tensor
v_buffer: torch.Tensor
k_scale: Optional[float]
v_scale: Optional[float]
cache_loc: torch.Tensor
def _view_3d(x, head_size):
return x.view(x.shape[0], -1, head_size)
def apply_rope_with_cos_sin_cache_inplace(
positions: torch.Tensor,
query: torch.Tensor,
key: torch.Tensor,
head_size: int,
@register_custom_op(mutates_args=["q", "k"])
def apply_rope_inplace(
q: torch.Tensor,
k: torch.Tensor,
cos_sin_cache: torch.Tensor,
is_neox: bool = True,
fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
enable_pdl: Optional[bool] = None,
positions: torch.Tensor,
*,
is_neox: bool,
rope_dim: int = 0,
) -> None:
r"""
Apply rotary embedding to keys and queries with precomputed cos/sin values.
This is designed to be compatible with the SGL/vLLM implementation.
The result is inplace applied to the input tensors.
Parameters
----------
positions : torch.Tensor
Position indices, shape: ``(nnz)``.
query : torch.Tensor
Query tensor, shape: ``(nnz, num_q_heads * head_size)``.
key : torch.Tensor
Key tensor, shape: ``(nnz, num_k_heads * head_size)``.
cos_sin_cache : torch.Tensor
Cosine and Sine cache tensor, shape: ``(max_seq_len, rotary_dim)``.
Cosine is the first half and Sine is the second half on rotary_dim.
is_neox : bool
Whether to use Neox style RoPE, default: ``True``.
* If ``True``, the last dimension of the query/key tensor is not interleaved, i.e.,
we rotate the first half dimensions ``([..., :head_dim//2])`` and the second half
dimensions ``([..., head_dim//2:])``.
* If ``False``, the last dimension of the query/key tensor is interleaved, i.e.,
we rotate the even dimensions ``([..., ::2])`` and odd dimensions ``([..., 1::2])``.
fused_set_kv_buffer_arg : FusedSetKVBufferArg
Fuse the set-kv-buffer operation into this kernel
Note
----
The rotary dimension is determined by the cosine cache and sine cache.
"""
if cos_sin_cache.dtype != torch.float32:
raise ValueError("cos_sin_cache should be float32")
Fused inplace rotary position embedding for query and key tensors.
if enable_pdl is None:
# the non-fused branch does not yet support PDL, but after we switch to our impl for that branch it will
enable_pdl = is_arch_support_pdl() and (fused_set_kv_buffer_arg is not None)
Args:
q: Query tensor of shape [num_tokens, num_qo_heads, rope_dim].
k: Key tensor of shape [num_tokens, num_kv_heads, rope_dim].
cos_sin_cache: Cosine/sine cache of shape [max_position, rope_dim],
where the first half along dim=-1 is cos and the second half is sin.
Must be float32.
positions: Position indices of shape [num_tokens], int32 or int64.
is_neox: Whether to use GPT-NeoX style (True) or GPT-J interleaved style (False).
rope_dim: Rotary embedding dimension. Defaults to cos_sin_cache.size(-1).
"""
rope_dim = rope_dim or cos_sin_cache.size(-1)
module = _jit_fused_rope_module(is_neox, rope_dim, q.dtype)
module.run_rope(q, k, cos_sin_cache, positions)
if (a := fused_set_kv_buffer_arg) is not None:
assert a.k_scale is None, "k_scale is not yet supported"
assert a.v_scale is None, "v_scale is not yet supported"
assert a.cache_loc.dtype == torch.int64, f"{a.cache_loc.dtype=}"
save_kv_cache = fused_set_kv_buffer_arg is not None
if save_kv_cache:
apply_rope_pos_ids_cos_sin_cache_with_kv_cache(
_view_3d(query, head_size),
_view_3d(key, head_size),
_view_3d(query, head_size),
_view_3d(key, head_size),
@register_custom_op(mutates_args=["q", "k_cache", "v_cache"])
def apply_rope_inplace_with_kvcache(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
k_cache: torch.Tensor,
v_cache: torch.Tensor,
cos_sin_cache: torch.Tensor,
positions: torch.Tensor,
out_loc: torch.Tensor,
*,
is_neox: bool,
rope_dim: int = 0,
) -> None:
"""
Fused inplace RoPE + KV cache store.
Applies rotary position embedding to q inplace. For k, applies RoPE and
stores the result in k_cache. The original v is also stored in v_cache.
Args:
q: Query tensor of shape [num_tokens, num_qo_heads, head_dim].
k: Key tensor of shape [num_tokens, num_kv_heads, head_dim].
v: Value tensor of shape [num_tokens, num_kv_heads, head_dim].
k_cache: Key cache of shape [cache_size, num_kv_heads * head_dim].
v_cache: Value cache of shape [cache_size, num_kv_heads * head_dim].
cos_sin_cache: Cosine/sine cache of shape [max_position, rope_dim], float32.
positions: Position indices of shape [num_tokens], int32 or int64.
out_loc: Cache write locations of shape [num_tokens], same dtype as positions.
is_neox: Whether to use GPT-NeoX style (True) or GPT-J interleaved (False).
rope_dim: Rotary embedding dimension. Defaults to cos_sin_cache.size(-1).
"""
rope_dim = rope_dim or cos_sin_cache.size(-1)
v = v.view_as(k)
module = _jit_fused_rope_module(is_neox, rope_dim, q.dtype)
module.run_rope_store(q, k, v, k_cache, v_cache, cos_sin_cache, positions, out_loc)
# NOTE: this name is intentionally set as the old kernel in `sgl_kernel`
def apply_rope_with_cos_sin_cache_inplace(
q: torch.Tensor,
k: torch.Tensor,
cos_sin_cache: torch.Tensor,
positions: torch.Tensor,
*,
is_neox: bool,
rope_dim: int = 0,
fused_args: Optional[FusedSetKVBufferArg] = None,
) -> None:
"""
Apply RoPE to q and k inplace, with optional fused kv cache store.
If `fused_args` is provided, it will perform fused RoPE and KV cache store.
Otherwise, it will only apply RoPE inplace.
Args:
q: Query tensor of shape [num_tokens, num_qo_heads, head_dim].
k: Key tensor of shape [num_tokens, num_kv_heads, head_dim].
cos_sin_cache: Cosine/sine cache of shape [max_position, rope_dim], float32.
positions: Position indices of shape [num_tokens], int32 or int64.
is_neox: Whether to use GPT-NeoX style (True) or GPT-J interleaved (False).
rope_dim: Rotary embedding dimension. Defaults to cos_sin_cache.size(-1).
fused_args: Optional arguments for fused RoPE + KV cache store. If None,
only RoPE will be applied inplace without touching kv cache.
"""
if fused_args is not None:
apply_rope_inplace_with_kvcache(
q,
k,
fused_args.value,
fused_args.k_buffer,
fused_args.v_buffer,
cos_sin_cache,
positions.long(),
_view_3d(fused_set_kv_buffer_arg.value, head_size),
_view_3d(fused_set_kv_buffer_arg.k_buffer, head_size),
_view_3d(fused_set_kv_buffer_arg.v_buffer, head_size),
(fused_set_kv_buffer_arg.cache_loc),
(not is_neox),
enable_pdl,
positions,
fused_args.cache_loc,
is_neox=is_neox,
rope_dim=rope_dim,
)
else:
apply_rope_pos_ids_cos_sin_cache_without_kv_cache(
_view_3d(query, head_size),
_view_3d(key, head_size),
_view_3d(query, head_size),
_view_3d(key, head_size),
cos_sin_cache,
positions.long(),
(not is_neox),
enable_pdl,
apply_rope_inplace(
q, k, cos_sin_cache, positions, is_neox=is_neox, rope_dim=rope_dim
)