diff --git a/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py new file mode 100644 index 000000000000..13a3b7681e68 --- /dev/null +++ b/python/sglang/jit_kernel/benchmark/bench_rotary_embedding.py @@ -0,0 +1,414 @@ +""" +Lean Rotary Embedding benchmark. + +Providers: +- jit_fused : SGL JIT with fused gather (positions + full cache) +- jit_unfused : SGL JIT with Python gather (indexing cos/sin in Python) +- aot_pos : SGL AOT with fused gather +- vllm_pos : vLLM +- flashinfer : FlashInfer +- torch_fp32 : Torch +""" + +import itertools +import os +from dataclasses import dataclass +from functools import lru_cache +from typing import Callable, Dict, Tuple + +import torch +import triton.testing + +DEVICE = "cuda" +DTYPE = torch.bfloat16 +MAX_SEQ_LEN = 8192 +NUM_HEADS = 32 + +IS_CI = ( + os.getenv("CI", "false").lower() == "true" + or os.getenv("GITHUB_ACTIONS", "false").lower() == "true" +) + +if IS_CI: + BS_RANGE = [1] + SEQ_RANGE = [1] + HEAD_SIZE_RANGE = [128] +else: + BS_RANGE = [1, 2, 4, 8, 16, 32, 64, 128] + SEQ_RANGE = [1, 16, 64, 256, 1024, 2048] + HEAD_SIZE_RANGE = [64, 128, 256] + +INTERLEAVED_RANGE = [True, False] + +ONLY_PROVIDER = os.getenv("SGL_BENCH_PROVIDER", "").strip() or None +INCLUDE_TORCH_REF = os.getenv("SGL_BENCH_INCLUDE_TORCH", "0").lower() in ( + "1", + "true", + "yes", + "y", +) +SHOW_SPEEDUP = os.getenv("SGL_BENCH_SPEEDUP", "1").lower() in ("1", "true", "yes", "y") + +try: + import sgl_kernel # noqa: F401 +except Exception: + sgl_kernel = None # type: ignore + +HAS_SGL_POS = hasattr(torch.ops, "sgl_kernel") and hasattr( + torch.ops.sgl_kernel, "rotary_embedding" +) + +try: + from vllm.model_executor.layers.rotary_embedding import ( + RotaryEmbedding as vLLMRotaryEmbedding, + ) + + HAS_VLLM = True +except Exception: + vLLMRotaryEmbedding = None + HAS_VLLM = False + +try: + from sglang.multimodal_gen.runtime.layers.rotary_embedding import ( + apply_flashinfer_rope_qk_inplace, + ) + + HAS_FLASHINFER = True +except Exception: + apply_flashinfer_rope_qk_inplace = None # type: ignore + HAS_FLASHINFER = False + + +def _compute_cos_sin_cache( + max_seq_len: int, + rotary_dim: int, + base: float = 10000.0, + dtype: torch.dtype = torch.float32, +): + inv_freq = 1.0 / ( + base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim) + ) + t = torch.arange(max_seq_len, dtype=torch.float32) + freqs = torch.einsum("i,j->ij", t, inv_freq) + return freqs.cos().to(dtype), freqs.sin().to(dtype) + + +@lru_cache(maxsize=None) +def _cos_sin_cache_half_cuda(rotary_dim: int) -> Tuple[torch.Tensor, torch.Tensor]: + cos, sin = _compute_cos_sin_cache(MAX_SEQ_LEN, rotary_dim, dtype=DTYPE) + return cos.to(device=DEVICE, dtype=DTYPE), sin.to(device=DEVICE, dtype=DTYPE) + + +@lru_cache(maxsize=None) +def _vllm_rope(head_size: int, rotary_dim: int, interleaved: bool, dtype: torch.dtype): + if not HAS_VLLM or vLLMRotaryEmbedding is None: + raise RuntimeError("vLLM not available") + return vLLMRotaryEmbedding( + head_size=head_size, + rotary_dim=rotary_dim, + max_position_embeddings=MAX_SEQ_LEN, + base=10000, + is_neox_style=interleaved, + dtype=dtype, + ).cuda() + + +@torch.no_grad() +def torch_impl_rotary_fp32(cos, sin, q, k, head_size: int, interleaved: bool) -> None: + orig_dtype = q.dtype + if interleaved and cos.shape[1] == head_size: + half = head_size // 2 + cos = cos.view(cos.shape[0], half, 2).select(2, 0).contiguous() + sin = sin.view(sin.shape[0], half, 2).select(2, 1).contiguous() + + cos_f, sin_f = cos.float(), sin.float() + q_f, k_f = q.float(), k.float() + + if interleaved: + embed_dim = int(cos_f.shape[1]) + rot_dim = embed_dim * 2 + cos_b = cos_f[:, None, :embed_dim] + sin_b = sin_f[:, None, :embed_dim] + + def _apply(x): + xr = x[..., :rot_dim] + xr2 = xr.view(xr.shape[0], xr.shape[1], embed_dim, 2) + x0 = xr2[..., 0].clone() + x1 = xr2[..., 1].clone() + xr2[..., 0].copy_(x0 * cos_b - x1 * sin_b) + xr2[..., 1].copy_(x1 * cos_b + x0 * sin_b) + + else: + if cos_f.shape[1] == head_size // 2: + embed_dim = int(cos_f.shape[1]) + rot_dim = embed_dim * 2 + cos_x, sin_x = cos_f[:, None, :], sin_f[:, None, :] + cos_y, sin_y = cos_x, sin_x + else: + embed_dim = int(cos_f.shape[1]) // 2 + rot_dim = embed_dim * 2 + cos_x, sin_x = cos_f[:, None, :embed_dim], sin_f[:, None, :embed_dim] + cos_y, sin_y = ( + cos_f[:, None, embed_dim:rot_dim], + sin_f[:, None, embed_dim:rot_dim], + ) + + def _apply(x): + xr = x[..., :rot_dim] + x0 = xr[..., :embed_dim].clone() + x1 = xr[..., embed_dim:rot_dim].clone() + xr[..., :embed_dim].copy_(x0 * cos_x - x1 * sin_x) + xr[..., embed_dim:rot_dim].copy_(x1 * cos_y + x0 * sin_y) + + _apply(q_f) + _apply(k_f) + q.copy_(q_f.to(orig_dtype)) + k.copy_(k_f.to(orig_dtype)) + + +def _assert_close_gpu(actual, expected, atol: float, rtol: float, name: str): + diff = (actual - expected).abs() + max_abs = float(diff.max().item()) + denom = atol + rtol * expected.abs() + max_rel = float((diff / denom).max().item()) + if max_abs > atol and max_rel > 1.0: + raise AssertionError( + f"{name}: not close max_abs={max_abs:.6g} max_rel={max_rel:.6g}" + ) + + +def sgl_jit_fused( + positions, cos_cache, sin_cache, q, k, head_size: int, interleaved: bool +): + from sglang.jit_kernel.rotary_embedding import rotary_embedding_cos_sin + + rotary_embedding_cos_sin( + cos_cache, sin_cache, q, k, head_size, interleaved, positions=positions + ) + + +def sgl_jit_unfused( + positions, cos_cache, sin_cache, q, k, head_size: int, interleaved: bool +): + from sglang.jit_kernel.rotary_embedding import rotary_embedding_cos_sin + + cos_g = cos_cache[positions].contiguous() + sin_g = sin_cache[positions].contiguous() + rotary_embedding_cos_sin(cos_g, sin_g, q, k, head_size, interleaved) + + +def sgl_aot_pos(positions, q, k, head_size: int, interleaved: bool, cos_sin_cache): + if not HAS_SGL_POS: + raise RuntimeError("torch.ops.sgl_kernel.rotary_embedding not available") + torch.ops.sgl_kernel.rotary_embedding( + positions, q, k, head_size, cos_sin_cache, not interleaved # is_neox flag + ) + + +def _is_flashinfer_unsupported(e: BaseException) -> bool: + msg = str(e) + return ("Unsupported head_dim" in msg) or ("cos_sin_cache should be float32" in msg) + + +@dataclass +class Case: + positions: torch.Tensor + cos_cache: torch.Tensor + sin_cache: torch.Tensor + cos_gathered: torch.Tensor + sin_gathered: torch.Tensor + cos_sin_cache_aot: torch.Tensor + q_base: torch.Tensor # [T, H, D] + k_base: torch.Tensor # [T, H, D] + q4_base: torch.Tensor # [B,S,H,D] for flashinfer + k4_base: torch.Tensor + cos_sin_cache_flashinfer: torch.Tensor | None + + +def prepare_case( + batch_size: int, seq_len: int, head_size: int, interleaved: bool +) -> Case: + num_tokens = batch_size * seq_len + rotary_dim = head_size + + cos_cache_half, sin_cache_half = _cos_sin_cache_half_cuda(rotary_dim) + positions = torch.arange(seq_len, device=DEVICE, dtype=torch.int64).repeat( + batch_size + ) + + cos_half = cos_cache_half[positions].contiguous() + sin_half = sin_cache_half[positions].contiguous() + + if interleaved: + cos_cache, sin_cache = cos_cache_half, sin_cache_half + cos_g, sin_g = cos_half, sin_half + else: + # non-interleaved: kernel expects R=2*embed_dim (x and y halves) + cos_cache = torch.cat([cos_cache_half, cos_cache_half], dim=-1).contiguous() + sin_cache = torch.cat([sin_cache_half, sin_cache_half], dim=-1).contiguous() + cos_g = torch.cat([cos_half, cos_half], dim=-1).contiguous() + sin_g = torch.cat([sin_half, sin_half], dim=-1).contiguous() + + cos_sin_cache_aot = torch.cat([cos_cache_half, sin_cache_half], dim=-1).contiguous() + + q4 = torch.randn( + batch_size, seq_len, NUM_HEADS, head_size, device=DEVICE, dtype=DTYPE + ).contiguous() + k4 = torch.randn( + batch_size, seq_len, NUM_HEADS, head_size, device=DEVICE, dtype=DTYPE + ).contiguous() + q = q4.view(num_tokens, NUM_HEADS, head_size) + k = k4.view(num_tokens, NUM_HEADS, head_size) + + cos_sin_cache_flashinfer = None + if HAS_FLASHINFER: + cos_f32, sin_f32 = _compute_cos_sin_cache( + MAX_SEQ_LEN, rotary_dim, dtype=torch.float32 + ) + cos_sin_cache_flashinfer = ( + torch.cat([cos_f32, sin_f32], dim=-1) + .to(device=DEVICE, dtype=torch.float32) + .contiguous() + ) + + return Case( + positions=positions, + cos_cache=cos_cache, + sin_cache=sin_cache, + cos_gathered=cos_g, + sin_gathered=sin_g, + cos_sin_cache_aot=cos_sin_cache_aot, + q_base=q, + k_base=k, + q4_base=q4, + k4_base=k4, + cos_sin_cache_flashinfer=cos_sin_cache_flashinfer, + ) + + +def make_fn( + provider: str, case: Case, head_size: int, interleaved: bool +) -> Callable[[], None]: + if provider in ("jit_fused", "jit_unfused", "aot_pos", "torch_fp32", "vllm_pos"): + q = case.q_base.clone() + k = case.k_base.clone() + + if provider == "jit_fused": + return lambda: sgl_jit_fused( + case.positions, + case.cos_cache, + case.sin_cache, + q, + k, + head_size, + interleaved, + ) + if provider == "jit_unfused": + return lambda: sgl_jit_unfused( + case.positions, + case.cos_cache, + case.sin_cache, + q, + k, + head_size, + interleaved, + ) + if provider == "aot_pos": + return lambda: sgl_aot_pos( + case.positions, q, k, head_size, interleaved, case.cos_sin_cache_aot + ) + if provider == "torch_fp32": + return lambda: torch_impl_rotary_fp32( + case.cos_gathered, case.sin_gathered, q, k, head_size, interleaved + ) + if provider == "vllm_pos": + rope = _vllm_rope(head_size, head_size, interleaved, DTYPE) + return lambda: rope.forward_cuda(case.positions, q, k) + + if provider == "flashinfer": + if not (HAS_FLASHINFER and apply_flashinfer_rope_qk_inplace is not None): + raise RuntimeError("FlashInfer not available") + q4 = case.q4_base.clone() + k4 = case.k4_base.clone() + return lambda: apply_flashinfer_rope_qk_inplace( + q4, + k4, + case.cos_sin_cache_flashinfer, + is_neox=not interleaved, + positions=case.positions, + ) + + raise ValueError(f"Unknown provider={provider}") + + +def bench(fn: Callable[[], None], use_cudagraph: bool) -> Tuple[float, float, float]: + qs = [0.5, 0.2, 0.8] + if use_cudagraph: + return triton.testing.do_bench_cudagraph(fn, quantiles=qs) # type: ignore + return triton.testing.do_bench(fn, quantiles=qs) # type: ignore + + +def available_providers() -> Dict[str, str]: + ps: Dict[str, str] = { + "jit_fused": "SGL JIT (fused gather)", + "jit_unfused": "SGL JIT (Python gather)", + } + if HAS_SGL_POS: + ps["aot_pos"] = "SGL AOT (fused gather)" + if HAS_VLLM: + ps["vllm_pos"] = "vLLM" + if HAS_FLASHINFER: + ps["flashinfer"] = "FlashInfer" + if INCLUDE_TORCH_REF: + ps["torch_fp32"] = "Torch" + return ps + + +PROVIDERS = available_providers() +if ONLY_PROVIDER is not None: + if ONLY_PROVIDER not in PROVIDERS: + raise ValueError( + f"Unknown provider={ONLY_PROVIDER}. Allowed: {list(PROVIDERS)}" + ) + PROVIDERS = {ONLY_PROVIDER: PROVIDERS[ONLY_PROVIDER]} + +configs = list( + itertools.product(BS_RANGE, SEQ_RANGE, HEAD_SIZE_RANGE, INTERLEAVED_RANGE) +) + +_SANITY_DONE = False + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size", "seq_len", "head_size", "interleaved"], + x_vals=configs, + line_arg="provider", + line_vals=list(PROVIDERS.keys()), + line_names=list(PROVIDERS.values()), + styles=None, + ylabel="us", + plot_name="rotary-embedding-performance", + args={}, + ) +) +def benchmark( + batch_size: int, seq_len: int, head_size: int, interleaved: bool, provider: str +): + case = prepare_case(batch_size, seq_len, head_size, interleaved) + fn = make_fn(provider, case, head_size, interleaved) + + use_cudagraph = provider != "flashinfer" + try: + ms, min_ms, max_ms = bench(fn, use_cudagraph=use_cudagraph) + except Exception as e: + if provider == "flashinfer" and _is_flashinfer_unsupported(e): + nan = float("nan") + return nan, nan, nan + raise + + return 1000 * ms, 1000 * min_ms, 1000 * max_ms + + +if __name__ == "__main__": + benchmark.run(print_data=True) diff --git a/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh new file mode 100644 index 000000000000..08a00124cfff --- /dev/null +++ b/python/sglang/jit_kernel/csrc/rotary_embedding_cos_sin.cuh @@ -0,0 +1,818 @@ +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace { + +// Convert scalar types to float for computation. +template +__device__ __forceinline__ float to_float(T v) { + return static_cast(v); +} +template <> +__device__ __forceinline__ float to_float(half v) { + return __half2float(v); +} +template <> +__device__ __forceinline__ float to_float(nv_bfloat16 v) { + return __bfloat162float(v); +} + +// Convert float back to scalar types (round-to-nearest for fp16/bf16). +template +__device__ __forceinline__ T from_float(float v) { + return static_cast(v); +} +template <> +__device__ __forceinline__ half from_float(float v) { + return __float2half_rn(v); +} +template <> +__device__ __forceinline__ nv_bfloat16 from_float(float v) { + return __float2bfloat16_rn(v); +} + +// Vectorize loads/stores in 16 bytes (float4) regardless of scalar type. +template +struct RotaryVecData; + +template +struct RotaryVecData { + float4 vec; + float2 cos_vec; + float2 sin_vec; +}; + +template +struct RotaryVecData { + float4 vec_x; + float4 vec_y; + float4 cos_x; + float4 sin_x; + float4 cos_y; + float4 sin_y; +}; + +// Per-thread temporary storage for a 16B tile and its cos/sin values. +// Layout depends on interleaved vs split-halves format. +template +__device__ __forceinline__ RotaryVecData load_rotary_vec( + const scalar_t* __restrict__ arr, + const scalar_t* __restrict__ cos_ptr, + const scalar_t* __restrict__ sin_ptr, + int rot_offset, + int embed_dim) { + RotaryVecData data; + + constexpr int kVecBytes = 16; + constexpr int kElePerVec = kVecBytes / sizeof(scalar_t); + + if constexpr (interleaved) { + const int base = rot_offset * 2; + + if constexpr (aligned_qk) { + data.vec = *reinterpret_cast(arr + base); + } else { + union VecU { + float4 v; + scalar_t e[kElePerVec]; + } tmp; +#pragma unroll + for (int i = 0; i < kElePerVec; ++i) + tmp.e[i] = arr[base + i]; + data.vec = tmp.v; + } + + data.cos_vec = *reinterpret_cast(cos_ptr + rot_offset); + data.sin_vec = *reinterpret_cast(sin_ptr + rot_offset); + } else { + const int bx = rot_offset; + const int by = rot_offset + embed_dim; + + if constexpr (aligned_qk) { + data.vec_x = *reinterpret_cast(arr + bx); + data.vec_y = *reinterpret_cast(arr + by); + } else { + union VecU { + float4 v; + scalar_t e[kElePerVec]; + } tx, ty; +#pragma unroll + for (int i = 0; i < kElePerVec; ++i) { + tx.e[i] = arr[bx + i]; + ty.e[i] = arr[by + i]; + } + data.vec_x = tx.v; + data.vec_y = ty.v; + } + + data.cos_x = *reinterpret_cast(cos_ptr + bx); + data.sin_x = *reinterpret_cast(sin_ptr + bx); + data.cos_y = *reinterpret_cast(cos_ptr + by); + data.sin_y = *reinterpret_cast(sin_ptr + by); + } + + return data; +} + +// Apply RoPE to the loaded 16B tile and write back to q/k +template +__device__ __forceinline__ void compute_store_rotary_vec( + scalar_t* __restrict__ arr, const RotaryVecData& data, int rot_offset, int embed_dim) { + constexpr int kVecBytes = 16; + constexpr int kElePerVec = kVecBytes / sizeof(scalar_t); + + if constexpr (interleaved) { + union VecU { + float4 v; + scalar_t e[kElePerVec]; + } v; + v.v = data.vec; + + union CSU { + float2 v; + scalar_t e[kElePerVec / 2]; + } c, s; + c.v = data.cos_vec; + s.v = data.sin_vec; + +#pragma unroll + for (int i = 0; i < kElePerVec; i += 2) { + const int idx = i / 2; + const float cos_val = to_float(c.e[idx]); + const float sin_val = to_float(s.e[idx]); + const float x = to_float(v.e[i]); + const float y = to_float(v.e[i + 1]); + v.e[i] = from_float(x * cos_val - y * sin_val); + v.e[i + 1] = from_float(y * cos_val + x * sin_val); + } + + const int base = rot_offset * 2; + if constexpr (aligned_qk) { + *reinterpret_cast(arr + base) = v.v; + } else { +#pragma unroll + for (int i = 0; i < kElePerVec; ++i) + arr[base + i] = v.e[i]; + } + } else { + union VecU { + float4 v; + scalar_t e[kElePerVec]; + } vx, vy; + vx.v = data.vec_x; + vy.v = data.vec_y; + + union CSU { + float4 v; + scalar_t e[kElePerVec]; + } cx, sx, cy, sy; + cx.v = data.cos_x; + sx.v = data.sin_x; + cy.v = data.cos_y; + sy.v = data.sin_y; + +#pragma unroll + for (int i = 0; i < kElePerVec; ++i) { + const float cos_x = to_float(cx.e[i]); + const float sin_x = to_float(sx.e[i]); + const float cos_y = to_float(cy.e[i]); + const float sin_y = to_float(sy.e[i]); + + const float x = to_float(vx.e[i]); + const float y = to_float(vy.e[i]); + + vx.e[i] = from_float(x * cos_x - y * sin_x); + vy.e[i] = from_float(y * cos_y + x * sin_y); + } + + const int bx = rot_offset; + const int by = rot_offset + embed_dim; + if constexpr (aligned_qk) { + *reinterpret_cast(arr + bx) = vx.v; + *reinterpret_cast(arr + by) = vy.v; + } else { +#pragma unroll + for (int i = 0; i < kElePerVec; ++i) { + arr[bx + i] = vx.e[i]; + arr[by + i] = vy.e[i]; + } + } + } +} + +// Scalar fallback: Apply RoPE for exactly one (x,y) pair per iteration. +template +inline __device__ void apply_token_rotary_embedding( + scalar_t* __restrict__ arr, // [head_size] + const scalar_t* __restrict__ cos_ptr, // [rot_dim] + const scalar_t* __restrict__ sin_ptr, // [rot_dim] + int rot_offset, + int embed_dim) { + if constexpr (interleaved) { + const int x_index = 2 * rot_offset; + const int y_index = x_index + 1; + + const float cos_val = to_float(cos_ptr[rot_offset]); + const float sin_val = to_float(sin_ptr[rot_offset]); + const float x = to_float(arr[x_index]); + const float y = to_float(arr[y_index]); + + arr[x_index] = from_float(x * cos_val - y * sin_val); + arr[y_index] = from_float(y * cos_val + x * sin_val); + } else { + const int x_index = rot_offset; + const int y_index = rot_offset + embed_dim; + + const float cos_val_x = to_float(cos_ptr[rot_offset]); + const float sin_val_x = to_float(sin_ptr[rot_offset]); + const float cos_val_y = to_float(cos_ptr[rot_offset + embed_dim]); + const float sin_val_y = to_float(sin_ptr[rot_offset + embed_dim]); + + const float x = to_float(arr[x_index]); + const float y = to_float(arr[y_index]); + + arr[x_index] = from_float(x * cos_val_x - y * sin_val_x); + arr[y_index] = from_float(y * cos_val_y + x * sin_val_y); + } +} + +// RoPE kernel with 2D grid: (token_idx, sub_block_idx). +// Each block processes a strided subset of rotary pairs across all heads for one token. +// Supports optional key tensor (nullptr means Q-only). +template +__global__ void rotary_embedding_kernel_2d( + const scalar_t* __restrict__ cos_data, // [num_tokens, rot_dim_arg] + const scalar_t* __restrict__ sin_data, // [num_tokens, rot_dim_arg] + scalar_t* __restrict__ query_total, // [num_tokens, num_heads, head_size] contiguous + scalar_t* __restrict__ key_total, // [num_tokens, num_kv_heads, head_size] contiguous or nullptr + const int64_t* __restrict__ positions, // [num_tokens] or nullptr + const int rot_dim_arg, + const int embed_dim_for_rotation_arg, + const int64_t query_token_stride, + const int64_t key_token_stride, + const int64_t head_stride_query, + const int64_t head_stride_key, + const int num_heads, + const int num_kv_heads, + const int head_size_arg, + const int blocks_per_token) { + const int token_idx = blockIdx.x; + if (token_idx >= gridDim.x) return; + + const int64_t rot_idx = (positions != nullptr) ? positions[token_idx] : token_idx; + const scalar_t* current_token_cos_ptr = cos_data + rot_idx * rot_dim_arg; + const scalar_t* current_token_sin_ptr = sin_data + rot_idx * rot_dim_arg; + + scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride; + scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr; + + const int local_block_idx = blockIdx.y; + const int embed_dim_for_rotation = (ROT_EMBED_DIM > 0) ? ROT_EMBED_DIM : embed_dim_for_rotation_arg; + + if constexpr (vectorized) { + constexpr int kVecBytes = 16; + constexpr int kElePerVec = kVecBytes / sizeof(scalar_t); + constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec; + + const int pair_stride = blockDim.x * blocks_per_token * pairs_per_step; + int i = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step; + const int nq_pairs = num_heads * embed_dim_for_rotation; + + RotaryVecData curr_data; + if (i < nq_pairs) { + const int head_idx = i / embed_dim_for_rotation; + const int rot_offset = i % embed_dim_for_rotation; + scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query; + curr_data = load_rotary_vec( + ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation); + } + + int next_i = i + pair_stride; + for (; i < nq_pairs; i += pair_stride, next_i += pair_stride) { + RotaryVecData next_data; + const bool active_next = (next_i < nq_pairs); + + if (active_next) { + const int head_idx_next = next_i / embed_dim_for_rotation; + const int rot_offset_next = next_i % embed_dim_for_rotation; + scalar_t* ptr_next = query_for_token + head_idx_next * (int)head_stride_query; + next_data = load_rotary_vec( + ptr_next, current_token_cos_ptr, current_token_sin_ptr, rot_offset_next, embed_dim_for_rotation); + } + + const int head_idx = i / embed_dim_for_rotation; + const int rot_offset = i % embed_dim_for_rotation; + scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query; + compute_store_rotary_vec(ptr, curr_data, rot_offset, embed_dim_for_rotation); + + if (active_next) curr_data = next_data; + } + + if (key_for_token != nullptr) { + const int nk_pairs = num_kv_heads * embed_dim_for_rotation; + int k_i = (local_block_idx * blockDim.x + threadIdx.x) * pairs_per_step; + + RotaryVecData curr_data_k; + if (k_i < nk_pairs) { + const int head_idx = k_i / embed_dim_for_rotation; + const int rot_offset = k_i % embed_dim_for_rotation; + scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key; + curr_data_k = load_rotary_vec( + ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation); + } + + int next_k_i = k_i + pair_stride; + for (; k_i < nk_pairs; k_i += pair_stride, next_k_i += pair_stride) { + RotaryVecData next_data_k; + const bool active_next = (next_k_i < nk_pairs); + if (active_next) { + const int head_idx_next = next_k_i / embed_dim_for_rotation; + const int rot_offset_next = next_k_i % embed_dim_for_rotation; + scalar_t* ptr_next = key_for_token + head_idx_next * (int)head_stride_key; + next_data_k = load_rotary_vec( + ptr_next, current_token_cos_ptr, current_token_sin_ptr, rot_offset_next, embed_dim_for_rotation); + } + + const int head_idx = k_i / embed_dim_for_rotation; + const int rot_offset = k_i % embed_dim_for_rotation; + scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key; + compute_store_rotary_vec( + ptr, curr_data_k, rot_offset, embed_dim_for_rotation); + + if (active_next) curr_data_k = next_data_k; + } + } + } else { + // Fallback to scalar implementation + const int pair_stride = blockDim.x * blocks_per_token; + const int thread_pair_offset = local_block_idx * blockDim.x + threadIdx.x; + + const int nq_pairs = num_heads * embed_dim_for_rotation; + for (int i = thread_pair_offset; i < nq_pairs; i += pair_stride) { + const int head_idx = i / embed_dim_for_rotation; + const int rot_offset = i % embed_dim_for_rotation; + scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query; + apply_token_rotary_embedding( + query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation); + } + + if (key_for_token != nullptr) { + const int nk_pairs = num_kv_heads * embed_dim_for_rotation; + for (int i = thread_pair_offset; i < nk_pairs; i += pair_stride) { + const int head_idx = i / embed_dim_for_rotation; + const int rot_offset = i % embed_dim_for_rotation; + scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key; + apply_token_rotary_embedding( + key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation); + } + } + } +} + +// RoPE kernel with 1D grid: one block per token. +template +__launch_bounds__(512) __global__ void rotary_embedding_kernel_1d( + const scalar_t* __restrict__ cos_data, // [num_tokens, rot_dim_arg] + const scalar_t* __restrict__ sin_data, // [num_tokens, rot_dim_arg] + scalar_t* __restrict__ query_total, + scalar_t* __restrict__ key_total, + const int64_t* __restrict__ positions, // [num_tokens] or nullptr + const int rot_dim_arg, + const int embed_dim_for_rotation_arg, + const int64_t query_token_stride, + const int64_t key_token_stride, + const int64_t head_stride_query, + const int64_t head_stride_key, + const int num_heads, + const int num_kv_heads, + const int head_size_arg) { + const int token_idx = blockIdx.x; + if (token_idx >= gridDim.x) return; + + const int64_t rot_idx = (positions != nullptr) ? positions[token_idx] : token_idx; + const scalar_t* current_token_cos_ptr = cos_data + rot_idx * rot_dim_arg; + const scalar_t* current_token_sin_ptr = sin_data + rot_idx * rot_dim_arg; + + scalar_t* query_for_token = query_total + token_idx * (int)query_token_stride; + scalar_t* key_for_token = (key_total != nullptr) ? (key_total + token_idx * (int)key_token_stride) : nullptr; + + const int embed_dim_for_rotation = (ROT_EMBED_DIM > 0) ? ROT_EMBED_DIM : embed_dim_for_rotation_arg; + + if constexpr (vectorized) { + constexpr int kVecBytes = 16; + constexpr int kElePerVec = kVecBytes / sizeof(scalar_t); + constexpr int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec; + + const int nq_pairs = num_heads * embed_dim_for_rotation; + const int stride = blockDim.x * pairs_per_step; + int i = threadIdx.x * pairs_per_step; + + RotaryVecData curr_data; + if (i < nq_pairs) { + const int head_idx = i / embed_dim_for_rotation; + const int rot_offset = i % embed_dim_for_rotation; + scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query; + curr_data = load_rotary_vec( + ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation); + } + + int next_i = i + stride; + for (; i < nq_pairs; i += stride, next_i += stride) { + RotaryVecData next_data = curr_data; + if (next_i < nq_pairs) { + const int head_idx = next_i / embed_dim_for_rotation; + const int rot_offset = next_i % embed_dim_for_rotation; + scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query; + next_data = load_rotary_vec( + ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation); + } + + if (i < nq_pairs) { + const int head_idx = i / embed_dim_for_rotation; + const int rot_offset = i % embed_dim_for_rotation; + scalar_t* ptr = query_for_token + head_idx * (int)head_stride_query; + compute_store_rotary_vec(ptr, curr_data, rot_offset, embed_dim_for_rotation); + } + curr_data = next_data; + } + + if (key_for_token != nullptr) { + const int nk_pairs = num_kv_heads * embed_dim_for_rotation; + int k_i = threadIdx.x * pairs_per_step; + RotaryVecData curr_data_k; + + if (k_i < nk_pairs) { + const int head_idx = k_i / embed_dim_for_rotation; + const int rot_offset = k_i % embed_dim_for_rotation; + scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key; + curr_data_k = load_rotary_vec( + ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation); + } + + int next_k_i = k_i + stride; + for (; k_i < nk_pairs; k_i += stride, next_k_i += stride) { + RotaryVecData next_data_k = curr_data_k; + if (next_k_i < nk_pairs) { + const int head_idx = next_k_i / embed_dim_for_rotation; + const int rot_offset = next_k_i % embed_dim_for_rotation; + scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key; + next_data_k = load_rotary_vec( + ptr, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation); + } + + if (k_i < nk_pairs) { + const int head_idx = k_i / embed_dim_for_rotation; + const int rot_offset = k_i % embed_dim_for_rotation; + scalar_t* ptr = key_for_token + head_idx * (int)head_stride_key; + compute_store_rotary_vec( + ptr, curr_data_k, rot_offset, embed_dim_for_rotation); + } + curr_data_k = next_data_k; + } + } + } else { + const int nq_pairs = num_heads * embed_dim_for_rotation; + for (int i = threadIdx.x; i < nq_pairs; i += blockDim.x) { + const int head_idx = i / embed_dim_for_rotation; + const int rot_offset = i % embed_dim_for_rotation; + scalar_t* query_for_token_head = query_for_token + head_idx * (int)head_stride_query; + apply_token_rotary_embedding( + query_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation); + } + + if (key_for_token != nullptr) { + const int nk_pairs = num_kv_heads * embed_dim_for_rotation; + for (int i = threadIdx.x; i < nk_pairs; i += blockDim.x) { + const int head_idx = i / embed_dim_for_rotation; + const int rot_offset = i % embed_dim_for_rotation; + scalar_t* key_for_token_head = key_for_token + head_idx * (int)head_stride_key; + apply_token_rotary_embedding( + key_for_token_head, current_token_cos_ptr, current_token_sin_ptr, rot_offset, embed_dim_for_rotation); + } + } + } +} + +// Dispatch kernel specialization by layout (interleaved), vectorization, and alignment. +// Selects 1D vs 2D grid based on launch configuration. +template +__forceinline__ void dispatch_rotary_launch( + bool use_grid_2d, + dim3 grid2d, + dim3 grid1d, + dim3 block, + cudaStream_t stream, + bool interleaved, + bool use_vec, + bool qk_aligned16, + const scalar_t* cos_ptr, + const scalar_t* sin_ptr, + scalar_t* q_ptr, + scalar_t* k_ptr, + const int64_t* positions_ptr, + int rot_dim_from_cache, + int embed_dim_for_rotation, + int64_t query_token_stride, + int64_t key_token_stride, + int64_t head_stride_query, + int64_t head_stride_key, + int num_heads, + int num_kv_heads, + int head_size, + int blocks_per_token) { + auto dispatch_all = [&](auto f) { + if (interleaved) { + if (use_vec) { + if (qk_aligned16) + f(std::true_type{}, std::true_type{}, std::true_type{}); + else + f(std::true_type{}, std::true_type{}, std::false_type{}); + } else { + f(std::true_type{}, std::false_type{}, std::true_type{}); + } + } else { + if (use_vec) { + if (qk_aligned16) + f(std::false_type{}, std::true_type{}, std::true_type{}); + else + f(std::false_type{}, std::true_type{}, std::false_type{}); + } else { + f(std::false_type{}, std::false_type{}, std::true_type{}); + } + } + }; + + if (use_grid_2d) { + dispatch_all([&](auto i_v, auto v_v, auto a_v) { + constexpr bool I = decltype(i_v)::value; + constexpr bool V = decltype(v_v)::value; + constexpr bool A = decltype(a_v)::value; + host::LaunchKernel(grid2d, block, stream)( + rotary_embedding_kernel_2d, + cos_ptr, + sin_ptr, + q_ptr, + k_ptr, + positions_ptr, + rot_dim_from_cache, + embed_dim_for_rotation, + query_token_stride, + key_token_stride, + head_stride_query, + head_stride_key, + num_heads, + num_kv_heads, + head_size, + blocks_per_token); + }); + } else { + dispatch_all([&](auto i_v, auto v_v, auto a_v) { + constexpr bool I = decltype(i_v)::value; + constexpr bool V = decltype(v_v)::value; + constexpr bool A = decltype(a_v)::value; + host::LaunchKernel(grid1d, block, stream)( + rotary_embedding_kernel_1d, + cos_ptr, + sin_ptr, + q_ptr, + k_ptr, + positions_ptr, + rot_dim_from_cache, + embed_dim_for_rotation, + query_token_stride, + key_token_stride, + head_stride_query, + head_stride_key, + num_heads, + num_kv_heads, + head_size); + }); + } +} + +// Validate inputs, choose vectorization/alignment path, and launch the RoPE kernel. +// cos/sin are indexed by either positions[t] (if provided) or token index t. +// Rotates the first 2*embed_dim_for_rotation elements of each head. +template +inline void launch_rotary( + const tvm::ffi::TensorView cos, + const tvm::ffi::TensorView sin, + const tvm::ffi::TensorView q, + const tvm::ffi::TensorView* k, + const tvm::ffi::TensorView* positions, + int64_t head_size, + bool interleaved) { + using namespace host; + + auto T = SymbolicSize{"T"}; + auto Hq = SymbolicSize{"Hq"}; + auto Hk = SymbolicSize{"Hk"}; + auto D = SymbolicSize{"D"}; + auto R = SymbolicSize{"R"}; + auto dtype = SymbolicDType{}; + auto device = SymbolicDevice{}; + + // Verify q: [T, Hq, D] first to establish device/dtype + TensorMatcher({T, Hq, D}).with_dtype(dtype).with_device(device).verify(q); + + const int64_t t = T.unwrap(); + const int64_t hq = Hq.unwrap(); + const int64_t d = D.unwrap(); + + // Verify cos/sin: [T_cache, R] + auto T_cache = SymbolicSize{"T_cache"}; + TensorMatcher({T_cache, R}) + .with_dtype(dtype) + .with_device(device) + .verify(cos) + .verify(sin); + + const int64_t r = R.unwrap(); + + const int64_t* positions_ptr = nullptr; + if (positions != nullptr) { + auto T_pos = SymbolicSize{"T_pos"}; + TensorMatcher({T_pos}).with_dtype().with_device(device).verify(*positions); + RuntimeCheck(T_pos.unwrap() == t, "positions length mismatch"); + positions_ptr = static_cast(positions->data_ptr()); + } else { + // If positions is null, cos/sin length must equal to t + RuntimeCheck( + T_cache.unwrap() == t, + "cos/sin length mismatch (expected ", + t, + ", got ", + T_cache.unwrap(), + ") when positions is None"); + } + + RuntimeCheck(d == head_size, "head_size mismatch: got ", d, " expected ", head_size); + RuntimeCheck(t > 0 && hq > 0 && d > 0 && r > 0, "invalid shape"); + if (!interleaved) { + RuntimeCheck(r % 2 == 0, "non-interleaved requires even R, got ", r); + } + const int embed_dim_for_rotation = interleaved ? (int)r : (int)(r / 2); + RuntimeCheck(embed_dim_for_rotation > 0, "embed_dim_for_rotation must be > 0"); + if constexpr (ROT > 0) { + RuntimeCheck( + embed_dim_for_rotation == ROT, + "embed_dim_for_rotation mismatch: got ", + embed_dim_for_rotation, + " expected ROT=", + ROT); + } + RuntimeCheck(2LL * embed_dim_for_rotation <= head_size, "rotate dim exceeds head_size"); + + const int64_t query_token_stride = hq * d; + const int64_t head_stride_query = d; + + int64_t key_token_stride = 0; + int64_t head_stride_key = d; + int hk = 0; + if (k != nullptr) { + TensorMatcher({T, Hk, D}).with_dtype(dtype).with_device(device).verify(*k); + hk = (int)Hk.unwrap(); + RuntimeCheck(hk > 0, "invalid key shape"); + key_token_stride = (int64_t)hk * d; + } + + const int max_pairs_to_rotate_per_token = + (k == nullptr) ? ((int)hq * embed_dim_for_rotation) + : std::max((int)hq * embed_dim_for_rotation, (int)hk * embed_dim_for_rotation); + + constexpr int kVecBytes = 16; + const int elem_bytes = (int)sizeof(scalar_t); + const int kElePerVec = kVecBytes / elem_bytes; + const int pairs_per_step = interleaved ? (kElePerVec / 2) : kElePerVec; + + bool can_vec_compute = true; + if ((embed_dim_for_rotation % pairs_per_step) != 0) can_vec_compute = false; + if ((reinterpret_cast(cos.data_ptr()) % kVecBytes) != 0) can_vec_compute = false; + if ((reinterpret_cast(sin.data_ptr()) % kVecBytes) != 0) can_vec_compute = false; + if (((r * elem_bytes) % kVecBytes) != 0) can_vec_compute = false; + + bool qk_aligned16 = true; + if ((reinterpret_cast(q.data_ptr()) % kVecBytes) != 0) qk_aligned16 = false; + if (((query_token_stride * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false; + if (((head_stride_query * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false; + if (k != nullptr) { + if ((reinterpret_cast(k->data_ptr()) % kVecBytes) != 0) qk_aligned16 = false; + if (((key_token_stride * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false; + if (((head_stride_key * elem_bytes) % kVecBytes) != 0) qk_aligned16 = false; + } + + const bool use_vec = can_vec_compute; + const int launch_pairs_per_thread = use_vec ? pairs_per_step : 1; + const int total_threads_needed = + (max_pairs_to_rotate_per_token + launch_pairs_per_thread - 1) / launch_pairs_per_thread; + + auto round_up32 = [](int x) { return ((x + 31) / 32) * 32; }; + + // Case 1: 2D Grid (only when num_tokens<=4 and needs multiple blocks per token) + const int threads_per_block_2d = std::min(512, std::max(128, round_up32(std::min(total_threads_needed, 512)))); + const int blocks_per_token_2d = (total_threads_needed + threads_per_block_2d - 1) / threads_per_block_2d; + const bool use_grid_2d = ((int)t <= 4) && (blocks_per_token_2d > 1); + + // Case 2: 1D Grid + const int threads_per_block_1d = std::min(512, std::max(128, round_up32(total_threads_needed))); + + const int threads_per_block = use_grid_2d ? threads_per_block_2d : threads_per_block_1d; + const int blocks_per_token = use_grid_2d ? blocks_per_token_2d : 1; + + const auto stream = LaunchKernel::resolve_device(device.unwrap()); + const scalar_t* cos_ptr = static_cast(cos.data_ptr()); + const scalar_t* sin_ptr = static_cast(sin.data_ptr()); + scalar_t* q_ptr = static_cast(q.data_ptr()); + scalar_t* k_ptr = (k != nullptr) ? static_cast(k->data_ptr()) : nullptr; + + const dim3 grid2d((int)t, std::max(1, blocks_per_token)); + const dim3 grid1d((int)t); + const dim3 block(threads_per_block); + + dispatch_rotary_launch( + use_grid_2d, + grid2d, + grid1d, + block, + stream, + interleaved, + use_vec, + qk_aligned16, + cos_ptr, + sin_ptr, + q_ptr, + k_ptr, + positions_ptr, + (int)r, + embed_dim_for_rotation, + query_token_stride, + key_token_stride, + head_stride_query, + head_stride_key, + (int)hq, + (int)hk, + (int)head_size, + blocks_per_token); + + RuntimeDeviceCheck(); +} + +} // namespace + +template +struct RotaryEmbeddingCosSinKernel { + static void run_q( + const tvm::ffi::TensorView cos, + const tvm::ffi::TensorView sin, + const tvm::ffi::TensorView query, + const tvm::ffi::Optional positions, + int64_t head_size, + bool interleaved) { + const tvm::ffi::TensorView* positions_ptr = positions.has_value() ? &positions.value() : nullptr; + + const auto dt = query.dtype(); + if (host::is_type(dt)) { + launch_rotary(cos, sin, query, nullptr, positions_ptr, head_size, interleaved); + } else if (host::is_type(dt)) { + launch_rotary(cos, sin, query, nullptr, positions_ptr, head_size, interleaved); + } else if (host::is_type(dt)) { + launch_rotary(cos, sin, query, nullptr, positions_ptr, head_size, interleaved); + } else { + host::Panic("Unsupported dtype for rotary_embedding_cos_sin"); + } + } + + static void run_qk( + const tvm::ffi::TensorView cos, + const tvm::ffi::TensorView sin, + const tvm::ffi::TensorView query, + const tvm::ffi::TensorView key, + const tvm::ffi::Optional positions, + int64_t head_size, + bool interleaved) { + const tvm::ffi::TensorView* positions_ptr = positions.has_value() ? &positions.value() : nullptr; + + const auto dt = query.dtype(); + if (host::is_type(dt)) { + launch_rotary(cos, sin, query, &key, positions_ptr, head_size, interleaved); + } else if (host::is_type(dt)) { + launch_rotary(cos, sin, query, &key, positions_ptr, head_size, interleaved); + } else if (host::is_type(dt)) { + launch_rotary(cos, sin, query, &key, positions_ptr, head_size, interleaved); + } else { + host::Panic("Unsupported dtype for rotary_embedding_cos_sin"); + } + } +}; diff --git a/python/sglang/jit_kernel/rotary_embedding.py b/python/sglang/jit_kernel/rotary_embedding.py new file mode 100644 index 000000000000..da43dd59e32e --- /dev/null +++ b/python/sglang/jit_kernel/rotary_embedding.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.jit_kernel.utils import cache_once, load_jit + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +@cache_once +def _jit_rotary_embedding_cos_sin_module(rot: int) -> Module: + return load_jit( + "rotary_embedding_cos_sin", + str(rot), + cuda_files=["rotary_embedding_cos_sin.cuh"], + cuda_wrappers=[ + ( + "rotary_embedding_cos_sin_q", + f"RotaryEmbeddingCosSinKernel<{rot}>::run_q", + ), + ( + "rotary_embedding_cos_sin_qk", + f"RotaryEmbeddingCosSinKernel<{rot}>::run_qk", + ), + ], + ) + + +@cache_once +def can_use_rotary_embedding_cos_sin(rot: int) -> bool: + if rot <= 0: + logging.getLogger(__name__).warning(f"Invalid rot={rot}") + return False + try: + _jit_rotary_embedding_cos_sin_module(rot) + return True + except Exception as e: + logging.getLogger(__name__).warning(f"JIT compile failed (rot={rot}): {e}") + return False + + +def rotary_embedding_cos_sin( + cos: torch.Tensor, + sin: torch.Tensor, + query: torch.Tensor, + key: Optional[torch.Tensor] = None, + head_size: int = 0, + interleaved: Optional[bool] = None, + positions: Optional[torch.Tensor] = None, + *, + is_neox: Optional[bool] = None, +) -> None: + if is_neox is not None: + if interleaved is not None and interleaved != is_neox: + raise ValueError(f"Mismatch: is_neox={is_neox}, interleaved={interleaved}") + interleaved = is_neox + interleaved = interleaved if interleaved is not None else True + + if head_size == 0: + if query.dim() < 3: + raise ValueError("head_size must be provided when query is 2D") + head_size = query.shape[-1] + + def _prepare(t: torch.Tensor, name: str) -> torch.Tensor: + if t.device.type != "cuda": + raise ValueError(f"{name} must be CUDA") + if not t.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + if t.dtype != query.dtype: + raise ValueError(f"{name} dtype mismatch") + + # Reshape to [tokens, heads, head_size] + if t.dim() == 2: + return t.view(t.shape[0], -1, head_size) + if t.dim() == 3: + if t.shape[-1] != head_size: + raise ValueError(f"{name} head_size mismatch") + return t + if t.dim() == 4: + if t.shape[-1] != head_size: + raise ValueError(f"{name} head_size mismatch") + return t.flatten(0, 1) + raise ValueError(f"{name} must be 2D, 3D or 4D") + + q_3d = _prepare(query, "query") + k_3d = _prepare(key, "key") if key is not None else None + + if cos.device != query.device or sin.device != query.device: + raise ValueError("cos/sin device mismatch") + if cos.dtype != query.dtype or sin.dtype != query.dtype: + raise ValueError("cos/sin dtype mismatch") + if positions is None and cos.shape[0] != q_3d.shape[0]: + raise ValueError(f"cos/sin shape {cos.shape} mismatches tokens {q_3d.shape[0]}") + + if interleaved and cos.shape[1] == head_size: + if head_size % 2 != 0: + raise ValueError("interleaved layout requires even head_size") + half = head_size // 2 + cos = cos.view(cos.shape[0], half, 2).select(2, 0).contiguous() + sin = sin.view(sin.shape[0], half, 2).select(2, 1).contiguous() + else: + cos, sin = cos.contiguous(), sin.contiguous() + + rot_dim = cos.shape[1] + if rot_dim <= 0: + raise ValueError("rot_dim must be > 0") + + if interleaved: + if 2 * rot_dim > head_size: + raise ValueError(f"rot_dim {rot_dim} too large for interleaved") + embed_dim = rot_dim + else: + if rot_dim % 2 != 0: + raise ValueError("non-interleaved requires even rot_dim") + if rot_dim > head_size: + raise ValueError(f"rot_dim {rot_dim} too large") + embed_dim = rot_dim // 2 + + module = _jit_rotary_embedding_cos_sin_module(embed_dim) + if k_3d is not None: + module.rotary_embedding_cos_sin_qk( + cos, sin, q_3d, k_3d, positions, head_size, interleaved + ) + else: + module.rotary_embedding_cos_sin_q( + cos, sin, q_3d, positions, head_size, interleaved + ) diff --git a/python/sglang/jit_kernel/tests/test_rotary_embedding.py b/python/sglang/jit_kernel/tests/test_rotary_embedding.py new file mode 100644 index 000000000000..9c01c91fae8b --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_rotary_embedding.py @@ -0,0 +1,228 @@ +# pyright: reportMissingImports=false + +import torch +import triton + +try: + import sgl_kernel # noqa: F401 +except Exception: + sgl_kernel = None # type: ignore + +HAS_SGL_POS = hasattr(torch.ops.sgl_kernel, "rotary_embedding") + + +def _compute_cos_sin_cache_half( + max_seq_len: int, + rotary_dim: int, + base: float = 10000.0, + dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq = 1.0 / ( + base ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32) / rotary_dim) + ) + t = torch.arange(max_seq_len, dtype=torch.float32) + freqs = torch.einsum("i,j->ij", t, inv_freq) + cos = freqs.cos().to(dtype) + sin = freqs.sin().to(dtype) + return cos, sin + + +def sglang_aot_rotary_positions( + positions: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + head_size: int, + interleaved: bool, + cos_sin_cache: torch.Tensor, +) -> None: + if not HAS_SGL_POS: + raise RuntimeError("torch.ops.sgl_kernel.rotary_embedding is not available") + torch.ops.sgl_kernel.rotary_embedding( + positions, + q, + k, + head_size, + cos_sin_cache, + not interleaved, # sgl-kernel positions op uses is_neox (split-halves) flag + ) + + +def sglang_jit_rotary_with_positions( + positions: torch.Tensor, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + head_size: int, + interleaved: bool, +) -> None: + from sglang.jit_kernel.rotary_embedding import rotary_embedding_cos_sin + + rotary_embedding_cos_sin( + cos_cache, sin_cache, q, k, head_size, interleaved, positions=positions + ) + + +@torch.no_grad() +def torch_impl_rotary_fp32( + cos: torch.Tensor, + sin: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + head_size: int, + interleaved: bool, +) -> None: + orig_dtype = q.dtype + if interleaved and cos.shape[1] == head_size: + half = head_size // 2 + cos = cos.view(cos.shape[0], half, 2).select(2, 0).contiguous() + sin = sin.view(sin.shape[0], half, 2).select(2, 1).contiguous() + + cos_f, sin_f = cos.float(), sin.float() + q_f, k_f = q.float(), k.float() + + if interleaved: + embed_dim = int(cos_f.shape[1]) + rot_dim = embed_dim * 2 + cos_b = cos_f[:, None, :embed_dim] + sin_b = sin_f[:, None, :embed_dim] + + def _apply(x: torch.Tensor) -> None: + xr = x[..., :rot_dim] + xr2 = xr.view(xr.shape[0], xr.shape[1], embed_dim, 2) + x0 = xr2[..., 0].clone() + x1 = xr2[..., 1].clone() + xr2[..., 0].copy_(x0 * cos_b - x1 * sin_b) + xr2[..., 1].copy_(x1 * cos_b + x0 * sin_b) + + else: + if cos_f.shape[1] == head_size // 2: + embed_dim = int(cos_f.shape[1]) + rot_dim = embed_dim * 2 + cos_x, sin_x = cos_f[:, None, :], sin_f[:, None, :] + cos_y, sin_y = cos_x, sin_x + else: + embed_dim = int(cos_f.shape[1]) // 2 + rot_dim = embed_dim * 2 + cos_x, sin_x = cos_f[:, None, :embed_dim], sin_f[:, None, :embed_dim] + cos_y, sin_y = ( + cos_f[:, None, embed_dim:rot_dim], + sin_f[:, None, embed_dim:rot_dim], + ) + + def _apply(x: torch.Tensor) -> None: + xr = x[..., :rot_dim] + x0 = xr[..., :embed_dim].clone() + x1 = xr[..., embed_dim:rot_dim].clone() + xr[..., :embed_dim].copy_(x0 * cos_x - x1 * sin_x) + xr[..., embed_dim:rot_dim].copy_(x1 * cos_y + x0 * sin_y) + + _apply(q_f) + _apply(k_f) + q.copy_(q_f.to(orig_dtype)) + k.copy_(k_f.to(orig_dtype)) + + +def main(): + DEVICE = "cuda" + if not torch.cuda.is_available(): + raise SystemExit("CUDA not available") + + NUM_Q_HEADS = 32 + NUM_KV_HEADS = 8 + BS_LIST = [2**n for n in range(0, 13)] + BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)] + + for DTYPE in [torch.bfloat16, torch.float16]: + for HEAD_SIZE in [64, 80, 96, 128, 256, 320]: + rotary_dim = HEAD_SIZE + max_seq_len = 8192 + cos_cache, sin_cache = _compute_cos_sin_cache_half( + max_seq_len, rotary_dim, dtype=DTYPE + ) + cos_cache, sin_cache = cos_cache.to(DEVICE), sin_cache.to(DEVICE) + + for INTERLEAVED in [True, False]: + aot_cos_sin_cache = torch.cat( + [cos_cache, sin_cache], dim=-1 + ).contiguous() + + for BS in BS_LIST: + positions = ( + torch.arange(BS, device=DEVICE, dtype=torch.int64) % max_seq_len + ) + cos_half = cos_cache[positions].contiguous() + sin_half = sin_cache[positions].contiguous() + + if INTERLEAVED: + # interleaved: cache R = embed_dim = head_size/2 + cos_g, sin_g = cos_half, sin_half + cos_cache_for_jit, sin_cache_for_jit = cos_cache, sin_cache + else: + # non-interleaved: kernel expects R = head_size (x/y halves) + cos_g = torch.cat([cos_half, cos_half], dim=-1).contiguous() + sin_g = torch.cat([sin_half, sin_half], dim=-1).contiguous() + cos_cache_for_jit = torch.cat( + [cos_cache, cos_cache], dim=-1 + ).contiguous() + sin_cache_for_jit = torch.cat( + [sin_cache, sin_cache], dim=-1 + ).contiguous() + + q = torch.randn( + BS, NUM_Q_HEADS, HEAD_SIZE, device=DEVICE, dtype=DTYPE + ) + k = torch.randn( + BS, NUM_KV_HEADS, HEAD_SIZE, device=DEVICE, dtype=DTYPE + ) + + q_ref_fp32, k_ref_fp32 = q.clone(), k.clone() + torch_impl_rotary_fp32( + cos_g, sin_g, q_ref_fp32, k_ref_fp32, HEAD_SIZE, INTERLEAVED + ) + + q_k_aot = (q.clone(), k.clone()) + q_k_jit = (q.clone(), k.clone()) + if HAS_SGL_POS: + sglang_aot_rotary_positions( + positions, + q_k_aot[0], + q_k_aot[1], + HEAD_SIZE, + INTERLEAVED, + aot_cos_sin_cache, + ) + sglang_jit_rotary_with_positions( + positions, + cos_cache_for_jit, + sin_cache_for_jit, + q_k_jit[0], + q_k_jit[1], + HEAD_SIZE, + INTERLEAVED, + ) + + ref_atol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3 + ref_rtol = 2e-2 if DTYPE == torch.bfloat16 else 2e-3 + + if HAS_SGL_POS: + triton.testing.assert_close( + q_ref_fp32, q_k_aot[0], atol=ref_atol, rtol=ref_rtol + ) + triton.testing.assert_close( + k_ref_fp32, q_k_aot[1], atol=ref_atol, rtol=ref_rtol + ) + triton.testing.assert_close( + q_ref_fp32, q_k_jit[0], atol=ref_atol, rtol=ref_rtol + ) + triton.testing.assert_close( + k_ref_fp32, q_k_jit[1], atol=ref_atol, rtol=ref_rtol + ) + + print( + f"HEAD_SIZE={HEAD_SIZE} interleaved={INTERLEAVED} dtype={DTYPE} passed." + ) + + +if __name__ == "__main__": + main() diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py index 2ef943229c6b..72fd65d836d3 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py +++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding.py @@ -123,6 +123,124 @@ def apply_flashinfer_rope_qk_inplace( return q_flat.view(bsz, seqlen, nheads, d), k_flat.view(bsz, seqlen, nheads, d) +def _split_cos_sin_from_cache( + cache: torch.Tensor, *, dtype: torch.dtype +) -> tuple[torch.Tensor, torch.Tensor]: + """Convert a 2D RoPE cache into (cos, sin) tensors.""" + if cache.is_complex(): + cos = cache.real + sin = cache.imag + if cos.dtype != dtype: + cos = cos.to(dtype) + if sin.dtype != dtype: + sin = sin.to(dtype) + return cos, sin + if cache.shape[1] % 2 != 0: + raise ValueError( + f"Expected complex freqs_cis or cat([cos,sin]) cache; got real cache with odd last dim={cache.shape[1]}" + ) + half = cache.shape[1] // 2 + cos = cache[:, :half] + sin = cache[:, half:] + if cos.dtype != dtype: + cos = cos.to(dtype) + if sin.dtype != dtype: + sin = sin.to(dtype) + return cos, sin + + +def apply_sglang_jit_rope_qk_inplace( + q: torch.Tensor, + k: torch.Tensor, + cos_sin_cache: torch.Tensor, + *, + head_size: Optional[int] = None, + is_neox: bool = False, + positions: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Apply SGLang JIT RoPE on GPU using `sglang.jit_kernel.rotary_embedding_cos_sin`.""" + if q.dim() != 4 or k.dim() != 4: + raise ValueError( + f"Expected q/k to be 4D [bsz, seqlen, nheads, head_size], " + f"got q:{tuple(q.shape)} k:{tuple(k.shape)}" + ) + if q.shape != k.shape: + raise ValueError( + f"q and k must have the same shape, got {q.shape} vs {k.shape}" + ) + + if not (isinstance(cos_sin_cache, torch.Tensor) and cos_sin_cache.dim() == 2): + raise ValueError("cos_sin_cache must be a 2D torch.Tensor") + + bsz, seqlen, nheads, d = q.shape + if head_size is None: + head_size = d + if head_size != d: + raise ValueError(f"head_size mismatch: inferred {d}, but head_size={head_size}") + + from sglang.jit_kernel.rotary_embedding import ( + rotary_embedding_cos_sin as sglang_jit_rotary_embedding_cos_sin, + ) + + num_tokens = bsz * seqlen + + cache_tok: torch.Tensor + if cos_sin_cache.shape[0] < seqlen: + # Should not happen if cache is sufficient + pass + + if positions is None: + if bsz > 1: + # For bsz > 1, we need explicit positions [0..seqlen, 0..seqlen] + # Create directly on device to avoid CPU overhead + pos_1d = torch.arange(seqlen, device=q.device, dtype=torch.long) + positions = pos_1d.repeat(bsz) + else: + # For bsz == 1, positions=None implies 0..seqlen, which Kernel supports natively + pass + else: + if not ( + isinstance(positions, torch.Tensor) + and positions.dtype == torch.long + and positions.dim() == 1 + ): + raise ValueError("positions must be a 1D torch.long Tensor") + if positions.numel() != num_tokens: + raise ValueError( + f"positions length must be bsz*seqlen={num_tokens}, got {positions.numel()}" + ) + positions = positions.to(q.device, non_blocking=True) + + q3 = q.view(num_tokens, nheads, d) + k3 = k.view(num_tokens, nheads, d) + + # Use the FULL cache (no index_select) + cache_full = cos_sin_cache + + # Include is_neox in key because it affects how we process (cat vs contiguous) + interleaved = not is_neox + + # Split and Cast the FULL cache + cos_full, sin_full = _split_cos_sin_from_cache(cache_full, dtype=q.dtype) + + # Process layout (Contiguous or Cat) + if interleaved: + cos = cos_full.contiguous() + sin = sin_full.contiguous() + else: + if cos_full.shape[1] * 2 == head_size: + cos = torch.cat([cos_full, cos_full], dim=-1).contiguous() + sin = torch.cat([sin_full, sin_full], dim=-1).contiguous() + else: + cos = cos_full.contiguous() + sin = sin_full.contiguous() + + sglang_jit_rotary_embedding_cos_sin( + cos, sin, q3, k3, head_size, interleaved, positions=positions + ) + return q3.view(bsz, seqlen, nheads, d), k3.view(bsz, seqlen, nheads, d) + + def _rotate_neox(x: torch.Tensor) -> torch.Tensor: x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :]