diff --git a/python/sglang/jit_kernel/benchmark/bench_rope.py b/python/sglang/jit_kernel/benchmark/bench_rope.py new file mode 100644 index 000000000..8d5cdae27 --- /dev/null +++ b/python/sglang/jit_kernel/benchmark/bench_rope.py @@ -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) diff --git a/python/sglang/jit_kernel/csrc/elementwise/rope.cuh b/python/sglang/jit_kernel/csrc/elementwise/rope.cuh index c19e9abc7..27b4e7ec8 100644 --- a/python/sglang/jit_kernel/csrc/elementwise/rope.cuh +++ b/python/sglang/jit_kernel/csrc/elementwise/rope.cuh @@ -1,655 +1,463 @@ -/* - * Copyright (c) 2024 by FlashInfer team. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - #include #include +#include +#include +#include #include +#include -#include // upstream -#include +#include -#include -#include - -namespace flashinfer { - -namespace kv_buffer_saver { - -template -__device__ __forceinline__ void prepare( - vec_t& v_vec, - IdType& kv_cache_offset, - DType* v, - IdType* kv_cache_loc, - uint32_t idx, - uint32_t tx, - uint32_t kv_head_idx, - size_t v_stride_n, - size_t v_stride_h) { - kv_cache_offset = kv_cache_loc[idx]; - - DType* v_ptr = v + get_elem_offset_impl(idx, kv_head_idx, 0, v_stride_n, v_stride_h); - v_vec.cast_load(v_ptr + tx * vec_size); -} - -template -__device__ __forceinline__ void save( - IdType& kv_cache_offset, - vec_t& k_vec, - vec_t& v_vec, - DType* k_buffer, - DType* v_buffer, - uint32_t idx, - uint32_t tx, - uint32_t kv_head_idx, - size_t k_buffer_stride_n, - size_t k_buffer_stride_h, - size_t v_buffer_stride_n, - size_t v_buffer_stride_h) { - DType* k_buffer_ptr = - k_buffer + get_elem_offset_impl(kv_cache_offset, kv_head_idx, 0, k_buffer_stride_n, k_buffer_stride_h); - DType* v_buffer_ptr = - v_buffer + get_elem_offset_impl(kv_cache_offset, kv_head_idx, 0, v_buffer_stride_n, v_buffer_stride_h); - k_vec.cast_store(k_buffer_ptr + tx * vec_size); - v_vec.cast_store(v_buffer_ptr + tx * vec_size); -} - -} // namespace kv_buffer_saver - -template < - bool save_kv_cache, - bool interleave, - uint32_t head_dim, - uint32_t vec_size, - uint32_t bdx, - typename DType, - typename IdType> -__global__ void BatchQKApplyRotaryPosIdsCosSinCacheEnhancedHeadParallelismKernel( - DType* q, - DType* k, - DType* v, - DType* q_rope, - DType* k_rope, - DType* k_buffer, - DType* v_buffer, - float* __restrict__ cos_sin_cache, - IdType* __restrict__ pos_ids, - uint32_t nnz, - uint32_t num_qo_heads, - uint32_t num_kv_heads, - uint32_t rotary_dim, - size_t q_stride_n, - size_t q_stride_h, - size_t k_stride_n, - size_t k_stride_h, - size_t v_stride_n, - size_t v_stride_h, - size_t q_rope_stride_n, - size_t q_rope_stride_h, - size_t k_rope_stride_n, - size_t k_rope_stride_h, - size_t k_buffer_stride_n, - size_t k_buffer_stride_h, - size_t v_buffer_stride_n, - size_t v_buffer_stride_h, - IdType* __restrict__ kv_cache_loc) { - uint32_t bx = blockIdx.x, tx = threadIdx.x, ty = threadIdx.y; - uint32_t by = blockIdx.y; - const uint32_t bdy = blockDim.y; - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); -#endif - - vec_t cos, sin; - if (bx * bdy + ty < nnz) { - const uint32_t idx = bx * bdy + ty; - const IdType pos = pos_ids[idx]; - - const int half_rotary_dim = rotary_dim / 2; - - // 1. if interleave: - // - cos = cos_sin_cache[pos_id][tx * vec_size // 2] - // - sin = cos_sin_cache[pos_id][(rot_dim // 2) + tx * vec_size // 2] - // 2. if not interleave - // - cos = cos_cache[pos_id][(tx * vec_size) % (rot_dim // 2)] - // - sin = sin_cache[pos_id][(rot_dim // 2) + (tx * vec_size) % (rot_dim // 2)] - if (tx * vec_size < rotary_dim) { - int sin_offset = rotary_dim / 2; - int vec_idx; - if constexpr (interleave) { - vec_idx = (tx * vec_size) / 2; // Force integer division - } else { - vec_idx = (tx * vec_size) % half_rotary_dim; // Use half_rotary_dim - } - cos.load(cos_sin_cache + (pos * rotary_dim) + vec_idx); - sin.load(cos_sin_cache + (pos * rotary_dim) + (sin_offset + vec_idx)); - } - - if (by < num_qo_heads) { - uint32_t qo_head_idx = by; - DType* q_ptr = q + get_elem_offset_impl(idx, qo_head_idx, 0, q_stride_n, q_stride_h); - DType* q_rope_ptr = q_rope + get_elem_offset_impl(idx, qo_head_idx, 0, q_rope_stride_n, q_rope_stride_h); - vec_t q_vec; - if constexpr (interleave) { - q_vec = vec_apply_llama_rope_cos_sin_interleave_reuse_half(q_ptr, cos, sin, rotary_dim); - } else { - q_vec = vec_apply_llama_rope_cos_sin(q_ptr, cos, sin, rotary_dim); - } - q_vec.cast_store(q_rope_ptr + tx * vec_size); - } else { - uint32_t kv_head_idx = by - num_qo_heads; - DType* k_ptr = k + get_elem_offset_impl(idx, kv_head_idx, 0, k_stride_n, k_stride_h); - - DType* k_rope_ptr = k_rope + get_elem_offset_impl(idx, kv_head_idx, 0, k_rope_stride_n, k_rope_stride_h); - - vec_t v_vec; - IdType kv_cache_offset; - if constexpr (save_kv_cache) { - kv_buffer_saver::prepare( - v_vec, kv_cache_offset, v, kv_cache_loc, idx, tx, kv_head_idx, v_stride_n, v_stride_h); - } - - vec_t k_vec; - if constexpr (interleave) { - k_vec = vec_apply_llama_rope_cos_sin_interleave_reuse_half(k_ptr, cos, sin, rotary_dim); - } else { - k_vec = vec_apply_llama_rope_cos_sin(k_ptr, cos, sin, rotary_dim); - } - k_vec.cast_store(k_rope_ptr + tx * vec_size); - - if constexpr (save_kv_cache) { - kv_buffer_saver::save( - kv_cache_offset, - k_vec, - v_vec, - k_buffer, - v_buffer, - idx, - tx, - kv_head_idx, - k_buffer_stride_n, - k_buffer_stride_h, - v_buffer_stride_n, - v_buffer_stride_h); - } - } - } - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); -#endif -} - -template < - bool save_kv_cache, - bool interleave, - uint32_t head_dim, - uint32_t vec_size, - uint32_t bdx, - typename DType, - typename IdType> -__global__ void BatchQKApplyRotaryPosIdsCosSinCacheEnhancedKernel( - DType* q, - DType* k, - DType* v, - DType* q_rope, - DType* k_rope, - DType* k_buffer, - DType* v_buffer, - float* __restrict__ cos_sin_cache, - IdType* __restrict__ pos_ids, - uint32_t nnz, - uint32_t num_qo_heads, - uint32_t num_kv_heads, - uint32_t rotary_dim, - size_t q_stride_n, - size_t q_stride_h, - size_t k_stride_n, - size_t k_stride_h, - size_t v_stride_n, - size_t v_stride_h, - size_t q_rope_stride_n, - size_t q_rope_stride_h, - size_t k_rope_stride_n, - size_t k_rope_stride_h, - size_t k_buffer_stride_n, - size_t k_buffer_stride_h, - size_t v_buffer_stride_n, - size_t v_buffer_stride_h, - IdType* __restrict__ kv_cache_loc) { - uint32_t bx = blockIdx.x, tx = threadIdx.x, ty = threadIdx.y; - const uint32_t bdy = blockDim.y; - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.wait;"); -#endif - - vec_t cos, sin; - if (bx * bdy + ty < nnz) { - const uint32_t idx = bx * bdy + ty; - const IdType pos = pos_ids[idx]; - const int half_rotary_dim = rotary_dim / 2; - - // 1. if interleave: - // - cos = cos_sin_cache[pos_id][tx * vec_size // 2] - // - sin = cos_sin_cache[pos_id][(rot_dim // 2) + tx * vec_size // 2] - // 2. if not interleave - // - cos = cos_cache[pos_id][(tx * vec_size) % (rot_dim // 2)] - // - sin = sin_cache[pos_id][(rot_dim // 2) + (tx * vec_size) % (rot_dim // 2)] - if (tx * vec_size < rotary_dim) { - int sin_offset = rotary_dim / 2; - int vec_idx; - if constexpr (interleave) { - vec_idx = (tx * vec_size) / 2; // Force integer division - } else { - vec_idx = (tx * vec_size) % half_rotary_dim; // Use half_rotary_dim - } - cos.load(cos_sin_cache + (pos * rotary_dim) + vec_idx); - sin.load(cos_sin_cache + (pos * rotary_dim) + (sin_offset + vec_idx)); - } - - // not to unroll the loop, because num head might be large and might lead to worse performance -#pragma unroll 1 - for (uint32_t qo_head_idx = 0; qo_head_idx < num_qo_heads; ++qo_head_idx) { - DType* q_ptr = q + get_elem_offset_impl(idx, qo_head_idx, 0, q_stride_n, q_stride_h); - DType* q_rope_ptr = q_rope + get_elem_offset_impl(idx, qo_head_idx, 0, q_rope_stride_n, q_rope_stride_h); - vec_t q_vec; - if constexpr (interleave) { - q_vec = vec_apply_llama_rope_cos_sin_interleave_reuse_half(q_ptr, cos, sin, rotary_dim); - } else { - q_vec = vec_apply_llama_rope_cos_sin(q_ptr, cos, sin, rotary_dim); - } - q_vec.cast_store(q_rope_ptr + tx * vec_size); - } - -#pragma unroll 1 - for (uint32_t kv_head_idx = 0; kv_head_idx < num_kv_heads; ++kv_head_idx) { - DType* k_ptr = k + get_elem_offset_impl(idx, kv_head_idx, 0, k_stride_n, k_stride_h); - - DType* k_rope_ptr = k_rope + get_elem_offset_impl(idx, kv_head_idx, 0, k_rope_stride_n, k_rope_stride_h); - - vec_t v_vec; - IdType kv_cache_offset; - if constexpr (save_kv_cache) { - kv_buffer_saver::prepare( - v_vec, kv_cache_offset, v, kv_cache_loc, idx, tx, kv_head_idx, v_stride_n, v_stride_h); - } - - vec_t k_vec; - if constexpr (interleave) { - k_vec = vec_apply_llama_rope_cos_sin_interleave_reuse_half(k_ptr, cos, sin, rotary_dim); - } else { - k_vec = vec_apply_llama_rope_cos_sin(k_ptr, cos, sin, rotary_dim); - } - k_vec.cast_store(k_rope_ptr + tx * vec_size); - - if constexpr (save_kv_cache) { - kv_buffer_saver::save( - kv_cache_offset, - k_vec, - v_vec, - k_buffer, - v_buffer, - idx, - tx, - kv_head_idx, - k_buffer_stride_n, - k_buffer_stride_h, - v_buffer_stride_n, - v_buffer_stride_h); - } - } - } - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - asm volatile("griddepcontrol.launch_dependents;"); -#endif -} - -#define DISPATCH_SAVE_KV_CACHE(save_kv_cache, SAVE_KV_CACHE, ...) \ - if (save_kv_cache) { \ - const bool SAVE_KV_CACHE = true; \ - __VA_ARGS__ \ - } else { \ - const bool SAVE_KV_CACHE = false; \ - __VA_ARGS__ \ - } - -template -cudaError_t BatchQKApplyRotaryPosIdsCosSinCacheEnhanced( - DType* q, - DType* k, - DType* v, - DType* q_rope, - DType* k_rope, - DType* k_buffer, - DType* v_buffer, - float* cos_sin_cache, - IdType* pos_ids, - uint32_t nnz, - uint32_t num_qo_heads, - uint32_t num_kv_heads, - uint32_t rotary_dim, - uint32_t head_dim, - size_t q_stride_n, - size_t q_stride_h, - size_t k_stride_n, - size_t k_stride_h, - size_t v_stride_n, - size_t v_stride_h, - size_t q_rope_stride_n, - size_t q_rope_stride_h, - size_t k_rope_stride_n, - size_t k_rope_stride_h, - size_t k_buffer_stride_n, - size_t k_buffer_stride_h, - size_t v_buffer_stride_n, - size_t v_buffer_stride_h, - IdType* kv_cache_loc, - bool interleave, - bool save_kv_cache, - bool enable_pdl, - cudaStream_t stream = nullptr) { - int dev_id = 0; - int num_sms = 0; - FLASHINFER_CUDA_CALL(cudaGetDevice(&dev_id)); - FLASHINFER_CUDA_CALL(cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, dev_id)); - -#define LAUNCH_KERNEL_RAW(kernel_name) \ - do { \ - cudaLaunchConfig_t config = {}; \ - config.gridDim = nblks; \ - config.blockDim = nthrs; \ - config.dynamicSmemBytes = 0; \ - config.stream = stream; \ - cudaLaunchAttribute attrs[1] = {}; \ - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \ - attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl; \ - config.numAttrs = 1; \ - config.attrs = attrs; \ - \ - FLASHINFER_CUDA_CALL(cudaLaunchKernelEx( \ - &config, \ - kernel_name, \ - q, \ - k, \ - v, \ - q_rope, \ - k_rope, \ - k_buffer, \ - v_buffer, \ - cos_sin_cache, \ - pos_ids, \ - nnz, \ - num_qo_heads, \ - num_kv_heads, \ - rotary_dim, \ - q_stride_n, \ - q_stride_h, \ - k_stride_n, \ - k_stride_h, \ - v_stride_n, \ - v_stride_h, \ - q_rope_stride_n, \ - q_rope_stride_h, \ - k_rope_stride_n, \ - k_rope_stride_h, \ - k_buffer_stride_n, \ - k_buffer_stride_h, \ - v_buffer_stride_n, \ - v_buffer_stride_h, \ - kv_cache_loc)); \ - } while (0) - - DISPATCH_SAVE_KV_CACHE(save_kv_cache, SAVE_KV_CACHE, { - DISPATCH_INTERLEAVE(interleave, INTERLEAVE, { - DISPATCH_HEAD_DIM(head_dim, HEAD_DIM, { - // operate on 16 Bytes at a time - constexpr uint32_t vec_size = std::max(16 / sizeof(DType), HEAD_DIM / 32); - // how many threads needed per head_dim - constexpr uint32_t bdx = HEAD_DIM / vec_size; - // how many threads needed per block - uint32_t num_threads = std::max(128U, bdx); - // how many tokens can we process in a block - uint32_t bdy = num_threads / bdx; - // how many blocks needed to process all tokens - uint32_t nblks_x = (nnz + bdy - 1) / bdy; - - auto kernel_0 = BatchQKApplyRotaryPosIdsCosSinCacheEnhancedKernel< - SAVE_KV_CACHE, - INTERLEAVE, - HEAD_DIM, - vec_size, - bdx, - DType, - IdType>; - - int num_blocks_per_sm_0 = 0; - FLASHINFER_CUDA_CALL(cudaOccupancyMaxActiveBlocksPerMultiprocessor( - &num_blocks_per_sm_0, kernel_0, num_threads, /*smem_size=*/0)); - uint32_t num_ctas_0 = num_blocks_per_sm_0 * num_sms; - - if ((nnz + bdy - 1) / bdy >= num_ctas_0) { - dim3 nblks(nblks_x); - dim3 nthrs(bdx, bdy); - LAUNCH_KERNEL_RAW(kernel_0); - } else { - dim3 nblks(nblks_x, num_qo_heads + num_kv_heads); - dim3 nthrs(bdx, bdy); - auto kernel_1 = BatchQKApplyRotaryPosIdsCosSinCacheEnhancedHeadParallelismKernel< - SAVE_KV_CACHE, - INTERLEAVE, - HEAD_DIM, - vec_size, - bdx, - DType, - IdType>; - LAUNCH_KERNEL_RAW(kernel_1); - } - }); - }); - }); -#undef LAUNCH_KERNEL_RAW - - return cudaSuccess; -} - -} // namespace flashinfer +#include namespace { -#define DISPATCH_TVM_DTYPE_TO_CTYPE(tvm_dtype_code, tvm_dtype_bits, c_type, ...) \ - [&]() -> bool { \ - if (tvm_dtype_code == kDLFloat && tvm_dtype_bits == 32) { \ - using c_type = float; \ - return __VA_ARGS__(); \ - } \ - if (tvm_dtype_code == kDLFloat && tvm_dtype_bits == 16) { \ - using c_type = half; \ - return __VA_ARGS__(); \ - } \ - if (tvm_dtype_code == kDLBfloat && tvm_dtype_bits == 16) { \ - using c_type = nv_bfloat16; \ - return __VA_ARGS__(); \ - } \ - RuntimeCheck(false, "Unsupported data type. Only float32, float16, and bfloat16 are supported."); \ - return false; \ - }() +struct FusedRopeParams { + void* __restrict__ q_ptr; + void* __restrict__ k_ptr; // NOTE: this k is pre-offset in host code to reduce computation in kernel + const void* __restrict__ cos_sin_cache_ptr; + const void* __restrict__ positions; + int64_t q_stride_bytes; + int64_t k_stride_bytes; + int64_t head_stride_bytes; + uint32_t num_qo_heads; + uint32_t num_kv_heads; + uint32_t num_tokens; +}; -inline void check_cuda_contiguous(tvm::ffi::TensorView x) { - using namespace host; +struct FusedRopeStoreParams { + FusedRopeParams base_params; + void* v_ptr; + void* __restrict__ k_cache; + void* __restrict__ v_cache; + const void* __restrict__ out_loc; + int64_t v_stride_bytes; + int64_t cache_stride_bytes; +}; - RuntimeCheck(x.device().device_type == kDLCUDA); - RuntimeCheck(x.is_contiguous()); +constexpr uint32_t kBlockSize = 128; + +[[maybe_unused]] +constexpr auto next_pow2(uint32_t target, uint32_t factor = 1) { + uint32_t power = 1; + while (power * factor < target) + power *= 2; + return power; } -struct ApplyRopePosIdsCosSinCacheKernel { +template +__global__ void fused_rope_kernel(const __grid_constant__ FusedRopeParams params) { + using namespace device; + + constexpr int64_t kCosSinStrideBytes = kRopeDim * sizeof(float); + constexpr int64_t kVecSize = next_pow2(kRopeDim, (2 * kWorkThreads * (1 + kIsNeox))); + using DType2 = packed_t; + using InputStorage = AlignedVector; + constexpr int64_t kDimPerThread = kVecSize * 2 * (1 + kIsNeox); + constexpr uint32_t kLaneCount = kRopeDim / kDimPerThread; + static_assert(kRopeDim % kDimPerThread == 0 && kLaneCount <= kWorkThreads); + + const auto &[ + q, k, cos_sin_cache_ptr, positions, // pointers + q_stride_bytes, k_stride_bytes, head_stride_bytes, // strides + num_qo_heads, num_kv_heads, num_tokens // dimensions + ] = params; + + const auto num_blks = gridDim.x; + constexpr auto kWorkersPerBlock = kBlockSize / kWorkThreads; + const auto num_workers = num_blks * kWorkersPerBlock; + const auto num_q_and_k_heads = num_qo_heads + num_kv_heads; + const auto num_works = num_q_and_k_heads * num_tokens; + const auto start_worker_id = (blockIdx.x * kBlockSize + threadIdx.x) / kWorkThreads; + const auto cos_cache_ptr = cos_sin_cache_ptr; + const auto sin_cache_ptr = pointer::offset(cos_sin_cache_ptr, kCosSinStrideBytes / 2); + + uint32_t lane_id = threadIdx.x % kWorkThreads; + if constexpr (kLaneCount < kWorkThreads) { + if (lane_id >= kLaneCount) return; + } + + PDLWaitPrimary(); + + for (auto idx = start_worker_id; idx < num_works; idx += num_workers) { + const int64_t token_id = idx / num_q_and_k_heads; + const int64_t head_id = idx % num_q_and_k_heads; + const auto pos = static_cast(positions)[token_id]; + const auto load_q = head_id < num_qo_heads; + const auto input_ = load_q ? pointer::offset(q, token_id * q_stride_bytes) // + : pointer::offset(k, token_id * k_stride_bytes); + const auto input = pointer::offset(input_, head_id * head_stride_bytes); + const auto cos_ptr = pointer::offset(cos_cache_ptr, pos * kCosSinStrideBytes); + const auto sin_ptr = pointer::offset(sin_cache_ptr, pos * kCosSinStrideBytes); + if constexpr (kIsNeox) { + using CacheStorage = AlignedVector; + const auto input_x = input; + const auto input_y = pointer::offset(input, (kRopeDim / 2) * sizeof(DType)); + auto input_vec_x = load_as(input_x, lane_id); + auto input_vec_y = load_as(input_y, lane_id); + const auto cos_pair = load_as(cos_ptr, lane_id); + const auto sin_pair = load_as(sin_ptr, lane_id); +#pragma unroll + for (int64_t j = 0; j < kVecSize; ++j) { + const auto [x0, x1] = cast(input_vec_x[j]); + const auto [y0, y1] = cast(input_vec_y[j]); + const auto [cos_0, cos_1] = cos_pair[j]; + const auto [sin_0, sin_1] = sin_pair[j]; + const auto out_x0 = x0 * cos_0 - y0 * sin_0; + const auto out_y0 = x0 * sin_0 + y0 * cos_0; + const auto out_x1 = x1 * cos_1 - y1 * sin_1; + const auto out_y1 = x1 * sin_1 + y1 * cos_1; + input_vec_x[j] = cast({out_x0, out_x1}); + input_vec_y[j] = cast({out_y0, out_y1}); + } + store_as(input_x, input_vec_x, lane_id); + store_as(input_y, input_vec_y, lane_id); + } else { + using CacheStorage = AlignedVector; + auto input_vec = load_as(input, lane_id); + const auto cos_vec = load_as(cos_ptr, lane_id); + const auto sin_vec = load_as(sin_ptr, lane_id); +#pragma unroll + for (int64_t j = 0; j < kVecSize; ++j) { + const auto [x, y] = cast(input_vec[j]); + const auto cos = cos_vec[j]; + const auto sin = sin_vec[j]; + const auto out_x = x * cos - y * sin; + const auto out_y = x * sin + y * cos; + input_vec[j] = cast({out_x, out_y}); + } + store_as(input, input_vec, lane_id); + } + } + + PDLTriggerSecondary(); +} + +template +__global__ void fused_rope_store_kernel(const __grid_constant__ FusedRopeStoreParams params) { + using namespace device; + + constexpr int64_t kCosSinStrideBytes = kRopeDim * sizeof(float); + constexpr int64_t kVecSize = kRopeDim / (2 * kWorkThreads * (1 + kIsNeox)); + using DType2 = packed_t; + using InputStorage = AlignedVector; + constexpr int64_t kDimPerThread = kVecSize * 2 * (1 + kIsNeox); + static_assert(kRopeDim == kDimPerThread * kWorkThreads); + + const auto& [base_params, v_ptr, k_cache, v_cache, out_loc, v_stride_bytes, cache_stride_bytes] = params; + const auto &[ + q, k, cos_sin_cache_ptr, positions, // pointers + q_stride_bytes, k_stride_bytes, head_stride_bytes, // strides + num_qo_heads, num_kv_heads, num_tokens // dimensions + ] = base_params; + + const auto num_blks = gridDim.x; + constexpr auto kWorkersPerBlock = kBlockSize / kWorkThreads; + const auto num_workers = num_blks * kWorkersPerBlock; + const auto num_q_and_k_heads = num_qo_heads + num_kv_heads; + const auto num_works = num_q_and_k_heads * num_tokens; + const auto num_extra_works = num_kv_heads * num_tokens; // rope works + v store works + const auto start_worker_id = (blockIdx.x * kBlockSize + threadIdx.x) / kWorkThreads; + const auto lane_id = threadIdx.x % kWorkThreads; + const auto cos_cache_ptr = cos_sin_cache_ptr; + const auto sin_cache_ptr = pointer::offset(cos_sin_cache_ptr, kCosSinStrideBytes / 2); + + auto idx = start_worker_id; + + PDLWaitPrimary(); + // in this case, head_dim = rope_dim must be true + __builtin_assume(head_stride_bytes == kRopeDim * sizeof(DType)); + + for (; idx < num_works; idx += num_workers) { + const int64_t token_id = idx / num_q_and_k_heads; + const int64_t head_id = idx % num_q_and_k_heads; + const auto pos = static_cast(positions)[token_id]; + const auto loc = static_cast(out_loc)[token_id]; + const auto load_q = head_id < num_qo_heads; + const auto input_ = load_q ? pointer::offset(q, token_id * q_stride_bytes) // + : pointer::offset(k, token_id * k_stride_bytes); + const auto input = pointer::offset(input_, head_id * head_stride_bytes); + const auto cos_ptr = pointer::offset(cos_cache_ptr, pos * kCosSinStrideBytes); + const auto sin_ptr = pointer::offset(sin_cache_ptr, pos * kCosSinStrideBytes); + if constexpr (kIsNeox) { + using CacheStorage = AlignedVector; + const auto input_x = input; + const auto input_y = pointer::offset(input, (kRopeDim / 2) * sizeof(DType)); + auto input_vec_x = load_as(input_x, lane_id); + auto input_vec_y = load_as(input_y, lane_id); + const auto cos_pair = load_as(cos_ptr, lane_id); + const auto sin_pair = load_as(sin_ptr, lane_id); +#pragma unroll + for (int64_t j = 0; j < kVecSize; ++j) { + const auto [x0, x1] = cast(input_vec_x[j]); + const auto [y0, y1] = cast(input_vec_y[j]); + const auto [cos_0, cos_1] = cos_pair[j]; + const auto [sin_0, sin_1] = sin_pair[j]; + const auto out_x0 = x0 * cos_0 - y0 * sin_0; + const auto out_y0 = x0 * sin_0 + y0 * cos_0; + const auto out_x1 = x1 * cos_1 - y1 * sin_1; + const auto out_y1 = x1 * sin_1 + y1 * cos_1; + input_vec_x[j] = cast({out_x0, out_x1}); + input_vec_y[j] = cast({out_y0, out_y1}); + } + const auto k_out = pointer::offset(k_cache, loc * cache_stride_bytes, head_id * head_stride_bytes); + const auto output_x = load_q ? input : k_out; + store_as(output_x, input_vec_x, lane_id); + const auto output_y = pointer::offset(output_x, (kRopeDim / 2) * sizeof(DType)); + store_as(output_y, input_vec_y, lane_id); + } else { + using CacheStorage = AlignedVector; + auto input_vec = load_as(input, lane_id); + const auto cos_vec = load_as(cos_ptr, lane_id); + const auto sin_vec = load_as(sin_ptr, lane_id); +#pragma unroll + for (int64_t j = 0; j < kVecSize; ++j) { + const auto [x, y] = cast(input_vec[j]); + const auto cos = cos_vec[j]; + const auto sin = sin_vec[j]; + const auto out_x = x * cos - y * sin; + const auto out_y = x * sin + y * cos; + input_vec[j] = cast({out_x, out_y}); + } + const auto k_out = pointer::offset(k_cache, loc * cache_stride_bytes, head_id * head_stride_bytes); + const auto output = load_q ? input : k_out; + store_as(output, input_vec, lane_id); + } + } + + __syncwarp(); // to avoid warp divergence + idx -= num_works; + for (; idx < num_extra_works; idx += num_workers) { + using VStorage = AlignedVector; + const int64_t token_id = idx / num_kv_heads; + const int64_t head_id = idx % num_kv_heads; + const auto loc = static_cast(out_loc)[token_id]; + const auto input = pointer::offset(v_ptr, token_id * v_stride_bytes, head_id * head_stride_bytes); + const auto input_vec = load_as(input, lane_id); + const auto output = pointer::offset(v_cache, loc * cache_stride_bytes, head_id * head_stride_bytes); + store_as(output, input_vec, lane_id); + } + PDLTriggerSecondary(); +} + +template +struct FusedRopeKernel { + static constexpr uint32_t kDimPerThread = std::gcd(16 / sizeof(DType), kRopeDim); + static constexpr uint32_t kWorkThreads = next_pow2(kRopeDim, kDimPerThread); + static constexpr bool kSupportFused = kWorkThreads * kDimPerThread == kRopeDim; + static_assert(kRopeDim % kDimPerThread == 0); + static_assert(kBlockSize % kWorkThreads == 0); + + template + static constexpr auto _kernel_0 = fused_rope_kernel; + template + static constexpr auto _kernel_1 = fused_rope_store_kernel; + + static auto get_num_sm(DLDevice device) { + static const auto kNumSM = host::runtime::get_sm_count(device.device_id); + return kNumSM; + } + static void - run(tvm::ffi::TensorView q, // [nnz, H_Q, D] - tvm::ffi::TensorView k, // [nnz, H_K, D] - tvm::ffi::TensorView q_rope, - tvm::ffi::TensorView k_rope, - tvm::ffi::TensorView cos_sin_cache, // [max_seq_len, R] - tvm::ffi::TensorView pos_ids, - bool interleave, - bool enable_pdl, - tvm::ffi::Optional v, // null or [nnz, H_V, D] - tvm::ffi::Optional k_buffer, // null or [nnz, H_K, D] - tvm::ffi::Optional v_buffer, // null or [nnz, H_V, D] - tvm::ffi::Optional kv_cache_loc // null or [n] - ) { + run(const tvm::ffi::TensorView q, + const tvm::ffi::TensorView k, + const tvm::ffi::TensorView cos_sin_cache, + const tvm::ffi::TensorView positions) { + using namespace host; + auto N = SymbolicSize{"num_tokens"}; + auto Q = SymbolicSize{"num_qo_heads"}; + auto K = SymbolicSize{"num_kv_heads"}; + auto D = SymbolicSize{"rope_dim"}; + auto Dq = SymbolicSize{"q_stride"}; + auto Dk = SymbolicSize{"k_stride"}; + auto Dd = SymbolicSize{"head_stride"}; + auto device = SymbolicDevice{}; + auto id_type = SymbolicDType{}; + D.set_value(kRopeDim); + device.set_options(); + TensorMatcher({N, Q, D}) // q input + .with_strides({Dq, Dd, 1}) + .with_dtype() + .with_device(device) + .verify(q); + TensorMatcher({N, K, D}) // k input + .with_strides({Dk, Dd, 1}) + .with_dtype() + .with_device(device) + .verify(k); + TensorMatcher({-1, D}) // cos_sin_cache + .with_dtype() + .with_device(device) + .verify(cos_sin_cache); + TensorMatcher({N}) // positions + .with_dtype(id_type) + .with_device(device) + .verify(positions); + + const auto num_tokens = static_cast(N.unwrap()); + const auto num_qo_heads = static_cast(Q.unwrap()); + const auto num_kv_heads = static_cast(K.unwrap()); + const auto q_stride_bytes = static_cast(Dq.unwrap() * sizeof(DType)); + const auto k_stride_bytes = static_cast(Dk.unwrap() * sizeof(DType)); + const auto head_stride_bytes = static_cast(Dd.unwrap() * sizeof(DType)); + + // NOTE: we offset the k here to reduce computation cost in the kernel + const int64_t k_offset = static_cast(num_qo_heads) * head_stride_bytes; + const auto params = FusedRopeParams{ + .q_ptr = q.data_ptr(), + .k_ptr = pointer::offset(k.data_ptr(), -k_offset), + .cos_sin_cache_ptr = cos_sin_cache.data_ptr(), + .positions = positions.data_ptr(), + .q_stride_bytes = q_stride_bytes, + .k_stride_bytes = k_stride_bytes, + .head_stride_bytes = head_stride_bytes, + .num_qo_heads = num_qo_heads, + .num_kv_heads = num_kv_heads, + .num_tokens = num_tokens, + }; + + const auto is_int32 = id_type.is_type(); + const auto kernel = is_int32 ? _kernel_0 : _kernel_0; + const uint32_t kNumSM = get_num_sm(device.unwrap()); + static const uint32_t kOccupancyTable[2] = { + runtime::get_blocks_per_sm(_kernel_0, kBlockSize), + runtime::get_blocks_per_sm(_kernel_0, kBlockSize), + }; + const auto max_blocks = kOccupancyTable[is_int32 ? 0 : 1] * kNumSM; + const auto num_works = (num_qo_heads + num_kv_heads) * num_tokens; + const auto needed_blocks = div_ceil(num_works, (kBlockSize / kWorkThreads)); + const auto num_blocks = std::min(max_blocks, needed_blocks); + LaunchKernel(num_blocks, kBlockSize, device.unwrap()) // + .enable_pdl(kUsePDL)(kernel, params); + } + + static void run_fused( + const tvm::ffi::TensorView q, + const tvm::ffi::TensorView k, + const tvm::ffi::TensorView v, + const tvm::ffi::TensorView k_cache, + const tvm::ffi::TensorView v_cache, + const tvm::ffi::TensorView cos_sin_cache, + const tvm::ffi::TensorView positions, + const tvm::ffi::TensorView out_loc) { + if constexpr (kSupportFused) { + return _run_fused_impl(q, k, v, k_cache, v_cache, cos_sin_cache, positions, out_loc); + } else { + host::Panic("Fused rope + store is not supported for rope_dim ", kRopeDim); + } + } + + static void _run_fused_impl( + const tvm::ffi::TensorView q, + const tvm::ffi::TensorView k, + const tvm::ffi::TensorView v, + const tvm::ffi::TensorView k_cache, + const tvm::ffi::TensorView v_cache, + const tvm::ffi::TensorView cos_sin_cache, + const tvm::ffi::TensorView positions, + const tvm::ffi::TensorView out_loc) { using namespace host; - RuntimeCheck(q.strides().back() == 1); - RuntimeCheck(k.strides().back() == 1); + auto N = SymbolicSize{"num_tokens"}; + auto Q = SymbolicSize{"num_qo_heads"}; + auto K = SymbolicSize{"num_kv_heads"}; + auto D = SymbolicSize{"rope_dim"}; + auto R = SymbolicSize{"row_size"}; + auto Dq = SymbolicSize{"q_stride"}; + auto Dk = SymbolicSize{"k_stride"}; + auto Dv = SymbolicSize{"v_stride"}; + auto Dd = SymbolicSize{"head_stride"}; + auto Dc = SymbolicSize{"cache_stride"}; + auto device = SymbolicDevice{}; + auto id_type = SymbolicDType{}; + D.set_value(kRopeDim); + device.set_options(); - const bool save_kv_cache = v.has_value(); - if (save_kv_cache) { - RuntimeCheck(v.has_value()); - RuntimeCheck(k_buffer.has_value()); - RuntimeCheck(v_buffer.has_value()); - RuntimeCheck(kv_cache_loc.has_value()); - // CHECK_LAST_DIM_CONTIGUOUS - RuntimeCheck(v.value().strides().back() == 1); - RuntimeCheck(k_buffer.value().strides().back() == 1); - RuntimeCheck(v_buffer.value().strides().back() == 1); - // CHECK_DIM - RuntimeCheck(k_buffer.value().ndim() == 3); - RuntimeCheck(v_buffer.value().ndim() == 3); - RuntimeCheck(v.value().ndim() == 3); - RuntimeCheck(kv_cache_loc.value().ndim() == 1); + TensorMatcher({N, Q, D}) // q input + .with_strides({Dq, Dd, 1}) + .with_dtype() + .with_device(device) + .verify(q); + TensorMatcher({N, K, D}) // k input + .with_strides({Dk, Dd, 1}) + .with_dtype() + .with_device(device) + .verify(k); + TensorMatcher({N, K, D}) // v input + .with_strides({Dv, Dd, 1}) + .with_dtype() + .with_device(device) + .verify(v); + TensorMatcher({-1, D}) // cos_sin_cache + .with_dtype() + .with_device(device) + .verify(cos_sin_cache); + TensorMatcher({N}) // positions, out_loc + .with_dtype(id_type) + .with_device(device) + .verify(positions) + .verify(out_loc); + TensorMatcher({-1, R}) // k_cache + .with_strides({Dc, 1}) + .with_dtype() + .with_device(device) + .verify(k_cache) + .verify(v_cache); - check_cuda_contiguous(kv_cache_loc.value()); - } + const auto num_tokens = static_cast(N.unwrap()); + const auto num_qo_heads = static_cast(Q.unwrap()); + const auto num_kv_heads = static_cast(K.unwrap()); + const auto q_stride_bytes = static_cast(Dq.unwrap() * sizeof(DType)); + const auto k_stride_bytes = static_cast(Dk.unwrap() * sizeof(DType)); + const auto head_stride = Dd.unwrap(); + const auto row_dim = R.unwrap(); + const auto head_stride_bytes = static_cast(Dd.unwrap() * sizeof(DType)); - size_t k_buffer_stride_n = save_kv_cache ? k_buffer.value().stride(0) : 0; - size_t k_buffer_stride_h = save_kv_cache ? k_buffer.value().stride(1) : 0; - size_t v_buffer_stride_n = save_kv_cache ? v_buffer.value().stride(0) : 0; - size_t v_buffer_stride_h = save_kv_cache ? v_buffer.value().stride(1) : 0; - size_t v_stride_n = save_kv_cache ? v.value().stride(0) : 0; - size_t v_stride_h = save_kv_cache ? v.value().stride(1) : 0; - auto kv_cache_loc_ptr = save_kv_cache ? static_cast(kv_cache_loc.value().data_ptr()) : nullptr; + RuntimeCheck(kRopeDim == head_stride, "rope_dim ", kRopeDim, " should = head_stride ", head_stride); + RuntimeCheck(num_kv_heads * kRopeDim == row_dim, "invalid kvcache"); - check_cuda_contiguous(cos_sin_cache); - check_cuda_contiguous(pos_ids); + // NOTE: we offset the k here to reduce computation cost in the kernel + const int64_t k_offset = static_cast(num_qo_heads) * head_stride_bytes; + const auto params = FusedRopeParams{ + .q_ptr = q.data_ptr(), + .k_ptr = pointer::offset(k.data_ptr(), -k_offset), + .cos_sin_cache_ptr = cos_sin_cache.data_ptr(), + .positions = positions.data_ptr(), + .q_stride_bytes = q_stride_bytes, + .k_stride_bytes = k_stride_bytes, + .head_stride_bytes = head_stride_bytes, + .num_qo_heads = num_qo_heads, + .num_kv_heads = num_kv_heads, + .num_tokens = num_tokens, + }; - auto device = q.device(); - RuntimeCheck(k.device() == device); - RuntimeCheck(cos_sin_cache.device() == device); - RuntimeCheck(pos_ids.device() == device); - RuntimeCheck(q.ndim() == 3); - RuntimeCheck(k.ndim() == 3); + const auto v_stride_bytes = static_cast(Dv.unwrap() * sizeof(DType)); + const auto cache_stride_bytes = static_cast(Dc.unwrap() * sizeof(DType)); + const auto store_params = FusedRopeStoreParams{ + .base_params = params, + .v_ptr = v.data_ptr(), + .k_cache = pointer::offset(k_cache.data_ptr(), -k_offset), + .v_cache = v_cache.data_ptr(), + .out_loc = out_loc.data_ptr(), + .v_stride_bytes = v_stride_bytes, + .cache_stride_bytes = cache_stride_bytes, + }; - // cos_sin_cache: (max_seq_len, R) - // First half of R is cos, second half is sin - RuntimeCheck(cos_sin_cache.ndim() == 2); - RuntimeCheck(q.size(0) == k.size(0)); - RuntimeCheck(q.size(2) == k.size(2)); - - unsigned int rotary_dim = cos_sin_cache.size(1); - unsigned int num_qo_heads = q.size(1); - unsigned int num_kv_heads = k.size(1); - unsigned int head_dim = q.size(2); - unsigned int nnz = q.size(0); - size_t q_stride_n = q.stride(0); - size_t q_stride_h = q.stride(1); - size_t k_stride_n = k.stride(0); - size_t k_stride_h = k.stride(1); - - size_t q_rope_stride_n = q_rope.stride(0); - size_t q_rope_stride_h = q_rope.stride(1); - size_t k_rope_stride_n = k_rope.stride(0); - size_t k_rope_stride_h = k_rope.stride(1); - - auto query_dtype = q.dtype(); - const cudaStream_t stream = LaunchKernel::resolve_device(device); - DISPATCH_TVM_DTYPE_TO_CTYPE(query_dtype.code, query_dtype.bits, c_type, [&] { - // TODO temporarily only use `BatchQKApplyRotaryPosIdsCosSinCacheEnhanced` when save_kv_cache - // to avoid changing original code path; but this branch is feature-complete and should switch to this later - if (save_kv_cache) { - cudaError_t status = flashinfer::BatchQKApplyRotaryPosIdsCosSinCacheEnhanced( - static_cast(q.data_ptr()), - static_cast(k.data_ptr()), - save_kv_cache ? static_cast(v.value().data_ptr()) : nullptr, - static_cast(q_rope.data_ptr()), - static_cast(k_rope.data_ptr()), - save_kv_cache ? static_cast(k_buffer.value().data_ptr()) : nullptr, - save_kv_cache ? static_cast(v_buffer.value().data_ptr()) : nullptr, - static_cast(cos_sin_cache.data_ptr()), - static_cast(pos_ids.data_ptr()), - nnz, - num_qo_heads, - num_kv_heads, - rotary_dim, - head_dim, - q_stride_n, - q_stride_h, - k_stride_n, - k_stride_h, - v_stride_n, - v_stride_h, - q_rope_stride_n, - q_rope_stride_h, - k_rope_stride_n, - k_rope_stride_h, - k_buffer_stride_n, - k_buffer_stride_h, - v_buffer_stride_n, - v_buffer_stride_h, - kv_cache_loc_ptr, - interleave, - save_kv_cache, - enable_pdl, - stream); - RuntimeCheck( - status == cudaSuccess, - "BatchQKApplyRotaryPosIdsCosSinCacheEnhanced failed with error code " + - std::string(cudaGetErrorString(status))); - } else { - RuntimeCheck(!enable_pdl); - cudaError_t status = flashinfer::BatchQKApplyRotaryPosIdsCosSinCache( - static_cast(q.data_ptr()), - static_cast(k.data_ptr()), - static_cast(q_rope.data_ptr()), - static_cast(k_rope.data_ptr()), - static_cast(cos_sin_cache.data_ptr()), - static_cast(pos_ids.data_ptr()), - nnz, - num_qo_heads, - num_kv_heads, - rotary_dim, - head_dim, - q_stride_n, - q_stride_h, - k_stride_n, - k_stride_h, - q_rope_stride_n, - q_rope_stride_h, - k_rope_stride_n, - k_rope_stride_h, - interleave, - stream); - RuntimeCheck( - status == cudaSuccess, - "BatchQKApplyRotaryPosIdsCosSinCache failed with error code " + std::string(cudaGetErrorString(status))); - } - return true; - }); + const auto is_int32 = id_type.is_type(); + const auto kernel = is_int32 ? _kernel_1 : _kernel_1; + const uint32_t kNumSM = get_num_sm(device.unwrap()); + static const uint32_t kOccupancyTable[2] = { + runtime::get_blocks_per_sm(_kernel_1, kBlockSize), + runtime::get_blocks_per_sm(_kernel_1, kBlockSize), + }; + const auto max_blocks = kOccupancyTable[is_int32 ? 0 : 1] * kNumSM; + // rope works for q+k heads, plus v store works for kv heads + const auto num_total_works = (num_qo_heads + 2 * num_kv_heads) * num_tokens; + const auto needed_blocks = div_ceil(num_total_works, (kBlockSize / kWorkThreads)); + const auto num_blocks = std::min(max_blocks, needed_blocks); + LaunchKernel(num_blocks, kBlockSize, device.unwrap()) // + .enable_pdl(kUsePDL)(kernel, store_params); } }; diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh index 56e836d19..56e38da76 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh @@ -11,6 +11,7 @@ #include #include #include +#include #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 +SGL_DEVICE T load_as(const void* ptr, int64_t offset = 0) { + return static_cast(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 +SGL_DEVICE void store_as(void* ptr, std::type_identity_t val, int64_t offset = 0) { + static_cast(ptr)[offset] = val; +} + namespace pointer { // we only allow void * pointer arithmetic for safety diff --git a/python/sglang/jit_kernel/rope.py b/python/sglang/jit_kernel/rope.py index 69115fb21..02c4e5fb6 100644 --- a/python/sglang/jit_kernel/rope.py +++ b/python/sglang/jit_kernel/rope.py @@ -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 ) diff --git a/python/sglang/jit_kernel/tests/test_rope.py b/python/sglang/jit_kernel/tests/test_rope.py index fd34b2765..99c271966 100644 --- a/python/sglang/jit_kernel/tests/test_rope.py +++ b/python/sglang/jit_kernel/tests/test_rope.py @@ -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"]) diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index ae0614635..11cbbfcdd 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -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 ( diff --git a/python/sglang/srt/models/gpt_oss.py b/python/sglang/srt/models/gpt_oss.py index 2cf813bce..96caaa65b 100644 --- a/python/sglang/srt/models/gpt_oss.py +++ b/python/sglang/srt/models/gpt_oss.py @@ -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" diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index cf0095551..5f5cc438c 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -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 diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index a74218462..ca3f6c09d 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -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, )