chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[sgl-kernel][Deepseek V3.2] Add row_starts to topk kernel (#12582)

Browse Source 
Signed-off-by: Hao Lu <14827759+hlu1@users.noreply.github.com>
This commit is contained in:
hlu1
2025-11-07 18:33:27 -08:00
committed by GitHub
parent 0b88d520a0
commit b8ddc296f4
5 changed files with 211 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/common_extension.cc
+3 -3
View File
@@ -107,15 +107,15 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
m.def("concat_mla_absorb_q(Tensor a, Tensor b, Tensor! out) -> ()");
m.impl("concat_mla_absorb_q", torch::kCUDA, &concat_mla_absorb_q);
m.def("fast_topk(Tensor score, Tensor indices, Tensor lengths) -> ()");
m.def("fast_topk(Tensor score, Tensor indices, Tensor lengths, Tensor? row_starts) -> ()");
m.impl("fast_topk", torch::kCUDA, &fast_topk_interface);
m.def(
"fast_topk_transform_fused(Tensor score, Tensor lengths, Tensor dst_page_table, Tensor src_page_table, Tensor "
"cu_seqlens_q) -> ()");
"cu_seqlens_q, Tensor? row_starts) -> ()");
m.impl("fast_topk_transform_fused", torch::kCUDA, &fast_topk_transform_interface);
m.def(
"fast_topk_transform_ragged_fused(Tensor score, Tensor lengths, Tensor topk_indices_ragged, Tensor "
"topk_indices_offset) -> ()");
"topk_indices_offset, Tensor ? row_starts) -> ()");
m.impl("fast_topk_transform_ragged_fused", torch::kCUDA, &fast_topk_transform_ragged_interface);
/*
sgl-kernel/csrc/elementwise/topk.cu
+51 -24
View File
@@ -25,9 +25,10 @@ constexpr int kThreadsPerBlock = 1024;
constexpr size_t kSmem = 32 * 1024 * sizeof(uint32_t); // 128KB
struct FastTopKParams {
const float* __restrict__ input; // [B, input_stride]
int32_t* __restrict__ indices; // [B, TopK]
int32_t* __restrict__ lengths; // [B]
const float* __restrict__ input; // [B, input_stride]
const int32_t* __restrict__ row_starts; // [B]
int32_t* __restrict__ indices; // [B, TopK]
int32_t* __restrict__ lengths; // [B]
int64_t input_stride;
};
@@ -72,7 +73,7 @@ __device__ __forceinline__ auto convert_to_uint32(float x) -> uint32_t {
return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
}
__device__ void fast_topk_cuda_tl(const float* __restrict__ input, int* __restrict__ index, int length) {
__device__ void fast_topk_cuda_tl(const float* __restrict__ input, int* __restrict__ index, int row_start, int length) {
// An optimized topk kernel copied from tilelang kernel
// We assume length > TopK here, or it will crash
int topk = TopK;
@@ -96,7 +97,7 @@ __device__ void fast_topk_cuda_tl(const float* __restrict__ input, int* __restri
__syncthreads();
for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
const auto bin = convert_to_uint8(input[idx]);
const auto bin = convert_to_uint8(input[idx + row_start]);
::atomicAdd(&s_histogram[bin], 1);
}
__syncthreads();
@@ -131,7 +132,7 @@ __device__ void fast_topk_cuda_tl(const float* __restrict__ input, int* __restri
if (topk == 0) {
for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
const auto bin = static_cast<int>(convert_to_uint8(input[idx]));
const auto bin = static_cast<int>(convert_to_uint8(input[idx + row_start]));
if (bin > threshold_bin) {
const auto pos = ::atomicAdd(&s_counter, 1);
index[pos] = idx;
@@ -147,7 +148,7 @@ __device__ void fast_topk_cuda_tl(const float* __restrict__ input, int* __restri
__syncthreads();
for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
const auto raw_input = input[idx];
const auto raw_input = input[idx + row_start];
const auto bin = static_cast<int>(convert_to_uint8(raw_input));
if (bin > threshold_bin) {
const auto pos = ::atomicAdd(&s_counter, 1);
@@ -191,7 +192,7 @@ __device__ void fast_topk_cuda_tl(const float* __restrict__ input, int* __restri
for (int i = tx; i < num_input; i += BLOCK_SIZE) {
const auto idx = s_input_idx[r_idx][i];
const auto offset = 24 - round * 8;
const auto bin = (convert_to_uint32(input[idx]) >> offset) & 0xFF;
const auto bin = (convert_to_uint32(input[idx + row_start]) >> offset) & 0xFF;
if (bin > threshold_bin) {
const auto pos = ::atomicAdd(&s_counter, 1);
index[pos] = idx;
@@ -207,7 +208,7 @@ __device__ void fast_topk_cuda_tl(const float* __restrict__ input, int* __restri
__syncthreads();
for (int i = tx; i < num_input; i += BLOCK_SIZE) {
const auto idx = s_input_idx[r_idx][i];
const auto raw_input = input[idx];
const auto raw_input = input[idx + row_start];
const auto offset = 24 - round * 8;
const auto bin = (convert_to_uint32(raw_input) >> offset) & 0xFF;
if (bin > threshold_bin) {
@@ -238,15 +239,16 @@ __device__ void fast_topk_cuda_tl(const float* __restrict__ input, int* __restri
__global__ __launch_bounds__(kThreadsPerBlock) // topk
void topk_kernel(const FastTopKParams params) {
const auto& [input, indices, lengths, input_stride] = params;
const auto& [input, row_starts, indices, lengths, input_stride] = params;
const auto bid = static_cast<uint64_t>(blockIdx.x);
const auto row_start = row_starts == nullptr ? 0 : row_starts[bid];
const auto length = lengths[bid];
const auto indice = indices + bid * TopK;
const auto score = input + bid * input_stride;
if (length <= TopK) {
return naive_topk_cuda(score, indice, length);
} else {
return fast_topk_cuda_tl(score, indice, length);
return fast_topk_cuda_tl(score, indice, row_start, length);
}
}
@@ -256,9 +258,10 @@ __global__ __launch_bounds__(kThreadsPerBlock) // decode
int32_t* __restrict__ dst_page_table,
const int32_t* __restrict__ src_page_table,
const int64_t src_stride) {
const auto& [input, _, lengths, input_stride] = params;
const auto& [input, _1, _2, lengths, input_stride] = params;
const auto bid = static_cast<uint64_t>(blockIdx.x);
const auto tid = threadIdx.x;
const auto row_start = 0;
const auto length = lengths[bid];
const auto src_page_entry = src_page_table + bid * src_stride;
const auto dst_page_entry = dst_page_table + bid * TopK;
@@ -267,7 +270,7 @@ __global__ __launch_bounds__(kThreadsPerBlock) // decode
return naive_topk_transform(score, length, dst_page_entry, src_page_entry);
} else {
__shared__ int s_indices[TopK];
fast_topk_cuda_tl(score, s_indices, length);
fast_topk_cuda_tl(score, s_indices, row_start, length);
// copy src[s_indices] to dst, we manually unroll here
static_assert(TopK % kThreadsPerBlock == 0);
static_assert(TopK / kThreadsPerBlock == 2);
@@ -288,10 +291,11 @@ __global__ __launch_bounds__(kThreadsPerBlock) // prefill
const int64_t src_stride,
const int32_t* __restrict__ cu_seqlens_q,
const int64_t prefill_bs) {
const auto& [input, _, lengths, input_stride] = params;
const auto& [input, row_starts, _, lengths, input_stride] = params;
const auto bid = static_cast<uint64_t>(blockIdx.x);
const auto tid = threadIdx.x;
const auto length = lengths[bid];
const auto row_start = row_starts == nullptr ? 0 : row_starts[bid];
const auto dst_page_entry = dst_page_table + bid * TopK;
const auto score = input + bid * input_stride;
@@ -318,7 +322,7 @@ __global__ __launch_bounds__(kThreadsPerBlock) // prefill
return naive_topk_transform(score, length, dst_page_entry, src_page_entry);
} else {
__shared__ int s_indices[TopK];
fast_topk_cuda_tl(score, s_indices, length);
fast_topk_cuda_tl(score, s_indices, row_start, length);
// copy src[s_indices] to dst, we manually unroll here
static_assert(TopK % kThreadsPerBlock == 0);
static_assert(TopK / kThreadsPerBlock == 2);
@@ -336,9 +340,10 @@ __global__ __launch_bounds__(kThreadsPerBlock) // prefill, ragged kv
const FastTopKParams params,
int32_t* __restrict__ topk_indices_ragged,
const int32_t* __restrict__ topk_indices_offset) {
const auto& [input, _, lengths, input_stride] = params;
const auto& [input, row_starts, _, lengths, input_stride] = params;
const auto bid = static_cast<uint64_t>(blockIdx.x);
const auto tid = threadIdx.x;
const auto row_start = row_starts == nullptr ? 0 : row_starts[bid];
const auto length = lengths[bid];
const auto dst_indices_entry = topk_indices_ragged + bid * TopK;
const auto score = input + bid * input_stride;
@@ -348,7 +353,7 @@ __global__ __launch_bounds__(kThreadsPerBlock) // prefill, ragged kv
return naive_topk_transform_ragged(score, length, dst_indices_entry, offset);
} else {
__shared__ int s_indices[TopK];
fast_topk_cuda_tl(score, s_indices, length);
fast_topk_cuda_tl(score, s_indices, row_start, length);
// copy src[s_indices] to dst, we manually unroll here
static_assert(TopK % kThreadsPerBlock == 0);
static_assert(TopK / kThreadsPerBlock == 2);
@@ -364,9 +369,15 @@ __global__ __launch_bounds__(kThreadsPerBlock) // prefill, ragged kv
auto get_params(
const at::Tensor& score,
const at::Tensor& lengths,
std::optional<at::Tensor> row_starts_opt = std::nullopt,
std::optional<at::Tensor> indices_opt = std::nullopt) -> FastTopKParams {
const auto B = score.size(0);
TORCH_CHECK(score.dim() == 2 && score.stride(1) == 1);
if (row_starts_opt.has_value()) {
const auto& row_starts = row_starts_opt.value();
TORCH_CHECK(row_starts.dim() == 1);
TORCH_CHECK(row_starts.size(0) == B);
}
TORCH_CHECK(lengths.dim() == 1 && lengths.is_contiguous());
TORCH_CHECK(lengths.size(0) == B);
int32_t* indices_data_ptr = nullptr;
@@ -380,6 +391,7 @@ auto get_params(
return FastTopKParams{
.input = score.data_ptr<float>(),
.row_starts = row_starts_opt.has_value() ? row_starts_opt->data_ptr<int32_t>() : nullptr,
.indices = indices_data_ptr,
.lengths = lengths.data_ptr<int32_t>(),
.input_stride = score.stride(0),
@@ -398,11 +410,15 @@ void setup_kernel_smem_once() {
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
void fast_topk_interface(const at::Tensor& score, at::Tensor& indices, const at::Tensor& lengths) {
void fast_topk_interface(
const at::Tensor& score, at::Tensor& indices, const at::Tensor& lengths, std::optional<at::Tensor> row_starts_opt) {
CHECK_CUDA(score);
CHECK_CUDA(indices);
if (row_starts_opt.has_value()) {
CHECK_CUDA(row_starts_opt.value());
}
CHECK_CUDA(lengths);
const auto params = get_params(score, lengths, indices);
const auto params = get_params(score, lengths, row_starts_opt, indices);
const auto B = score.size(0);
const auto stream = at::cuda::getCurrentCUDAStream().stream();
const auto grid = dim3{static_cast<uint32_t>(B)};
@@ -418,13 +434,17 @@ void fast_topk_transform_interface(
const at::Tensor& lengths,
at::Tensor& dst_page_table,
const at::Tensor& src_page_table,
const at::Tensor& cu_seqlens_q) {
const at::Tensor& cu_seqlens_q,
std::optional<at::Tensor> row_starts_opt) {
CHECK_CUDA(score);
CHECK_CUDA(lengths);
CHECK_CUDA(dst_page_table);
CHECK_CUDA(src_page_table);
CHECK_CUDA(cu_seqlens_q);
const auto params = get_params(score, lengths);
if (row_starts_opt.has_value()) {
CHECK_CUDA(row_starts_opt.value());
}
const auto params = get_params(score, lengths, row_starts_opt);
const auto B = score.size(0);
TORCH_CHECK(dst_page_table.dim() == 2 && dst_page_table.is_contiguous());
TORCH_CHECK(src_page_table.dim() == 2 && src_page_table.stride(1) == 1);
@@ -442,7 +462,10 @@ void fast_topk_transform_interface(
const auto src_stride = src_page_table.stride(0);
// dispatch to decode or prefill
const auto is_decode = (prefill_bs == B);
// extend and draft extend: row_starts_opt is not null, invokes the prefill kernel
// decode: row_starts_opt is null, invokes the decode kernel
// target verify: row_starts_opt is null, invokes the prefill kernel
const auto is_decode = !row_starts_opt.has_value() && prefill_bs == B;
if (is_decode) {
setup_kernel_smem_once<topk_transform_decode_kernel, kSmem>();
topk_transform_decode_kernel<<<grid, block, kSmem, stream>>>(
@@ -466,13 +489,17 @@ void fast_topk_transform_ragged_interface(
const at::Tensor& score,
const at::Tensor& lengths,
at::Tensor& topk_indices_ragged,
const at::Tensor& topk_indices_offset) {
const at::Tensor& topk_indices_offset,
std::optional<at::Tensor> row_starts_opt) {
CHECK_CUDA(score);
CHECK_CUDA(lengths);
CHECK_CUDA(topk_indices_ragged);
CHECK_CUDA(topk_indices_offset);
if (row_starts_opt.has_value()) {
CHECK_CUDA(row_starts_opt.value());
}
const auto params = get_params(score, lengths);
const auto params = get_params(score, lengths, row_starts_opt);
const auto B = score.size(0);
TORCH_CHECK(topk_indices_ragged.dim() == 2 && topk_indices_ragged.is_contiguous());
TORCH_CHECK(topk_indices_offset.dim() == 1);