minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[kernel slimming] Clean many useless sgl-kernel deprecated kernels (#20277)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2026-03-14 16:45:54 +08:00
committed by GitHub
parent 75a7879fd4
commit 25e38216b6
26 changed files with 60 additions and 1483 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/common_extension.cc
-12
View File
@@ -84,12 +84,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
m.def("gelu_and_mul(Tensor! out, Tensor input) -> ()");
m.impl("gelu_and_mul", torch::kCUDA, &gelu_and_mul);
m.def(
"apply_rope_pos_ids_cos_sin_cache(Tensor q, Tensor k, Tensor! q_rope, Tensor! k_rope, Tensor cos_sin_cache, "
"Tensor pos_ids, bool interleave, bool enable_pdl, "
"Tensor? v, Tensor!? k_buffer, Tensor!? v_buffer, Tensor? kv_cache_loc) -> ()");
m.impl("apply_rope_pos_ids_cos_sin_cache", torch::kCUDA, &apply_rope_pos_ids_cos_sin_cache);
m.def(
"rotary_embedding(Tensor positions, Tensor! query,"
" Tensor!? key, int head_size,"
@@ -151,9 +145,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
" float eps, float fp8_min, float fp8_max, bool scale_ue8m0, bool fuse_silu_and_mul, Tensor? masked_m) -> ()");
m.impl("sgl_per_token_group_quant_8bit_v2", torch::kCUDA, &sgl_per_token_group_quant_8bit_v2);
m.def("sgl_per_tensor_quant_fp8(Tensor input, Tensor! output_q, Tensor! output_s, bool is_static) -> ()");
m.impl("sgl_per_tensor_quant_fp8", torch::kCUDA, &sgl_per_tensor_quant_fp8);
m.def("sgl_per_token_quant_fp8(Tensor input, Tensor! output_q, Tensor! output_s) -> ()");
m.impl("sgl_per_token_quant_fp8", torch::kCUDA, &sgl_per_token_quant_fp8);
@@ -355,9 +346,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
/*
* From csrc/memory
*/
m.def("store_kv_cache(Tensor k_cache, Tensor v_cache, Tensor out_loc, Tensor k, Tensor v) -> ()");
m.impl("store_kv_cache", &store_kv_cache);
m.def("weak_ref_tensor(Tensor tensor) -> Tensor");
m.impl("weak_ref_tensor", torch::kCUDA, &weak_ref_tensor);
sgl-kernel/csrc/elementwise/rope.cu
-168
View File
@@ -1,168 +0,0 @@
/*
* 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 <ATen/cuda/Exceptions.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <torch/all.h>
#include "pos_enc.cuh"
#include "utils.h"
using namespace flashinfer;
void apply_rope_pos_ids_cos_sin_cache(
at::Tensor q,
at::Tensor k,
at::Tensor q_rope,
at::Tensor k_rope,
at::Tensor cos_sin_cache,
at::Tensor pos_ids,
bool interleave,
bool enable_pdl,
const std::optional<at::Tensor>& v,
const std::optional<at::Tensor>& k_buffer,
const std::optional<at::Tensor>& v_buffer,
const std::optional<at::Tensor>& kv_cache_loc) {
CHECK_LAST_DIM_CONTIGUOUS(q);
CHECK_LAST_DIM_CONTIGUOUS(k);
const bool save_kv_cache = v.has_value();
if (save_kv_cache) {
TORCH_CHECK(v.has_value());
TORCH_CHECK(k_buffer.has_value());
TORCH_CHECK(v_buffer.has_value());
TORCH_CHECK(kv_cache_loc.has_value());
CHECK_LAST_DIM_CONTIGUOUS(v.value());
CHECK_LAST_DIM_CONTIGUOUS(k_buffer.value());
CHECK_LAST_DIM_CONTIGUOUS(v_buffer.value());
CHECK_DIM(3, k_buffer.value()); // k_buffer: (nnz, H_K, D)
CHECK_DIM(3, v_buffer.value()); // v_buffer: (nnz, H_V, D)
CHECK_DIM(3, v.value()); // v: (nnz, H_V, D)
CHECK_DIM(1, kv_cache_loc.value()); // v: (n)
CHECK_INPUT(kv_cache_loc.value());
}
size_t k_buffer_stride_n = save_kv_cache ? k_buffer->stride(0) : 0;
size_t k_buffer_stride_h = save_kv_cache ? k_buffer->stride(1) : 0;
size_t v_buffer_stride_n = save_kv_cache ? v_buffer->stride(0) : 0;
size_t v_buffer_stride_h = save_kv_cache ? v_buffer->stride(1) : 0;
size_t v_stride_n = save_kv_cache ? v->stride(0) : 0;
size_t v_stride_h = save_kv_cache ? v->stride(1) : 0;
auto kv_cache_loc_ptr = save_kv_cache ? static_cast<int64_t*>(kv_cache_loc->data_ptr()) : nullptr;
CHECK_INPUT(cos_sin_cache);
CHECK_INPUT(pos_ids);
auto device = q.device();
CHECK_EQ(k.device(), device);
CHECK_EQ(cos_sin_cache.device(), device);
CHECK_EQ(pos_ids.device(), device);
CHECK_DIM(3, q); // q: (nnz, H_Q, D)
CHECK_DIM(3, k); // k: (nnz, H_K, D)
// cos_sin_cache: (max_seq_len, R)
// First half of R is cos, second half is sin
CHECK_DIM(2, cos_sin_cache);
CHECK_EQ(q.size(0), k.size(0));
CHECK_EQ(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);
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(q.scalar_type(), 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 = BatchQKApplyRotaryPosIdsCosSinCacheEnhanced(
static_cast<c_type*>(q.data_ptr()),
static_cast<c_type*>(k.data_ptr()),
save_kv_cache ? static_cast<c_type*>(v->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->data_ptr()) : nullptr,
save_kv_cache ? static_cast<c_type*>(v_buffer->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);
TORCH_CHECK(
status == cudaSuccess,
"BatchQKApplyRotaryPosIdsCosSinCacheEnhanced failed with error code " +
std::string(cudaGetErrorString(status)));
} else {
TORCH_CHECK(!enable_pdl);
cudaError_t status = 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);
TORCH_CHECK(
status == cudaSuccess,
"BatchQKApplyRotaryPosIdsCosSinCache failed with error code " + std::string(cudaGetErrorString(status)));
}
return true;
});
}
sgl-kernel/csrc/gemm/per_tensor_quant_fp8.cu
-123
View File
@@ -1,123 +0,0 @@
#include <ATen/cuda/CUDAContext.h>
#include <c10/util/Float8_e4m3fn.h>
#include <cmath>
#include <cub/block/block_reduce.cuh>
#include <flashinfer/vec_dtypes.cuh>
#include "utils.h"
template <typename T>
__global__ void
per_tensor_absmax_kernel(const T* __restrict__ input, float* __restrict__ output_s, const int64_t num_elements) {
float max_value = 0.0f;
unsigned int tid = threadIdx.x;
unsigned int gid = blockIdx.x * blockDim.x + threadIdx.x;
const int grid_size = blockDim.x * gridDim.x;
constexpr uint32_t vec_size = 16 / sizeof(T);
using vec_t = flashinfer::vec_t<T, vec_size>;
const int32_t num_vec_elems = num_elements / vec_size;
for (int32_t i = gid; i < num_vec_elems; i += grid_size) {
vec_t input_vec;
input_vec.cast_load(input + i * vec_size);
#pragma unroll
for (uint32_t j = 0; j < vec_size; ++j) {
float val = static_cast<float>(input_vec[j]);
max_value = fmaxf(max_value, fabsf(val));
}
}
const int32_t remaining_start = num_vec_elems * vec_size;
for (int32_t idx = remaining_start + gid; idx < num_elements; idx += grid_size) {
float val = static_cast<float>(input[idx]);
max_value = fmaxf(max_value, fabsf(val));
}
max_value = blockReduceMax(max_value);
if (tid == 0) {
atomicMaxFloat(output_s, max_value / FP8_E4M3_MAX);
}
}
template <typename T, typename DST_DTYPE>
__global__ void per_tensor_quant_fp8_kernel(
const T* __restrict__ input,
DST_DTYPE* __restrict__ output,
const float* __restrict__ scale,
const int64_t num_elements) {
const int gid = blockIdx.x * blockDim.x + threadIdx.x;
const int grid_size = blockDim.x * gridDim.x;
const float scale_val = 1.0f / (*scale);
// We want to store 128 bits of data at a time. 16 = 128 / 8 bits
// Load is already vectorized, so 16 elements work for T.
const uint32_t VEC_SIZE = 16;
using vec_t = flashinfer::vec_t<T, VEC_SIZE>;
const int32_t num_vec_elems = num_elements / VEC_SIZE;
for (int32_t i = gid; i < num_vec_elems; i += grid_size) {
vec_t input_vec;
input_vec.cast_load(input + i * VEC_SIZE);
DST_DTYPE output_arr[VEC_SIZE];
#pragma unroll
for (uint32_t j = 0; j < VEC_SIZE; ++j) {
float val = fmax(fmin(static_cast<float>(input_vec[j]) * scale_val, FP8_E4M3_MAX), -FP8_E4M3_MAX);
#if !defined(USE_ROCM) || defined(HIP_FP8_TYPE_E4M3)
output_arr[j] = static_cast<DST_DTYPE>(val);
#else
output_arr[j] = c10::Float8_e4m3fnuz(
__hip_cvt_float_to_fp8(val, fp8::fp8_type::__default_saturation, fp8::fp8_type::__default_interpret),
c10::Float8_e4m3fnuz::from_bits());
#endif
}
*(uint4*)(output + i * VEC_SIZE) = *(uint4*)output_arr;
}
const int32_t remaining_start = num_vec_elems * VEC_SIZE;
for (int32_t idx = remaining_start + gid; idx < num_elements; idx += grid_size) {
float val = fmax(-FP8_E4M3_MAX, fmin(static_cast<float>(input[idx]) * scale_val, FP8_E4M3_MAX));
#if !defined(USE_ROCM) || defined(HIP_FP8_TYPE_E4M3)
output[idx] = static_cast<DST_DTYPE>(val);
#else
output[idx] = c10::Float8_e4m3fnuz(
__hip_cvt_float_to_fp8(val, fp8::fp8_type::__default_saturation, fp8::fp8_type::__default_interpret),
c10::Float8_e4m3fnuz::from_bits());
#endif
}
}
void sgl_per_tensor_quant_fp8(torch::Tensor input, torch::Tensor output_q, torch::Tensor output_s, bool is_static) {
CHECK_INPUT(input);
CHECK_INPUT(output_q);
CHECK_INPUT(output_s);
const int block_size = 256;
const int num_elements = input.numel();
const int num_blocks = min((num_elements + block_size - 1) / block_size, 1024);
dim3 grid(num_blocks);
dim3 block(block_size);
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(input.scalar_type(), scalar_t, [&] {
if (is_static == false) {
per_tensor_absmax_kernel<scalar_t><<<grid, block, 0, stream>>>(
static_cast<scalar_t*>(input.data_ptr()), static_cast<float*>(output_s.data_ptr()), num_elements);
}
per_tensor_quant_fp8_kernel<scalar_t, __nv_fp8_e4m3><<<grid, block, 0, stream>>>(
static_cast<scalar_t*>(input.data_ptr()),
static_cast<__nv_fp8_e4m3*>(output_q.data_ptr()),
static_cast<float*>(output_s.data_ptr()),
num_elements);
return true;
});
}
sgl-kernel/csrc/memory/store.cu
-147
View File
@@ -1,147 +0,0 @@
#include <ATen/Dispatch.h>
#include <ATen/core/TensorBody.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/util/Exception.h>
#include <cstddef>
#include <cstdint>
#include <type_traits>
namespace {
using std::size_t;
using std::uint64_t;
// Each warp will process 256 bytes per loop iteration
template <typename T>
__global__ void store_kv_cache_256x1(
uint64_t* __restrict__ k_cache,
uint64_t* __restrict__ v_cache,
const T* __restrict__ out_loc,
const size_t length,
const uint64_t* __restrict__ k,
const uint64_t* __restrict__ v,
const size_t kv_cache_stride,
const size_t kv_input_stride,
const size_t num_items) {
const auto idx = blockIdx.x * blockDim.x + threadIdx.x;
const auto warp_id = idx / 32;
const auto lane_id = idx % 32;
if (warp_id >= length) return;
const auto offset = out_loc[warp_id];
const auto k_dst = k_cache + offset * kv_cache_stride;
const auto v_dst = v_cache + offset * kv_cache_stride;
const auto k_src = k + warp_id * kv_input_stride;
const auto v_src = v + warp_id * kv_input_stride;
for (size_t i = 0; i < num_items; ++i) {
k_dst[lane_id + i * 32] = k_src[lane_id + i * 32];
v_dst[lane_id + i * 32] = v_src[lane_id + i * 32];
}
}
// Each warp will process 128 bytes per loop iteration
template <typename T>
__global__ void store_kv_cache_128x2(
uint64_t* __restrict__ k_cache,
uint64_t* __restrict__ v_cache,
const T* __restrict__ out_loc,
const size_t length,
const uint64_t* __restrict__ k,
const uint64_t* __restrict__ v,
const size_t kv_cache_stride,
const size_t kv_input_stride,
const size_t num_items) {
const auto idx = blockIdx.x * blockDim.x + threadIdx.x;
const auto warp_id = idx / 32;
const auto lane_id = idx % 32;
if (warp_id >= length) return;
const auto offset = out_loc[warp_id];
const auto copy_k = lane_id < 16;
const auto copy_id = lane_id % 16;
const auto cache = copy_k ? k_cache : v_cache;
const auto input = copy_k ? k : v;
const auto dst = cache + offset * kv_cache_stride;
const auto src = input + warp_id * kv_input_stride;
for (size_t i = 0; i < num_items; ++i) {
dst[copy_id + i * 16] = src[copy_id + i * 16];
}
}
} // namespace
auto store_kv_cache(at::Tensor k_cache, at::Tensor v_cache, at::Tensor out_loc, at::Tensor k, at::Tensor v) -> void {
const auto max_tokens = k_cache.size(0);
const auto num_tokens = out_loc.size(0);
k_cache = k_cache.view({max_tokens, -1});
v_cache = v_cache.view({max_tokens, -1});
k = k.view({num_tokens, -1});
v = v.view({num_tokens, -1});
TORCH_CHECK(
k_cache.is_cuda() && v_cache.is_cuda() && out_loc.is_cuda() && k.is_cuda() && v.is_cuda(),
"All tensors must be CUDA tensors");
TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must have the same size");
TORCH_CHECK(k_cache.strides() == v_cache.strides(), "k_cache and v_cache must have the same strides");
TORCH_CHECK(k.sizes() == v.sizes(), "k and v must have the same size");
TORCH_CHECK(k.strides() == v.strides(), "k and v must have the same strides");
TORCH_CHECK(k.stride(-1) == 1 && k_cache.stride(-1) == 1, "k and k_cache must be contiguous in head.");
TORCH_CHECK(k.size(-1) == k_cache.size(-1), "k and k_cache must have the same head size");
TORCH_CHECK(out_loc.dim() == 1 && out_loc.is_contiguous(), "out_loc must be a 1D contiguous tensor");
static_assert(sizeof(uint64_t) == 8, "uint64_t must be 8 bytes, our code assumes that");
const auto length = out_loc.size(0);
const auto elem_size = k.element_size();
const auto size_bytes = elem_size * k.size(-1);
const auto kv_cache_stride_bytes = elem_size * k_cache.stride(-2);
const auto kv_input_stride_bytes = elem_size * k.stride(-2);
const auto kv_cache_stride = kv_cache_stride_bytes / 8;
const auto kv_input_stride = kv_input_stride_bytes / 8;
const auto k_cache_ptr = static_cast<uint64_t*>(k_cache.data_ptr());
const auto v_cache_ptr = static_cast<uint64_t*>(v_cache.data_ptr());
const auto k_ptr = static_cast<const uint64_t*>(k.data_ptr());
const auto v_ptr = static_cast<const uint64_t*>(v.data_ptr());
const auto num_threads = 256;
const auto num_warps = num_threads / 32;
const auto num_blocks = (length + num_warps - 1) / num_warps;
const auto stream = at::cuda::getCurrentCUDAStream();
AT_DISPATCH_INTEGRAL_TYPES(out_loc.scalar_type(), "store_kv_cache", [&] {
if constexpr (!std::is_same_v<scalar_t, int32_t> && !std::is_same_v<scalar_t, int64_t>) {
// do not instantiate the kernel if out_loc is not int32 or int64
TORCH_CHECK(false, "out_loc must be of type int32 or int64, got: ", out_loc.scalar_type());
} else {
if (size_bytes % 256 == 0) {
const auto items_per_warp = size_bytes / 256;
store_kv_cache_256x1<<<num_blocks, num_threads, 0, stream>>>(
k_cache_ptr,
v_cache_ptr,
out_loc.data_ptr<scalar_t>(),
length,
k_ptr,
v_ptr,
kv_cache_stride,
kv_input_stride,
items_per_warp);
} else if (size_bytes % 128 == 0) {
const auto items_per_warp = size_bytes / 128;
store_kv_cache_128x2<<<num_blocks, num_threads, 0, stream>>>(
k_cache_ptr,
v_cache_ptr,
out_loc.data_ptr<scalar_t>(),
length,
k_ptr,
v_ptr,
kv_cache_stride,
kv_input_stride,
items_per_warp);
} else {
TORCH_CHECK(
false,
"The last dimension size bytes of k and v must be"
" divisible by 128 at least, got: ",
size_bytes);
}
}
});
}