[Kernel] Add JIT apply_rope_with_cos_sin_cache_inplace (#18155)
This commit is contained in:
pansicheng
656
python/sglang/jit_kernel/csrc/elementwise/rope.cuh
Normal file
656
python/sglang/jit_kernel/csrc/elementwise/rope.cuh
Normal file
@@ -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 <sgl_kernel/tensor.h>
|
||||
#include <sgl_kernel/utils.h>
|
||||
|
||||
#include <sgl_kernel/utils.cuh>
|
||||
|
||||
#include <flashinfer/pos_enc.cuh> // upstream
|
||||
#include <tvm/ffi/container/tensor.h>
|
||||
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
namespace flashinfer {
|
||||
|
||||
namespace kv_buffer_saver {
|
||||
|
||||
template <typename DType, typename IdType, uint32_t vec_size>
|
||||
__device__ __forceinline__ void prepare(
|
||||
vec_t<float, vec_size>& 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 <typename DType, typename IdType, uint32_t vec_size>
|
||||
__device__ __forceinline__ void save(
|
||||
IdType& kv_cache_offset,
|
||||
vec_t<float, vec_size>& k_vec,
|
||||
vec_t<float, vec_size>& 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<float, vec_size> 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<float, vec_size> q_vec;
|
||||
if constexpr (interleave) {
|
||||
q_vec = vec_apply_llama_rope_cos_sin_interleave_reuse_half<vec_size, bdx>(q_ptr, cos, sin, rotary_dim);
|
||||
} else {
|
||||
q_vec = vec_apply_llama_rope_cos_sin<vec_size, bdx>(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<float, vec_size> v_vec;
|
||||
IdType kv_cache_offset;
|
||||
if constexpr (save_kv_cache) {
|
||||
kv_buffer_saver::prepare<DType, IdType, vec_size>(
|
||||
v_vec, kv_cache_offset, v, kv_cache_loc, idx, tx, kv_head_idx, v_stride_n, v_stride_h);
|
||||
}
|
||||
|
||||
vec_t<float, vec_size> k_vec;
|
||||
if constexpr (interleave) {
|
||||
k_vec = vec_apply_llama_rope_cos_sin_interleave_reuse_half<vec_size, bdx>(k_ptr, cos, sin, rotary_dim);
|
||||
} else {
|
||||
k_vec = vec_apply_llama_rope_cos_sin<vec_size, bdx>(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<DType, IdType, vec_size>(
|
||||
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<float, vec_size> 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<float, vec_size> q_vec;
|
||||
if constexpr (interleave) {
|
||||
q_vec = vec_apply_llama_rope_cos_sin_interleave_reuse_half<vec_size, bdx>(q_ptr, cos, sin, rotary_dim);
|
||||
} else {
|
||||
q_vec = vec_apply_llama_rope_cos_sin<vec_size, bdx>(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<float, vec_size> v_vec;
|
||||
IdType kv_cache_offset;
|
||||
if constexpr (save_kv_cache) {
|
||||
kv_buffer_saver::prepare<DType, IdType, vec_size>(
|
||||
v_vec, kv_cache_offset, v, kv_cache_loc, idx, tx, kv_head_idx, v_stride_n, v_stride_h);
|
||||
}
|
||||
|
||||
vec_t<float, vec_size> k_vec;
|
||||
if constexpr (interleave) {
|
||||
k_vec = vec_apply_llama_rope_cos_sin_interleave_reuse_half<vec_size, bdx>(k_ptr, cos, sin, rotary_dim);
|
||||
} else {
|
||||
k_vec = vec_apply_llama_rope_cos_sin<vec_size, bdx>(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<DType, IdType, vec_size>(
|
||||
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 <typename DType, typename IdType>
|
||||
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<tvm::ffi::TensorView> v, // null or [nnz, H_V, D]
|
||||
tvm::ffi::Optional<tvm::ffi::TensorView> k_buffer, // null or [nnz, H_K, D]
|
||||
tvm::ffi::Optional<tvm::ffi::TensorView> v_buffer, // null or [nnz, H_V, D]
|
||||
tvm::ffi::Optional<tvm::ffi::TensorView> 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<int64_t*>(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<c_type*>(q.data_ptr()),
|
||||
static_cast<c_type*>(k.data_ptr()),
|
||||
save_kv_cache ? static_cast<c_type*>(v.value().data_ptr()) : nullptr,
|
||||
static_cast<c_type*>(q_rope.data_ptr()),
|
||||
static_cast<c_type*>(k_rope.data_ptr()),
|
||||
save_kv_cache ? static_cast<c_type*>(k_buffer.value().data_ptr()) : nullptr,
|
||||
save_kv_cache ? static_cast<c_type*>(v_buffer.value().data_ptr()) : nullptr,
|
||||
static_cast<float*>(cos_sin_cache.data_ptr()),
|
||||
static_cast<int64_t*>(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<c_type*>(q.data_ptr()),
|
||||
static_cast<c_type*>(k.data_ptr()),
|
||||
static_cast<c_type*>(q_rope.data_ptr()),
|
||||
static_cast<c_type*>(k_rope.data_ptr()),
|
||||
static_cast<float*>(cos_sin_cache.data_ptr()),
|
||||
static_cast<int64_t*>(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
|
||||
236
python/sglang/jit_kernel/rope.py
Normal file
236
python/sglang/jit_kernel/rope.py
Normal file
@@ -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,
|
||||
)
|
||||
301
python/sglang/jit_kernel/tests/test_rope.py
Normal file
301
python/sglang/jit_kernel/tests/test_rope.py
Normal file
@@ -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__])
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user