[Feature] rewrite rope kernel; remove flashinfer dependencies (#18844)
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
DarkSharpness
350
python/sglang/jit_kernel/benchmark/bench_rope.py
Normal file
350
python/sglang/jit_kernel/benchmark/bench_rope.py
Normal file
@@ -0,0 +1,350 @@
|
||||
import itertools
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import (
|
||||
DEFAULT_DEVICE,
|
||||
DEFAULT_DTYPE,
|
||||
get_benchmark_range,
|
||||
run_benchmark,
|
||||
)
|
||||
|
||||
MAX_SEQ_LEN = 131072
|
||||
ROPE_BASE = 10000.0
|
||||
ROPE_DIM = 128
|
||||
CACHE_SIZE = 1024 * 1024
|
||||
|
||||
|
||||
def create_cos_sin_cache(
|
||||
rotary_dim: int = ROPE_DIM,
|
||||
max_position: int = MAX_SEQ_LEN,
|
||||
base: float = ROPE_BASE,
|
||||
) -> torch.Tensor:
|
||||
inv_freq = 1.0 / (
|
||||
base
|
||||
** (
|
||||
torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=DEFAULT_DEVICE)
|
||||
/ rotary_dim
|
||||
)
|
||||
)
|
||||
t = torch.arange(max_position, dtype=torch.float32, device=DEFAULT_DEVICE)
|
||||
freqs = torch.einsum("i,j->ij", t, inv_freq)
|
||||
cos = freqs.cos()
|
||||
sin = freqs.sin()
|
||||
return torch.cat((cos, sin), dim=-1)
|
||||
|
||||
|
||||
# Pre-build the cache once
|
||||
COS_SIN_CACHE = create_cos_sin_cache()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RoPE-only provider implementations
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def flashinfer_rope(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
|
||||
|
||||
head_size = q.shape[-1]
|
||||
apply_rope_with_cos_sin_cache_inplace(
|
||||
positions=positions,
|
||||
query=q.view(q.shape[0], -1),
|
||||
key=k.view(k.shape[0], -1),
|
||||
head_size=head_size,
|
||||
cos_sin_cache=COS_SIN_CACHE,
|
||||
is_neox=is_neox,
|
||||
)
|
||||
|
||||
|
||||
def sglang_rope_v0(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
from sglang.jit_kernel.pos_enc import rotary_embedding_with_key
|
||||
|
||||
head_size = q.shape[-1]
|
||||
rotary_embedding_with_key(
|
||||
positions=positions,
|
||||
query=q.view(q.shape[0], -1),
|
||||
key=k.view(k.shape[0], -1),
|
||||
head_size=head_size,
|
||||
cos_sin_cache=COS_SIN_CACHE,
|
||||
is_neox=is_neox,
|
||||
)
|
||||
|
||||
|
||||
def sglang_rope_v1(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
from sgl_kernel import apply_rope_with_cos_sin_cache_inplace
|
||||
|
||||
head_size = q.shape[-1]
|
||||
apply_rope_with_cos_sin_cache_inplace(
|
||||
positions=positions,
|
||||
query=q.view(q.shape[0], -1),
|
||||
key=k.view(k.shape[0], -1),
|
||||
head_size=head_size,
|
||||
cos_sin_cache=COS_SIN_CACHE,
|
||||
is_neox=is_neox,
|
||||
)
|
||||
|
||||
|
||||
def sglang_rope_v2(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
from sglang.jit_kernel.rope import apply_rope_inplace
|
||||
|
||||
apply_rope_inplace(q, k, COS_SIN_CACHE, positions, is_neox=is_neox)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RoPE + KV cache store provider implementations
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def rope_v0_store(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
k_cache: torch.Tensor,
|
||||
v_cache: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
out_loc: torch.Tensor,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
from sglang.jit_kernel.kvcache import store_cache
|
||||
from sglang.jit_kernel.rope import apply_rope_inplace
|
||||
|
||||
head_size = q.shape[-1]
|
||||
row_dim = k.shape[-2] * head_size
|
||||
apply_rope_inplace(
|
||||
positions=positions,
|
||||
q=q,
|
||||
k=k,
|
||||
rope_dim=head_size,
|
||||
cos_sin_cache=COS_SIN_CACHE,
|
||||
is_neox=is_neox,
|
||||
)
|
||||
store_cache(
|
||||
k.view(-1, row_dim),
|
||||
v.view(-1, row_dim),
|
||||
k_cache,
|
||||
v_cache,
|
||||
out_loc,
|
||||
)
|
||||
|
||||
|
||||
def rope_v1_store(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
k_cache: torch.Tensor,
|
||||
v_cache: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
out_loc: torch.Tensor,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
from sgl_kernel import FusedSetKVBufferArg, apply_rope_with_cos_sin_cache_inplace
|
||||
|
||||
head_size = q.shape[-1]
|
||||
apply_rope_with_cos_sin_cache_inplace(
|
||||
positions=positions,
|
||||
query=q.view(q.shape[0], -1),
|
||||
key=k.view(k.shape[0], -1),
|
||||
head_size=head_size,
|
||||
cos_sin_cache=COS_SIN_CACHE,
|
||||
is_neox=is_neox,
|
||||
fused_set_kv_buffer_arg=FusedSetKVBufferArg(
|
||||
value=v.view(v.shape[0], -1),
|
||||
k_buffer=k_cache,
|
||||
v_buffer=v_cache,
|
||||
k_scale=None,
|
||||
v_scale=None,
|
||||
cache_loc=out_loc,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def rope_v2_store(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
k_cache: torch.Tensor,
|
||||
v_cache: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
out_loc: torch.Tensor,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
from sglang.jit_kernel.rope import apply_rope_inplace_with_kvcache
|
||||
|
||||
apply_rope_inplace_with_kvcache(
|
||||
q, k, v, k_cache, v_cache, COS_SIN_CACHE, positions, out_loc, is_neox=is_neox
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark configuration (shared)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
BS_RANGE = get_benchmark_range(
|
||||
full_range=[2**n for n in range(0, 16)],
|
||||
ci_range=[16],
|
||||
)
|
||||
QK_HEAD_RANGE = get_benchmark_range(
|
||||
full_range=[(8, 1), (16, 2), (32, 8)],
|
||||
ci_range=[(16, 2)],
|
||||
)
|
||||
QK_HEAD_RANGE = [f"{q},{k}" for q, k in QK_HEAD_RANGE]
|
||||
IS_NEOX_RANGE = get_benchmark_range(
|
||||
full_range=[True, False],
|
||||
ci_range=[True],
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark 1: RoPE only
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
ROPE_LINE_VALS = ["fi", "rope_v0", "rope_v1", "rope_v2"]
|
||||
ROPE_LINE_NAMES = ["FlashInfer", "SGL RoPE v0", "SGL RoPE v1", "SGL RoPE v2"]
|
||||
ROPE_STYLES = [("green", "-."), ("red", "-"), ("orange", "-"), ("blue", "--")]
|
||||
|
||||
rope_configs = list(itertools.product(QK_HEAD_RANGE, IS_NEOX_RANGE, BS_RANGE))
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["num_q_k_heads", "is_neox", "batch_size"],
|
||||
x_vals=rope_configs,
|
||||
line_arg="provider",
|
||||
line_vals=ROPE_LINE_VALS,
|
||||
line_names=ROPE_LINE_NAMES,
|
||||
styles=ROPE_STYLES,
|
||||
ylabel="us",
|
||||
plot_name="rope-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(batch_size: int, num_q_k_heads: str, is_neox: bool, provider: str):
|
||||
qo, kv = num_q_k_heads.split(",")
|
||||
num_qo_heads = int(qo)
|
||||
num_kv_heads = int(kv)
|
||||
q = torch.randn(
|
||||
(batch_size, num_qo_heads, ROPE_DIM),
|
||||
dtype=DEFAULT_DTYPE,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
k = torch.randn(
|
||||
(batch_size, num_kv_heads, ROPE_DIM),
|
||||
dtype=DEFAULT_DTYPE,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
seed = batch_size << 16 | num_qo_heads << 8 | num_kv_heads << 4 | is_neox
|
||||
torch.random.manual_seed(seed)
|
||||
positions = torch.randint(
|
||||
MAX_SEQ_LEN, (batch_size,), device=DEFAULT_DEVICE, dtype=torch.int64
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
FN_MAP = {
|
||||
"fi": flashinfer_rope,
|
||||
"rope_v0": sglang_rope_v0,
|
||||
"rope_v1": sglang_rope_v1,
|
||||
"rope_v2": sglang_rope_v2,
|
||||
}
|
||||
fn = lambda: FN_MAP[provider](q, k, positions, is_neox)
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark 2: RoPE + KV cache store
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
STORE_LINE_VALS = ["rope_v0_store", "rope_v1_store", "rope_v2_store"]
|
||||
STORE_LINE_NAMES = ["SGL RoPE v0 + Store", "SGL RoPE v1 + Store", "SGL RoPE v2 + Store"]
|
||||
STORE_STYLES = [("red", "-"), ("orange", "-"), ("blue", "--")]
|
||||
|
||||
store_configs = list(itertools.product(QK_HEAD_RANGE, IS_NEOX_RANGE, BS_RANGE))
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["num_q_k_heads", "is_neox", "batch_size"],
|
||||
x_vals=store_configs,
|
||||
line_arg="provider",
|
||||
line_vals=STORE_LINE_VALS,
|
||||
line_names=STORE_LINE_NAMES,
|
||||
styles=STORE_STYLES,
|
||||
ylabel="us",
|
||||
plot_name="rope-store-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark_store(batch_size: int, num_q_k_heads: str, is_neox: bool, provider: str):
|
||||
qo, kv = num_q_k_heads.split(",")
|
||||
num_qo_heads = int(qo)
|
||||
num_kv_heads = int(kv)
|
||||
q = torch.randn(
|
||||
(batch_size, num_qo_heads, ROPE_DIM),
|
||||
dtype=DEFAULT_DTYPE,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
k = torch.randn(
|
||||
(batch_size, num_kv_heads, ROPE_DIM),
|
||||
dtype=DEFAULT_DTYPE,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
v = torch.randn(
|
||||
(batch_size, num_kv_heads, ROPE_DIM),
|
||||
dtype=DEFAULT_DTYPE,
|
||||
device=DEFAULT_DEVICE,
|
||||
)
|
||||
row_size = num_kv_heads * ROPE_DIM
|
||||
k_cache = torch.zeros(
|
||||
CACHE_SIZE, row_size, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE
|
||||
)
|
||||
v_cache = torch.zeros(
|
||||
CACHE_SIZE, row_size, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE
|
||||
)
|
||||
out_loc = torch.randperm(CACHE_SIZE, device=DEFAULT_DEVICE, dtype=torch.int64)[
|
||||
:batch_size
|
||||
]
|
||||
seed = batch_size << 16 | num_qo_heads << 8 | num_kv_heads << 4 | is_neox
|
||||
torch.random.manual_seed(seed)
|
||||
positions = torch.randint(
|
||||
MAX_SEQ_LEN, (batch_size,), device=DEFAULT_DEVICE, dtype=torch.int64
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
FN_MAP = {
|
||||
"rope_v0_store": rope_v0_store,
|
||||
"rope_v1_store": rope_v1_store,
|
||||
"rope_v2_store": rope_v2_store,
|
||||
}
|
||||
fn = lambda: FN_MAP[provider](
|
||||
q, k, v, k_cache, v_cache, positions, out_loc, is_neox
|
||||
)
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("Running RoPE performance benchmark...")
|
||||
benchmark.run(print_data=True)
|
||||
print("\nRunning RoPE + KV cache store performance benchmark...")
|
||||
benchmark_store.run(print_data=True)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -11,6 +11,7 @@
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_fp8.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <type_traits>
|
||||
|
||||
#ifndef USE_ROCM
|
||||
using fp32_t = float;
|
||||
@@ -62,6 +63,31 @@ SGL_DEVICE void PDLTriggerSecondary() {
|
||||
#endif
|
||||
}
|
||||
|
||||
/**
|
||||
* \brief Load data with the specified type and offset from a void pointer.
|
||||
* \tparam T The type to load.
|
||||
* \param ptr The base pointer.
|
||||
* \param offset The offset in number of elements of type T.
|
||||
*/
|
||||
template <typename T>
|
||||
SGL_DEVICE T load_as(const void* ptr, int64_t offset = 0) {
|
||||
return static_cast<const T*>(ptr)[offset];
|
||||
}
|
||||
|
||||
/**
|
||||
* \brief Store data with the specified type and offset to a void pointer.
|
||||
* \tparam T The type to store.
|
||||
* \param ptr The base pointer.
|
||||
* \param val The value to store.
|
||||
* \param offset The offset in number of elements of type T.
|
||||
* \note we use type_identity_t to force the caller to explicitly specify
|
||||
* the template parameter `T`, which can avoid accidentally using the wrong type.
|
||||
*/
|
||||
template <typename T>
|
||||
SGL_DEVICE void store_as(void* ptr, std::type_identity_t<T> val, int64_t offset = 0) {
|
||||
static_cast<T*>(ptr)[offset] = val;
|
||||
}
|
||||
|
||||
namespace pointer {
|
||||
|
||||
// we only allow void * pointer arithmetic for safety
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -1,47 +1,20 @@
|
||||
import time
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
from sgl_kernel import FusedSetKVBufferArg as FusedSetKVBufferArgKernel
|
||||
from sgl_kernel import (
|
||||
apply_rope_with_cos_sin_cache_inplace as apply_rope_with_cos_sin_cache_inplace_kernel,
|
||||
)
|
||||
|
||||
from sglang.jit_kernel.rope import FusedSetKVBufferArg as FusedSetKVBufferArgJit
|
||||
from sglang.jit_kernel.rope import (
|
||||
apply_rope_with_cos_sin_cache_inplace as apply_rope_with_cos_sin_cache_inplace_jit,
|
||||
)
|
||||
|
||||
DEVICE = "cuda"
|
||||
DTYPE = torch.bfloat16
|
||||
MAX_SEQ_LEN = 131072 # common seq length
|
||||
ROPE_BASE = 10000.0
|
||||
CACHE_SIZE = 1024 * 128
|
||||
|
||||
|
||||
@triton.jit
|
||||
def burn_kernel(out_ptr, iters: tl.constexpr):
|
||||
pid = tl.program_id(0)
|
||||
x = tl.full((), pid + 1, dtype=tl.uint32)
|
||||
|
||||
a = tl.full((), 1664525, dtype=tl.uint32)
|
||||
c = tl.full((), 1013904223, dtype=tl.uint32)
|
||||
sh = tl.full((), 13, dtype=tl.uint32)
|
||||
|
||||
for _ in range(iters):
|
||||
x = x * a + c
|
||||
x = x ^ (x >> sh)
|
||||
|
||||
if pid == 0:
|
||||
tl.store(out_ptr, x)
|
||||
|
||||
|
||||
def triton_burn(ms: float, grid=(256,)):
|
||||
iters = int(ms * 20000)
|
||||
out = torch.empty((), device="cuda", dtype=torch.uint32)
|
||||
burn_kernel[grid](out, iters=iters)
|
||||
return out
|
||||
|
||||
|
||||
def create_cos_sin_cache(rotary_dim, max_position_embeddings, base, dtype):
|
||||
def create_cos_sin_cache(
|
||||
rotary_dim: int,
|
||||
max_position: int = MAX_SEQ_LEN,
|
||||
base: float = ROPE_BASE,
|
||||
) -> torch.Tensor:
|
||||
"""Create cos/sin cache compatible with SGLang layout: [max_pos, rotary_dim]."""
|
||||
inv_freq = 1.0 / (
|
||||
base
|
||||
** (
|
||||
@@ -49,253 +22,223 @@ def create_cos_sin_cache(rotary_dim, max_position_embeddings, base, dtype):
|
||||
/ rotary_dim
|
||||
)
|
||||
)
|
||||
|
||||
t = torch.arange(max_position_embeddings, dtype=torch.float32, device=DEVICE)
|
||||
freqs = torch.einsum("i,j -> ij", t, inv_freq)
|
||||
t = torch.arange(max_position, dtype=torch.float32, device=DEVICE)
|
||||
freqs = torch.einsum("i,j->ij", t, inv_freq)
|
||||
cos = freqs.cos()
|
||||
sin = freqs.sin()
|
||||
cache = torch.cat((cos, sin), dim=-1)
|
||||
cache = torch.cat((cos, sin), dim=-1) # [max_pos, rotary_dim]
|
||||
return cache
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bs", [1, 8])
|
||||
@pytest.mark.parametrize("seq_len", [1, 512])
|
||||
@pytest.mark.parametrize("num_qo_heads", [1, 16])
|
||||
@pytest.mark.parametrize("num_kv_heads", [1, 16])
|
||||
@pytest.mark.parametrize("head_dim", [64, 512])
|
||||
@pytest.mark.parametrize("rotary_dim", [64, 128])
|
||||
@pytest.mark.parametrize("interleave", [False, True])
|
||||
@pytest.mark.parametrize("enable_pdl", [False, True])
|
||||
@pytest.mark.parametrize("save_kv_cache", [False, True])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
|
||||
# ---------------------------------------------------------------------------
|
||||
# Implementation wrappers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def sglang_jit_rope(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
from sglang.jit_kernel.rope import apply_rope_inplace
|
||||
|
||||
apply_rope_inplace(q, k, cos_sin_cache, positions, is_neox=is_neox)
|
||||
|
||||
|
||||
def flashinfer_rope(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
from flashinfer.rope import apply_rope_with_cos_sin_cache_inplace
|
||||
|
||||
head_size = q.shape[-1]
|
||||
# flashinfer expects [nnz, num_heads * head_size]
|
||||
q_2d = q.view(q.shape[0], -1)
|
||||
k_2d = k.view(k.shape[0], -1)
|
||||
apply_rope_with_cos_sin_cache_inplace(
|
||||
positions=positions,
|
||||
query=q_2d,
|
||||
key=k_2d,
|
||||
head_size=head_size,
|
||||
cos_sin_cache=cos_sin_cache,
|
||||
is_neox=is_neox,
|
||||
)
|
||||
|
||||
|
||||
def torch_impl_rope(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
cos_sin_cache: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
# TODO: implement a pure-PyTorch reference for extra coverage
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test parameters
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
BS_LIST = [2**x for x in range(12)]
|
||||
BS_LIST += [x + 1 for x in BS_LIST] # odd sizes to stress non-aligned paths
|
||||
NUM_KV_HEADS_LIST = [1, 2, 8]
|
||||
GQA_RATIO = [1, 4, 8]
|
||||
ROPE_DIM_LIST = [64, 128, 256, 512]
|
||||
IS_NEOX_LIST = [False, True]
|
||||
DTYPE_LIST = [torch.bfloat16, torch.float16]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", BS_LIST)
|
||||
@pytest.mark.parametrize("gqa_ratio", GQA_RATIO)
|
||||
@pytest.mark.parametrize("num_kv_heads", NUM_KV_HEADS_LIST)
|
||||
@pytest.mark.parametrize("rope_dim", ROPE_DIM_LIST)
|
||||
@pytest.mark.parametrize("is_neox", IS_NEOX_LIST)
|
||||
@pytest.mark.parametrize("dtype", DTYPE_LIST)
|
||||
def test_rope(
|
||||
bs,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
rotary_dim,
|
||||
interleave: bool,
|
||||
enable_pdl: bool,
|
||||
save_kv_cache: bool,
|
||||
batch_size: int,
|
||||
gqa_ratio: int,
|
||||
num_kv_heads: int,
|
||||
rope_dim: int,
|
||||
is_neox: bool,
|
||||
dtype: torch.dtype,
|
||||
) -> None:
|
||||
if head_dim < rotary_dim:
|
||||
pytest.skip(f"{head_dim=} < {rotary_dim=}")
|
||||
if not save_kv_cache and enable_pdl:
|
||||
pytest.skip(f"({save_kv_cache=}, {enable_pdl=}) is not allowed")
|
||||
|
||||
q = torch.randn(bs * seq_len, num_qo_heads * head_dim, device=DEVICE, dtype=dtype)
|
||||
k = torch.randn(bs * seq_len, num_kv_heads * head_dim, device=DEVICE, dtype=dtype)
|
||||
v = torch.randn(bs * seq_len, num_kv_heads * head_dim, device=DEVICE, dtype=dtype)
|
||||
|
||||
KV_POOL_SIZE = bs * seq_len * 2
|
||||
k_buffer = torch.zeros(
|
||||
KV_POOL_SIZE, num_kv_heads, head_dim, device=DEVICE, dtype=dtype
|
||||
num_qo_heads = num_kv_heads * gqa_ratio
|
||||
q = torch.randn(batch_size, num_qo_heads, rope_dim, device=DEVICE, dtype=dtype)
|
||||
k = torch.randn(batch_size, num_kv_heads, rope_dim, device=DEVICE, dtype=dtype)
|
||||
positions = torch.randint(
|
||||
0, MAX_SEQ_LEN, (batch_size,), device=DEVICE, dtype=torch.int64
|
||||
)
|
||||
v_buffer = torch.zeros(
|
||||
KV_POOL_SIZE, num_kv_heads, head_dim, device=DEVICE, dtype=dtype
|
||||
)
|
||||
out_cache_loc = torch.randperm(KV_POOL_SIZE, dtype=torch.int64, device=DEVICE)[
|
||||
: bs * seq_len
|
||||
].clone()
|
||||
cos_sin_cache = create_cos_sin_cache(rope_dim)
|
||||
|
||||
pos_ids = torch.arange(seq_len, device=DEVICE).repeat(bs)
|
||||
q_fi, k_fi = q.clone(), k.clone()
|
||||
q_jit, k_jit = q.clone(), k.clone()
|
||||
|
||||
max_seq_len = seq_len
|
||||
base = 10000
|
||||
cos_sin_cache = create_cos_sin_cache(rotary_dim, max_seq_len, base, dtype)
|
||||
flashinfer_rope(q_fi, k_fi, cos_sin_cache, positions, is_neox)
|
||||
sglang_jit_rope(q_jit, k_jit, cos_sin_cache, positions, is_neox)
|
||||
|
||||
q_jit = q.clone()
|
||||
k_jit = k.clone()
|
||||
v_jit = v.clone()
|
||||
k_buffer_jit = k_buffer.clone()
|
||||
v_buffer_jit = v_buffer.clone()
|
||||
out_cache_loc_jit = out_cache_loc.clone()
|
||||
fused_set_kv_buffer_arg_jit = FusedSetKVBufferArgJit(
|
||||
value=v_jit,
|
||||
k_buffer=k_buffer_jit.view(k_buffer_jit.shape[0], -1),
|
||||
v_buffer=v_buffer_jit.view(v_buffer_jit.shape[0], -1),
|
||||
k_scale=None,
|
||||
v_scale=None,
|
||||
cache_loc=out_cache_loc_jit,
|
||||
)
|
||||
|
||||
q_kernel = q.clone()
|
||||
k_kernel = k.clone()
|
||||
v_kernel = v.clone()
|
||||
k_buffer_kernel = k_buffer.clone()
|
||||
v_buffer_kernel = v_buffer.clone()
|
||||
out_cache_loc_kernel = out_cache_loc.clone()
|
||||
fused_set_kv_buffer_arg_kernel = FusedSetKVBufferArgKernel(
|
||||
value=v_kernel,
|
||||
k_buffer=k_buffer_kernel.view(k_buffer_kernel.shape[0], -1),
|
||||
v_buffer=v_buffer_kernel.view(v_buffer_kernel.shape[0], -1),
|
||||
k_scale=None,
|
||||
v_scale=None,
|
||||
cache_loc=out_cache_loc_kernel,
|
||||
)
|
||||
|
||||
stream_jit = torch.cuda.Stream()
|
||||
stream_kernel = torch.cuda.Stream()
|
||||
|
||||
triton_burn(10, grid=(1024,))
|
||||
r = torch.randn_like(q)
|
||||
r_jit, r_kernel = r.clone(), r.clone()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
with torch.cuda.stream(stream_jit):
|
||||
# Test if rotary_embedding runs on stream_jit
|
||||
triton_burn(10, grid=(1024,))
|
||||
q_jit = q_jit + r_jit
|
||||
apply_rope_with_cos_sin_cache_inplace_jit(
|
||||
positions=pos_ids,
|
||||
query=q_jit,
|
||||
key=k_jit,
|
||||
head_size=head_dim,
|
||||
cos_sin_cache=cos_sin_cache,
|
||||
is_neox=(not interleave),
|
||||
fused_set_kv_buffer_arg=(
|
||||
fused_set_kv_buffer_arg_jit if save_kv_cache else None
|
||||
),
|
||||
enable_pdl=enable_pdl,
|
||||
)
|
||||
|
||||
with torch.cuda.stream(stream_kernel):
|
||||
triton_burn(10, grid=(1024,))
|
||||
q_kernel = q_kernel + r_kernel
|
||||
apply_rope_with_cos_sin_cache_inplace_kernel(
|
||||
positions=pos_ids,
|
||||
query=q_kernel,
|
||||
key=k_kernel,
|
||||
head_size=head_dim,
|
||||
cos_sin_cache=cos_sin_cache,
|
||||
is_neox=(not interleave),
|
||||
fused_set_kv_buffer_arg=(
|
||||
fused_set_kv_buffer_arg_kernel if save_kv_cache else None
|
||||
),
|
||||
enable_pdl=enable_pdl,
|
||||
)
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
atol = 1e-3 if dtype != torch.float32 else 1e-6
|
||||
rtol = 1e-3 if dtype != torch.float32 else 1e-6
|
||||
torch.testing.assert_close(q_jit, q_kernel, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(k_jit, k_kernel, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(k_buffer_jit, k_buffer_kernel, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(v_buffer_jit, v_buffer_kernel, atol=atol, rtol=rtol)
|
||||
atol = rtol = 1e-2
|
||||
triton.testing.assert_close(q_fi, q_jit, atol=atol, rtol=rtol)
|
||||
triton.testing.assert_close(k_fi, k_jit, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bs", [8])
|
||||
@pytest.mark.parametrize("seq_len", [256, 512, 1024])
|
||||
@pytest.mark.parametrize("num_qo_heads", [16])
|
||||
@pytest.mark.parametrize("num_kv_heads", [16])
|
||||
@pytest.mark.parametrize("head_dim", [64])
|
||||
@pytest.mark.parametrize("rotary_dim", [64])
|
||||
@pytest.mark.parametrize("interleave", [False])
|
||||
@pytest.mark.parametrize("enable_pdl", [False])
|
||||
@pytest.mark.parametrize("save_kv_cache", [False])
|
||||
@pytest.mark.parametrize("dtype", [torch.bfloat16])
|
||||
def test_bench_rope(
|
||||
bs,
|
||||
seq_len,
|
||||
num_qo_heads,
|
||||
num_kv_heads,
|
||||
head_dim,
|
||||
rotary_dim,
|
||||
interleave: bool,
|
||||
enable_pdl: bool,
|
||||
save_kv_cache: bool,
|
||||
dtype: torch.dtype,
|
||||
@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
|
||||
def test_rope_position_dtypes(dtype: torch.dtype) -> None:
|
||||
"""Ensure both int32 and int64 position tensors work correctly."""
|
||||
batch_size, num_qo_heads, num_kv_heads, rope_dim = 16384, 16, 2, 128
|
||||
is_neox = True
|
||||
|
||||
q = torch.randn(batch_size, num_qo_heads, rope_dim, device=DEVICE, dtype=DTYPE)
|
||||
k = torch.randn(batch_size, num_kv_heads, rope_dim, device=DEVICE, dtype=DTYPE)
|
||||
positions = torch.randint(0, MAX_SEQ_LEN, (batch_size,), device=DEVICE, dtype=dtype)
|
||||
cos_sin_cache = create_cos_sin_cache(rope_dim)
|
||||
|
||||
q_fi, k_fi = q.clone(), k.clone()
|
||||
q_jit, k_jit = q.clone(), k.clone()
|
||||
|
||||
flashinfer_rope(q_fi, k_fi, cos_sin_cache, positions.long(), is_neox)
|
||||
sglang_jit_rope(q_jit, k_jit, cos_sin_cache, positions, is_neox)
|
||||
|
||||
atol = rtol = 1e-2
|
||||
triton.testing.assert_close(q_fi, q_jit, atol=atol, rtol=rtol)
|
||||
triton.testing.assert_close(k_fi, k_jit, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", BS_LIST)
|
||||
@pytest.mark.parametrize("is_neox", IS_NEOX_LIST)
|
||||
@pytest.mark.parametrize("rope_dim", [64, 80, 96, 128])
|
||||
@pytest.mark.parametrize("head_dim", [64, 128, 256])
|
||||
def test_partial_rope(batch_size: int, is_neox: bool, rope_dim: int, head_dim: int):
|
||||
if head_dim < rope_dim:
|
||||
pytest.skip("Invalid config: head_dim must be >= rope_dim.")
|
||||
num_qo_heads, num_kv_heads = 8, 2
|
||||
|
||||
q = torch.randn(batch_size, num_qo_heads, head_dim, device=DEVICE, dtype=DTYPE)
|
||||
k = torch.randn(batch_size, num_kv_heads, head_dim, device=DEVICE, dtype=DTYPE)
|
||||
positions = torch.randint(0, MAX_SEQ_LEN, (batch_size,), device=DEVICE)
|
||||
cos_sin_cache = create_cos_sin_cache(rope_dim)
|
||||
|
||||
q_fi, k_fi = q.clone(), k.clone()
|
||||
q_jit, k_jit = q.clone(), k.clone()
|
||||
rope = ..., slice(rope_dim) # NOTE: flashinfer by default apply to first rope_dim
|
||||
|
||||
flashinfer_rope(q_fi, k_fi, cos_sin_cache, positions.long(), is_neox)
|
||||
sglang_jit_rope(q_jit[rope], k_jit[rope], cos_sin_cache, positions, is_neox)
|
||||
|
||||
atol = rtol = 1e-2
|
||||
triton.testing.assert_close(q_fi, q_jit, atol=atol, rtol=rtol)
|
||||
triton.testing.assert_close(k_fi, k_jit, atol=atol, rtol=rtol)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", BS_LIST)
|
||||
@pytest.mark.parametrize("gqa_ratio", GQA_RATIO)
|
||||
@pytest.mark.parametrize("num_kv_heads", NUM_KV_HEADS_LIST)
|
||||
@pytest.mark.parametrize("rope_dim", ROPE_DIM_LIST)
|
||||
@pytest.mark.parametrize("is_neox", IS_NEOX_LIST)
|
||||
def test_fused_rope_store(
|
||||
batch_size: int,
|
||||
gqa_ratio: int,
|
||||
num_kv_heads: int,
|
||||
rope_dim: int,
|
||||
is_neox: bool,
|
||||
) -> None:
|
||||
if head_dim < rotary_dim:
|
||||
pytest.skip(f"{head_dim=} < {rotary_dim=}")
|
||||
if not save_kv_cache and enable_pdl:
|
||||
pytest.skip(f"({save_kv_cache=}, {enable_pdl=}) is not allowed")
|
||||
"""Test fused RoPE + KV cache store against separate RoPE + manual store."""
|
||||
from sglang.jit_kernel.rope import apply_rope_inplace_with_kvcache
|
||||
|
||||
q = torch.randn(bs * seq_len, num_qo_heads * head_dim, device=DEVICE, dtype=dtype)
|
||||
k = torch.randn(bs * seq_len, num_kv_heads * head_dim, device=DEVICE, dtype=dtype)
|
||||
v = torch.randn(bs * seq_len, num_kv_heads * head_dim, device=DEVICE, dtype=dtype)
|
||||
num_qo_heads = num_kv_heads * gqa_ratio
|
||||
dtype = DTYPE
|
||||
|
||||
KV_POOL_SIZE = bs * seq_len * 2
|
||||
k_buffer = torch.zeros(
|
||||
KV_POOL_SIZE, num_kv_heads, head_dim, device=DEVICE, dtype=dtype
|
||||
q = torch.randn(batch_size, num_qo_heads, rope_dim, device=DEVICE, dtype=dtype)
|
||||
k = torch.randn(batch_size, num_kv_heads, rope_dim, device=DEVICE, dtype=dtype)
|
||||
v = torch.randn(batch_size, num_kv_heads, rope_dim, device=DEVICE, dtype=dtype)
|
||||
positions = torch.randint(
|
||||
0, MAX_SEQ_LEN, (batch_size,), device=DEVICE, dtype=torch.int64
|
||||
)
|
||||
v_buffer = torch.zeros(
|
||||
KV_POOL_SIZE, num_kv_heads, head_dim, device=DEVICE, dtype=dtype
|
||||
out_loc = torch.randperm(CACHE_SIZE, device=DEVICE, dtype=torch.int64)[:batch_size]
|
||||
cos_sin_cache = create_cos_sin_cache(rope_dim)
|
||||
|
||||
row_size = num_kv_heads * rope_dim
|
||||
k_cache_ref = torch.zeros(CACHE_SIZE, row_size, device=DEVICE, dtype=dtype)
|
||||
v_cache_ref = torch.zeros(CACHE_SIZE, row_size, device=DEVICE, dtype=dtype)
|
||||
k_cache_fused = torch.zeros(CACHE_SIZE, row_size, device=DEVICE, dtype=dtype)
|
||||
v_cache_fused = torch.zeros(CACHE_SIZE, row_size, device=DEVICE, dtype=dtype)
|
||||
|
||||
# --- reference: separate RoPE then manual scatter ---
|
||||
q_ref, k_ref = q.clone(), k.clone()
|
||||
flashinfer_rope(q_ref, k_ref, cos_sin_cache, positions, is_neox)
|
||||
k_cache_ref[out_loc] = k_ref.view(batch_size, -1)
|
||||
v_cache_ref[out_loc] = v.view(batch_size, -1)
|
||||
|
||||
# --- fused kernel ---
|
||||
q_fused, k_fused = q.clone(), k.clone()
|
||||
v_fused = v.clone()
|
||||
apply_rope_inplace_with_kvcache(
|
||||
q_fused,
|
||||
k_fused,
|
||||
v_fused,
|
||||
k_cache_fused,
|
||||
v_cache_fused,
|
||||
cos_sin_cache,
|
||||
positions,
|
||||
out_loc,
|
||||
is_neox=is_neox,
|
||||
)
|
||||
out_cache_loc = torch.randperm(KV_POOL_SIZE, dtype=torch.int64, device=DEVICE)[
|
||||
: bs * seq_len
|
||||
].clone()
|
||||
|
||||
pos_ids = torch.arange(seq_len, device=DEVICE).repeat(bs)
|
||||
|
||||
max_seq_len = seq_len
|
||||
base = 10000
|
||||
cos_sin_cache = create_cos_sin_cache(rotary_dim, max_seq_len, base, dtype)
|
||||
|
||||
q_jit = q.clone()
|
||||
k_jit = k.clone()
|
||||
v_jit = v.clone()
|
||||
k_buffer_jit = k_buffer.clone()
|
||||
v_buffer_jit = v_buffer.clone()
|
||||
out_cache_loc_jit = out_cache_loc.clone()
|
||||
|
||||
q_kernel = q.clone()
|
||||
k_kernel = k.clone()
|
||||
v_kernel = v.clone()
|
||||
k_buffer_kernel = k_buffer.clone()
|
||||
v_buffer_kernel = v_buffer.clone()
|
||||
out_cache_loc_kernel = out_cache_loc.clone()
|
||||
|
||||
jit_args = {
|
||||
"positions": pos_ids,
|
||||
"query": q_jit,
|
||||
"key": k_jit,
|
||||
"head_size": head_dim,
|
||||
"cos_sin_cache": cos_sin_cache,
|
||||
"is_neox": (not interleave),
|
||||
"fused_set_kv_buffer_arg": None,
|
||||
"enable_pdl": enable_pdl,
|
||||
}
|
||||
jit_time = bench_rope(
|
||||
apply_rope_with_cos_sin_cache_inplace_jit,
|
||||
jit_args,
|
||||
atol = rtol = 1e-2
|
||||
# q should match RoPE-only result
|
||||
triton.testing.assert_close(q_ref, q_fused, atol=atol, rtol=rtol)
|
||||
# k_cache should contain the rotated k
|
||||
triton.testing.assert_close(
|
||||
k_cache_ref[out_loc], k_cache_fused[out_loc], atol=atol, rtol=rtol
|
||||
)
|
||||
kernel_args = {
|
||||
"positions": pos_ids,
|
||||
"query": q_kernel,
|
||||
"key": k_kernel,
|
||||
"head_size": head_dim,
|
||||
"cos_sin_cache": cos_sin_cache,
|
||||
"is_neox": (not interleave),
|
||||
"fused_set_kv_buffer_arg": None,
|
||||
"enable_pdl": enable_pdl,
|
||||
}
|
||||
kernel_time = bench_rope(
|
||||
apply_rope_with_cos_sin_cache_inplace_kernel,
|
||||
kernel_args,
|
||||
)
|
||||
print(f"\nPerformance Test - Batch={bs}, SeqLen={seq_len}")
|
||||
print(f"JIT: {jit_time*1000:.9f}ms, SGL: {kernel_time*1000:.9f}ms")
|
||||
if kernel_time > 0:
|
||||
speedup = kernel_time / jit_time if jit_time > 0 else float("inf")
|
||||
print(f"Speedup (SGL/JIT): {speedup:.2f}x")
|
||||
|
||||
|
||||
def bench_rope(fn, args):
|
||||
warmup = 10
|
||||
iteration = 100
|
||||
for _ in range(warmup):
|
||||
fn(**args)
|
||||
torch.cuda.synchronize()
|
||||
start_time = time.time()
|
||||
for _ in range(iteration):
|
||||
fn(**args)
|
||||
torch.cuda.synchronize()
|
||||
return (time.time() - start_time) / iteration
|
||||
# v_cache should be an exact copy
|
||||
assert torch.all(v_cache_ref[out_loc] == v_cache_fused[out_loc]), "v_cache mismatch"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__])
|
||||
pytest.main([__file__, "-v"])
|
||||
|
||||
@@ -5,7 +5,7 @@ from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
import math
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
@@ -37,13 +37,11 @@ _is_cpu = is_cpu()
|
||||
_is_xpu = is_xpu()
|
||||
_is_musa = is_musa()
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.jit_kernel.rope import FusedSetKVBufferArg # For type check-only
|
||||
|
||||
if _is_cuda:
|
||||
from sglang.jit_kernel.rope import (
|
||||
FusedSetKVBufferArg,
|
||||
apply_rope_with_cos_sin_cache_inplace,
|
||||
)
|
||||
else:
|
||||
FusedSetKVBufferArg = None
|
||||
from sglang.jit_kernel.rope import apply_rope_with_cos_sin_cache_inplace
|
||||
|
||||
if _use_aiter:
|
||||
from aiter.rotary_embedding import get_rope as aiter_get_rope
|
||||
@@ -362,19 +360,19 @@ class RotaryEmbedding(MultiPlatformOp):
|
||||
fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
if not self.use_fallback_kernel:
|
||||
batch_size = positions.size(0)
|
||||
q_rope = query.view(batch_size, -1, self.head_size)
|
||||
k_rope = key.view(batch_size, -1, self.head_size)
|
||||
if self.head_size != self.rotary_dim:
|
||||
q_rope = q_rope[..., : self.rotary_dim]
|
||||
k_rope = k_rope[..., : self.rotary_dim]
|
||||
apply_rope_with_cos_sin_cache_inplace(
|
||||
positions=positions,
|
||||
query=query,
|
||||
key=key,
|
||||
head_size=self.head_size,
|
||||
q=q_rope,
|
||||
k=k_rope,
|
||||
cos_sin_cache=self.cos_sin_cache,
|
||||
is_neox=self.is_neox_style,
|
||||
# Compatible with old sgl-kernel
|
||||
**(
|
||||
dict(fused_set_kv_buffer_arg=fused_set_kv_buffer_arg)
|
||||
if fused_set_kv_buffer_arg is not None
|
||||
else {}
|
||||
),
|
||||
fused_args=fused_set_kv_buffer_arg,
|
||||
)
|
||||
else:
|
||||
assert (
|
||||
|
||||
@@ -75,17 +75,12 @@ from sglang.srt.models.utils import (
|
||||
enable_fused_set_kv_buffer,
|
||||
)
|
||||
from sglang.srt.server_args import get_global_server_args
|
||||
from sglang.srt.utils import LazyValue, add_prefix, is_cuda, is_npu, make_layers
|
||||
from sglang.srt.utils import LazyValue, add_prefix, is_npu, make_layers
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
_is_cuda = is_cuda()
|
||||
_is_npu = is_npu()
|
||||
|
||||
|
||||
if _is_cuda:
|
||||
from sgl_kernel import FusedSetKVBufferArg # noqa: F401
|
||||
|
||||
|
||||
class GptOssConfig(PretrainedConfig):
|
||||
model_type = "gpt_oss"
|
||||
|
||||
|
||||
@@ -513,6 +513,10 @@ class LLaDA2MoeAttention(nn.Module):
|
||||
head_dim=self.head_dim,
|
||||
alt_stream=self.alt_stream,
|
||||
)
|
||||
can_fuse_set_kv = (
|
||||
self.head_dim == self.rotary_emb.rotary_dim
|
||||
and enable_fused_set_kv_buffer(forward_batch)
|
||||
)
|
||||
q, k = self.rotary_emb(
|
||||
positions,
|
||||
q,
|
||||
@@ -523,7 +527,7 @@ class LLaDA2MoeAttention(nn.Module):
|
||||
layer=self.attn,
|
||||
forward_batch=forward_batch,
|
||||
)
|
||||
if enable_fused_set_kv_buffer(forward_batch)
|
||||
if can_fuse_set_kv
|
||||
else None
|
||||
),
|
||||
)
|
||||
@@ -532,7 +536,7 @@ class LLaDA2MoeAttention(nn.Module):
|
||||
k,
|
||||
v,
|
||||
forward_batch,
|
||||
save_kv_cache=not enable_fused_set_kv_buffer(forward_batch),
|
||||
save_kv_cache=not can_fuse_set_kv,
|
||||
)
|
||||
attn_output, _ = self.dense(context_layer)
|
||||
return attn_output
|
||||
|
||||
@@ -119,20 +119,18 @@ def create_fused_set_kv_buffer_arg(
|
||||
layer: RadixAttention,
|
||||
forward_batch: ForwardBatch,
|
||||
):
|
||||
from sgl_kernel import FusedSetKVBufferArg
|
||||
from sglang.jit_kernel.rope import FusedSetKVBufferArg
|
||||
|
||||
layer_id = layer.layer_id
|
||||
token_to_kv_pool = forward_batch.token_to_kv_pool
|
||||
|
||||
k_buffer = token_to_kv_pool.get_key_buffer(layer_id)
|
||||
v_buffer = token_to_kv_pool.get_value_buffer(layer_id)
|
||||
|
||||
assert layer.k_scale is None and layer.v_scale is None, "scale not supported"
|
||||
return FusedSetKVBufferArg(
|
||||
value=value,
|
||||
k_buffer=k_buffer.view(k_buffer.shape[0], -1),
|
||||
v_buffer=v_buffer.view(v_buffer.shape[0], -1),
|
||||
k_scale=layer.k_scale,
|
||||
v_scale=layer.v_scale,
|
||||
cache_loc=forward_batch.out_cache_loc,
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user