diff --git a/python/sglang/jit_kernel/csrc/elementwise/rope.cuh b/python/sglang/jit_kernel/csrc/elementwise/rope.cuh new file mode 100644 index 000000000..c19e9abc7 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/elementwise/rope.cuh @@ -0,0 +1,656 @@ +/* + * 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 // upstream +#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 + +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; \ + }() + +inline void check_cuda_contiguous(tvm::ffi::TensorView x) { + using namespace host; + + RuntimeCheck(x.device().device_type == kDLCUDA); + RuntimeCheck(x.is_contiguous()); +} + +struct ApplyRopePosIdsCosSinCacheKernel { + 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] + ) { + using namespace host; + + RuntimeCheck(q.strides().back() == 1); + RuntimeCheck(k.strides().back() == 1); + + 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); + + check_cuda_contiguous(kv_cache_loc.value()); + } + + 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; + + check_cuda_contiguous(cos_sin_cache); + check_cuda_contiguous(pos_ids); + + 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); + + // 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; + }); + } +}; + +} // namespace diff --git a/python/sglang/jit_kernel/rope.py b/python/sglang/jit_kernel/rope.py new file mode 100644 index 000000000..69115fb21 --- /dev/null +++ b/python/sglang/jit_kernel/rope.py @@ -0,0 +1,236 @@ +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.srt.utils.custom_op import register_custom_op + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +@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() + return load_jit( + "apply_rope_pos_ids_cos_sin_cache", + cuda_files=["elementwise/rope.cuh"], + cuda_wrappers=[ + ( + "apply_rope_pos_ids_cos_sin_cache", + "ApplyRopePosIdsCosSinCacheKernel::run", + ) + ], + 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: + """ + value : Optional[torch.Tensor] + Value tensor, shape: ``(nnz, num_v_heads * head_size)``. + k_buffer : Optional[torch.Tensor] + 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. + """ + + 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, + cos_sin_cache: torch.Tensor, + is_neox: bool = True, + fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None, + enable_pdl: Optional[bool] = None, +) -> 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") + + 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) + + 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), + 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, + ) + 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, + ) diff --git a/python/sglang/jit_kernel/tests/test_rope.py b/python/sglang/jit_kernel/tests/test_rope.py new file mode 100644 index 000000000..fd34b2765 --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_rope.py @@ -0,0 +1,301 @@ +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" + + +@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): + inv_freq = 1.0 / ( + base + ** ( + torch.arange(0, rotary_dim, 2, dtype=torch.float32, device=DEVICE) + / rotary_dim + ) + ) + + t = torch.arange(max_position_embeddings, 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) + 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]) +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, + 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 + ) + 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() + + 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() + 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) + + +@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, +) -> 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 + ) + 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() + + 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, + ) + 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 + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index 311df6d01..6db5f0987 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -37,7 +37,10 @@ _is_xpu = is_xpu() _is_musa = is_musa() if _is_cuda: - from sgl_kernel import FusedSetKVBufferArg, apply_rope_with_cos_sin_cache_inplace + from sglang.jit_kernel.rope import ( + FusedSetKVBufferArg, + apply_rope_with_cos_sin_cache_inplace, + ) else: FusedSetKVBufferArg = None