gahow/agentic-pd-hybrid
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: vendor sglang v0.5.10 snapshot

Browse Source
This commit is contained in:
Gahow Wang
2026-04-24 12:29:36 +00:00
parent 78f0d15221
commit bded08301f
4308 changed files with 1200894 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

170
third_party/sglang/sgl-kernel/csrc/elementwise/activation.cu vendored Normal file
View File

@@ -0,0 +1,170 @@
/*
* 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/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/all.h>
#ifndef USE_ROCM
#include <flashinfer/activation.cuh>
#include "utils.h"
#else
#include "hip/hip_act_and_mul.cuh"
#endif
// Adapted from flashinfer activation
// https://github.com/flashinfer-ai/flashinfer/blob/4e8eb1879f9c3ba6d75511e5893183bf8f289a62/csrc/activation.cu#L44
namespace detail {
template <typename T>
__device__ __forceinline__ float to_f32(const T& x) {
#if USE_ROCM
return castToFloat(x);
#else
return static_cast<float>(x);
#endif
}
template <typename T>
__device__ __forceinline__ T from_f32(float f32) {
#if USE_ROCM
return castFromFloat<T>(f32);
#else
return static_cast<T>(f32);
#endif
}
} // namespace detail
template <typename T>
__device__ __forceinline__ T silu(const T& x) {
float f32_val = detail::to_f32(x);
return detail::from_f32<T>(f32_val / (1.0f + expf(-f32_val)));
}
template <typename T>
__device__ __forceinline__ T gelu(const T& x) {
constexpr float kAlpha = M_SQRT1_2;
float f32_val = detail::to_f32(x);
return detail::from_f32<T>(f32_val * (0.5f * (1.0f + erf(f32_val * kAlpha))));
}
// gelu_quick(x) = x * torch.sigmoid(1.702 * x)
template <typename T>
__device__ __forceinline__ T gelu_quick_act(const T& x) {
float f32_val = detail::to_f32(x);
return detail::from_f32<T>(f32_val / (1.0f + expf(-f32_val * 1.702f)));
}
template <typename T>
__device__ __forceinline__ T gelu_tanh(const T& x) {
constexpr float kAlpha = 0.044715f;
constexpr float kBeta = 0.7978845608028654f;
float f32_val = detail::to_f32(x);
const float cdf = 0.5f * (1.0f + tanhf((kBeta * (f32_val + kAlpha * f32_val * f32_val * f32_val))));
return detail::from_f32<T>(f32_val * cdf);
}
void silu_and_mul(at::Tensor& out, at::Tensor& input) {
int d = input.size(-1) / 2;
int64_t num_tokens = input.numel() / input.size(-1);
dim3 grid(num_tokens);
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(input.scalar_type(), c_type, [&] {
uint32_t vec_size = 16 / sizeof(c_type);
dim3 block(std::min(d / vec_size, 1024U));
#if USE_ROCM
sgl_hip::activation::act_and_mul_kernel<c_type, silu>
<<<grid, block, 0, stream>>>(static_cast<c_type*>(out.data_ptr()), static_cast<c_type*>(input.data_ptr()), d);
#else
flashinfer::activation::act_and_mul_kernel<c_type, silu>
<<<grid, block, 0, stream>>>(static_cast<c_type*>(out.data_ptr()), static_cast<c_type*>(input.data_ptr()), d);
#endif
return true;
});
}
void gelu_tanh_and_mul(at::Tensor& out, at::Tensor& input) {
int d = input.size(-1) / 2;
int64_t num_tokens = input.numel() / input.size(-1);
dim3 grid(num_tokens);
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(input.scalar_type(), c_type, [&] {
uint32_t vec_size = 16 / sizeof(c_type);
dim3 block(std::min(d / vec_size, 1024U));
#if USE_ROCM
sgl_hip::activation::act_and_mul_kernel<c_type, gelu_tanh>
<<<grid, block, 0, stream>>>(static_cast<c_type*>(out.data_ptr()), static_cast<c_type*>(input.data_ptr()), d);
#else
flashinfer::activation::act_and_mul_kernel<c_type, gelu_tanh>
<<<grid, block, 0, stream>>>(static_cast<c_type*>(out.data_ptr()), static_cast<c_type*>(input.data_ptr()), d);
#endif
return true;
});
}
void gelu_and_mul(at::Tensor& out, at::Tensor& input) {
int d = input.size(-1) / 2;
int64_t num_tokens = input.numel() / input.size(-1);
dim3 grid(num_tokens);
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(input.scalar_type(), c_type, [&] {
uint32_t vec_size = 16 / sizeof(c_type);
dim3 block(std::min(d / vec_size, 1024U));
#if USE_ROCM
sgl_hip::activation::act_and_mul_kernel<c_type, gelu>
<<<grid, block, 0, stream>>>(static_cast<c_type*>(out.data_ptr()), static_cast<c_type*>(input.data_ptr()), d);
#else
flashinfer::activation::act_and_mul_kernel<c_type, gelu>
<<<grid, block, 0, stream>>>(static_cast<c_type*>(out.data_ptr()), static_cast<c_type*>(input.data_ptr()), d);
#endif
return true;
});
}
#if USE_ROCM
void gelu_quick(at::Tensor& out, const at::Tensor& input) {
int d = input.size(-1);
int64_t num_tokens = input.numel() / input.size(-1);
dim3 grid(num_tokens);
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(input.scalar_type(), c_type, [&] {
uint32_t vec_size = 16 / sizeof(c_type);
dim3 block(std::min(d / vec_size, 1024U));
sgl_hip::activation::act_only_kernel<c_type, gelu_quick_act>
<<<grid, block, 0, stream>>>(static_cast<c_type*>(out.data_ptr()), static_cast<c_type*>(input.data_ptr()), d);
return true;
});
}
#endif

218
third_party/sglang/sgl-kernel/csrc/elementwise/concat_mla.cu vendored Normal file
View File

@@ -0,0 +1,218 @@
#include <ATen/cuda/CUDAContext.h>
#include <ATen/cuda/CUDADataType.h>
#include <cuda_runtime.h>
#include "pytorch_extension_utils.h"
#include "utils.cuh"
constexpr int NUM_LOCAL_HEADS = 128;
constexpr int QK_NOPE_HEAD_DIM = 128;
constexpr int QK_ROPE_HEAD_DIM = 64;
constexpr int K_HEAD_DIM = QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM;
constexpr int HEAD_CHUNK_SIZE = 16;
constexpr int NUM_HEAD_CHUNKS = NUM_LOCAL_HEADS / HEAD_CHUNK_SIZE;
__global__ void concat_mla_k_kernel(
nv_bfloat16* __restrict__ k,
const nv_bfloat16* __restrict__ k_nope,
const nv_bfloat16* __restrict__ k_rope,
const int num_tokens,
const int64_t k_stride_0,
const int k_stride_1,
const int64_t k_nope_stride_0,
const int k_nope_stride_1,
const int64_t k_rope_stride_0) {
const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
const int token_id = flat_warp_id / NUM_HEAD_CHUNKS;
const int head_chunk_id = flat_warp_id % NUM_HEAD_CHUNKS;
const int lane_id = get_lane_id();
if (token_id >= num_tokens) return;
using NopeVec = int2; // 8B/thread，32 thread = 256B/row
using RopeVec = int; // 4B/thread，32 thread = 128B/row
static_assert(sizeof(NopeVec) * 32 == QK_NOPE_HEAD_DIM * sizeof(nv_bfloat16), "nope vec mismatch");
static_assert(sizeof(RopeVec) * 32 == QK_ROPE_HEAD_DIM * sizeof(nv_bfloat16), "rope vec mismatch");
const int head_row0 = head_chunk_id * HEAD_CHUNK_SIZE;
const int2* __restrict__ nope_src =
reinterpret_cast<const int2*>(k_nope + token_id * k_nope_stride_0 + head_row0 * k_nope_stride_1) + lane_id;
int2* __restrict__ nope_dst = reinterpret_cast<int2*>(k + token_id * k_stride_0 + head_row0 * k_stride_1) + lane_id;
int* __restrict__ rope_dst =
reinterpret_cast<int*>(k + token_id * k_stride_0 + head_row0 * k_stride_1 + QK_NOPE_HEAD_DIM) + lane_id;
const int nope_src_stride_v = (k_nope_stride_1 >> 2); // int2 covers 4 bf16
const int nope_dst_stride_v = (k_stride_1 >> 2);
const int rope_dst_stride_v = (k_stride_1 >> 1); // int covers 2 bf16
const int* rope_base = reinterpret_cast<const int*>(k_rope + token_id * k_rope_stride_0);
const RopeVec rope_val = ld_na_global_v1(rope_base + lane_id);
prefetch_L2(nope_src);
NopeVec cur = ld_na_global_v2(nope_src);
#pragma unroll
for (int i = 0; i < HEAD_CHUNK_SIZE; ++i) {
NopeVec next;
if (i + 1 < HEAD_CHUNK_SIZE) {
const int2* next_src = nope_src + nope_src_stride_v;
prefetch_L2(next_src);
next = ld_na_global_v2(next_src);
}
st_na_global_v2(nope_dst, cur);
st_na_global_v1(rope_dst, rope_val);
nope_src += nope_src_stride_v;
nope_dst += nope_dst_stride_v;
rope_dst += rope_dst_stride_v;
cur = next;
}
}
inline void check_tensor(const at::Tensor& t, int64_t shape0, int64_t shape1, int64_t shape2, c10::ScalarType dtype) {
TORCH_CHECK_EQ(t.dim(), 3);
TORCH_CHECK_EQ(t.size(0), shape0);
TORCH_CHECK_EQ(t.size(1), shape1);
TORCH_CHECK_EQ(t.size(2), shape2);
TORCH_CHECK_EQ(t.dtype(), dtype);
TORCH_CHECK(t.device().is_cuda());
TORCH_CHECK_EQ(((int64_t)t.data_ptr()) % 16, 0); // alignment
}
void concat_mla_k(at::Tensor k, at::Tensor k_nope, at::Tensor k_rope) {
const int num_tokens = k.size(0);
check_tensor(k, num_tokens, NUM_LOCAL_HEADS, K_HEAD_DIM, at::kBFloat16);
check_tensor(k_nope, num_tokens, NUM_LOCAL_HEADS, QK_NOPE_HEAD_DIM, at::kBFloat16);
check_tensor(k_rope, num_tokens, 1, QK_ROPE_HEAD_DIM, at::kBFloat16);
TORCH_CHECK_EQ(k.stride(2), 1);
TORCH_CHECK_EQ(k_nope.stride(2), 1);
TORCH_CHECK_EQ(k_rope.stride(2), 1);
const auto stream = at::cuda::getCurrentCUDAStream().stream();
constexpr int num_warps_per_block = 32;
const int grid_size = ceil_div(num_tokens * NUM_HEAD_CHUNKS, num_warps_per_block);
const int block_size = num_warps_per_block * 32;
concat_mla_k_kernel<<<grid_size, block_size, 0, stream>>>(
reinterpret_cast<nv_bfloat16*>(k.data_ptr()),
reinterpret_cast<nv_bfloat16*>(k_nope.data_ptr()),
reinterpret_cast<nv_bfloat16*>(k_rope.data_ptr()),
num_tokens,
k.stride(0),
k.stride(1),
k_nope.stride(0),
k_nope.stride(1),
k_rope.stride(0));
cudaError_t err = cudaGetLastError();
TORCH_CHECK(err == cudaSuccess, "CUDA kernel launch failed: ", cudaGetErrorString(err));
}
// ============================== concat_mla_absorb_q ==============================
// TODO give a name prefix, also maybe refactor code above
constexpr int A_LAST_DIM = 512;
constexpr int B_LAST_DIM = 64;
__global__ void concat_mla_absorb_q_kernel(
nv_bfloat16* a,
nv_bfloat16* b,
nv_bfloat16* out,
const int num_items,
const int dim_1,
const int64_t a_stride_0,
const int a_stride_1,
const int64_t b_stride_0,
const int b_stride_1,
const int64_t out_stride_0,
const int out_stride_1) {
const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
const int lane_id = get_lane_id();
const int idx_0 = flat_warp_id / dim_1;
const int idx_1 = flat_warp_id % dim_1;
if (flat_warp_id >= num_items) {
return;
}
using ABufType = int4;
constexpr int A_NUM_UNROLL = 2;
static_assert(sizeof(ABufType) * A_NUM_UNROLL == A_LAST_DIM * sizeof(a[0]) / 32);
ABufType a_buf[A_NUM_UNROLL];
using BBufType = int;
constexpr int B_NUM_UNROLL = 1;
static_assert(sizeof(BBufType) * B_NUM_UNROLL == B_LAST_DIM * sizeof(b[0]) / 32);
BBufType b_buf;
{
const BBufType* base_addr = reinterpret_cast<BBufType*>(b + idx_0 * b_stride_0 + idx_1 * b_stride_1);
b_buf = *(base_addr + lane_id);
}
#pragma unroll
for (int i = 0; i < A_NUM_UNROLL; ++i) {
const ABufType* base_addr = reinterpret_cast<ABufType*>(a + idx_0 * a_stride_0 + idx_1 * a_stride_1);
a_buf[i] = *(base_addr + i * 32 + lane_id);
}
{
BBufType* base_addr = reinterpret_cast<BBufType*>(out + idx_0 * out_stride_0 + idx_1 * out_stride_1 + A_LAST_DIM);
*(base_addr + lane_id) = b_buf;
}
#pragma unroll
for (int i = 0; i < A_NUM_UNROLL; ++i) {
ABufType* base_addr = reinterpret_cast<ABufType*>(out + idx_0 * out_stride_0 + idx_1 * out_stride_1);
*(base_addr + i * 32 + lane_id) = a_buf[i];
}
}
inline void check_tensor_concat_mla_absorb_q(const at::Tensor& t, int64_t shape2) {
TORCH_CHECK_EQ(t.dim(), 3);
TORCH_CHECK_EQ(t.size(2), shape2);
TORCH_CHECK_EQ(t.stride(2), 1);
TORCH_CHECK_EQ(t.dtype(), at::kBFloat16);
TORCH_CHECK(t.device().is_cuda());
TORCH_CHECK_EQ(((int64_t)t.data_ptr()) % 16, 0); // alignment
}
// TODO further optimize it later
void concat_mla_absorb_q(at::Tensor a, at::Tensor b, at::Tensor out) {
check_tensor_concat_mla_absorb_q(a, A_LAST_DIM);
check_tensor_concat_mla_absorb_q(b, B_LAST_DIM);
check_tensor_concat_mla_absorb_q(out, A_LAST_DIM + B_LAST_DIM);
const auto stream = at::cuda::getCurrentCUDAStream().stream();
TORCH_CHECK_EQ(a.size(0) * a.size(1), b.size(0) * b.size(1));
TORCH_CHECK_EQ(a.size(1), b.size(1));
const int num_items = a.size(0) * a.size(1);
constexpr int num_warps_per_block = 32;
const int grid_size = ceil_div(num_items, num_warps_per_block);
const int block_size = num_warps_per_block * 32;
concat_mla_absorb_q_kernel<<<grid_size, block_size, 0, stream>>>(
reinterpret_cast<nv_bfloat16*>(a.data_ptr()),
reinterpret_cast<nv_bfloat16*>(b.data_ptr()),
reinterpret_cast<nv_bfloat16*>(out.data_ptr()),
num_items,
a.size(1),
a.stride(0),
a.stride(1),
b.stride(0),
b.stride(1),
out.stride(0),
out.stride(1));
cudaError_t err = cudaGetLastError();
TORCH_CHECK(err == cudaSuccess, "CUDA kernel launch failed: ", cudaGetErrorString(err));
}
// test-1

58
third_party/sglang/sgl-kernel/csrc/elementwise/copy.cu vendored Normal file
View File

@@ -0,0 +1,58 @@
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/all.h>
#include <vector>
template <int N>
struct InputArray {
int values[N];
};
template <int N>
__global__ void copy_to_gpu_no_ce_kernel(const InputArray<N> input_array, int* output) {
int idx = threadIdx.x + blockIdx.x * blockDim.x;
if (idx < N) {
output[idx] = input_array.values[idx];
}
}
template <int N>
void copy_to_gpu_no_ce_impl(const at::Tensor& input, at::Tensor& output) {
TORCH_CHECK(input.dim() == 1, "input must be 1-D");
TORCH_CHECK(static_cast<int>(input.numel()) == N, "input numel must equal template N");
TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
TORCH_CHECK(input.dtype() == torch::kInt32, "input dtype must be int32");
TORCH_CHECK(output.dim() == 1, "output dim");
TORCH_CHECK(static_cast<int>(output.numel()) == N, "output size");
TORCH_CHECK(output.is_contiguous(), "output contiguous");
TORCH_CHECK(output.dtype() == torch::kInt32, "output dtype");
TORCH_CHECK(input.device().is_cpu(), "input must be a CPU tensor");
TORCH_CHECK(output.device().is_cuda(), "output must be a CUDA tensor");
InputArray<N> input_array;
const int* input_ptr = input.data_ptr<int>();
for (int i = 0; i < N; ++i)
input_array.values[i] = input_ptr[i];
// may use multi thread blocks if performance bottleneck
dim3 grid(1);
dim3 block(static_cast<int>(input.numel()));
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
copy_to_gpu_no_ce_kernel<<<grid, block, 0, stream>>>(input_array, output.data_ptr<int>());
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
void copy_to_gpu_no_ce(const at::Tensor& input, at::Tensor& output) {
int N = static_cast<int>(input.numel());
// Can use macro if there are more N needed
if (N == 72) {
copy_to_gpu_no_ce_impl<72>(input, output);
} else if (N == 64) {
copy_to_gpu_no_ce_impl<64>(input, output);
} else {
TORCH_CHECK(false, "unexpected N");
}
}

59
third_party/sglang/sgl-kernel/csrc/elementwise/fused_add_rms_norm_kernel.cu vendored Normal file
View File

@@ -0,0 +1,59 @@
/* Copyright 2025 SGLang Team. All Rights Reserved.
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/CUDAContext.h>
#include <flashinfer/norm.cuh>
#include "utils.h"
using namespace flashinfer;
void sgl_fused_add_rmsnorm(
torch::Tensor input, torch::Tensor residual, torch::Tensor weight, double eps, bool enable_pdl) {
CHECK_INPUT(input);
CHECK_INPUT(residual);
CHECK_INPUT(weight);
auto device = input.device();
CHECK_EQ(residual.device(), device);
CHECK_EQ(weight.device(), device);
CHECK_DIM(2, input); // input: (batch_size, hidden_size)
CHECK_DIM(2, residual); // residual: (batch_size, hidden_size)
CHECK_DIM(1, weight); // weight: (hidden_size)
CHECK_EQ(input.size(0), residual.size(0));
CHECK_EQ(input.size(1), residual.size(1));
CHECK_EQ(input.size(1), weight.size(0));
unsigned int batch_size = input.size(0);
unsigned int hidden_size = input.size(1);
cudaStream_t torch_current_stream = at::cuda::getCurrentCUDAStream();
// support float16, bfloat16 and float32
DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(input.scalar_type(), c_type, [&] {
cudaError_t status = norm::FusedAddRMSNorm(
static_cast<c_type*>(input.data_ptr()),
static_cast<c_type*>(residual.data_ptr()),
static_cast<c_type*>(weight.data_ptr()),
batch_size,
hidden_size,
input.stride(0),
residual.stride(0),
eps,
enable_pdl,
torch_current_stream);
TORCH_CHECK(
status == cudaSuccess, "FusedAddRMSNorm failed with error code " + std::string(cudaGetErrorString(status)));
return true;
});
}

208
third_party/sglang/sgl-kernel/csrc/elementwise/pos_enc.cu vendored Normal file
View File

@@ -0,0 +1,208 @@
// Adapted from
// https://github.com/vllm-project/vllm/blob/014ece97c7aa49084a1119dca792af081a18dbc1/csrc/pos_encoding_kernels.cu
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/all.h>
#include "utils.h"
template <typename scalar_t, bool IS_NEOX>
inline __device__ void apply_token_rotary_embedding(
scalar_t* __restrict__ arr,
const scalar_t* __restrict__ cos_ptr,
const scalar_t* __restrict__ sin_ptr,
int rot_offset,
int embed_dim) {
int x_index, y_index;
scalar_t cos, sin;
if (IS_NEOX) {
// GPT-NeoX style rotary embedding.
x_index = rot_offset;
y_index = embed_dim + rot_offset;
cos = SGLANG_LDG(cos_ptr + x_index);
sin = SGLANG_LDG(sin_ptr + x_index);
} else {
// GPT-J style rotary embedding.
x_index = 2 * rot_offset;
y_index = 2 * rot_offset + 1;
cos = SGLANG_LDG(cos_ptr + x_index / 2);
sin = SGLANG_LDG(sin_ptr + x_index / 2);
}
const scalar_t x = arr[x_index];
const scalar_t y = arr[y_index];
arr[x_index] = x * cos - y * sin;
arr[y_index] = y * cos + x * sin;
}
template <typename scalar_t, bool IS_NEOX>
inline __device__ void apply_rotary_embedding(
scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads,
// head_size] or [num_tokens, num_heads,
// head_size]
scalar_t* __restrict__ key, // nullptr or
// [batch_size, seq_len, num_kv_heads,
// head_size] or [num_tokens, num_kv_heads,
// head_size]
const scalar_t* cache_ptr,
const int head_size,
const int num_heads,
const int num_kv_heads,
const int rot_dim,
const int token_idx,
const int64_t query_stride,
const int64_t key_stride,
const int64_t head_stride) {
const int embed_dim = rot_dim / 2;
const scalar_t* cos_ptr = cache_ptr;
const scalar_t* sin_ptr = cache_ptr + embed_dim;
const int nq = num_heads * embed_dim;
for (int i = threadIdx.x; i < nq; i += blockDim.x) {
const int head_idx = i / embed_dim;
const int64_t token_head = token_idx * query_stride + head_idx * head_stride;
const int rot_offset = i % embed_dim;
apply_token_rotary_embedding<scalar_t, IS_NEOX>(query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
}
if (key != nullptr) {
const int nk = num_kv_heads * embed_dim;
for (int i = threadIdx.x; i < nk; i += blockDim.x) {
const int head_idx = i / embed_dim;
const int64_t token_head = token_idx * key_stride + head_idx * head_stride;
const int rot_offset = i % embed_dim;
apply_token_rotary_embedding<scalar_t, IS_NEOX>(key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
}
}
}
template <typename scalar_t, bool IS_NEOX>
__global__ void rotary_embedding_kernel(
const int64_t* __restrict__ positions, // [batch_size, seq_len] or
// [num_tokens]
scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads,
// head_size] or [num_tokens, num_heads,
// head_size]
scalar_t* __restrict__ key, // nullptr or
// [batch_size, seq_len, num_kv_heads,
// head_size] or [num_tokens, num_kv_heads,
// head_size]
const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim //
// 2]
const int rot_dim,
const int64_t query_stride,
const int64_t key_stride,
const int64_t head_stride,
const int num_heads,
const int num_kv_heads,
const int head_size) {
// Each thread block is responsible for one token.
const int token_idx = blockIdx.x;
int64_t pos = positions[token_idx];
const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
apply_rotary_embedding<scalar_t, IS_NEOX>(
query,
key,
cache_ptr,
head_size,
num_heads,
num_kv_heads,
rot_dim,
token_idx,
query_stride,
key_stride,
head_stride);
}
void rotary_embedding(
torch::Tensor& positions, // [batch_size, seq_len] or [num_tokens]
torch::Tensor& query, // [batch_size, seq_len, num_heads * head_size] or
// [num_tokens, num_heads * head_size] or
// [batch_size, seq_len, num_heads, head_size] or
// [num_tokens, num_heads, head_size]
std::optional<torch::Tensor> key,
// null or
// [batch_size, seq_len, num_kv_heads * head_size] or
// [num_tokens, num_kv_heads * head_size] or
// [batch_size, seq_len, num_heads, head_size] or
// [num_tokens, num_heads, head_size]
int64_t head_size,
torch::Tensor& cos_sin_cache, // [max_position, rot_dim]
bool is_neox) {
// num_tokens = batch_size * seq_len
int64_t num_tokens = positions.numel();
int positions_ndim = positions.dim();
// Make sure num_tokens dim is consistent across positions, query, and key
TORCH_CHECK(
positions_ndim == 1 || positions_ndim == 2, "positions must have shape [num_tokens] or [batch_size, seq_len]");
if (positions_ndim == 1) {
TORCH_CHECK(
query.size(0) == positions.size(0) && (!key.has_value() || key->size(0) == positions.size(0)),
"query, key and positions must have the same number of tokens");
}
if (positions_ndim == 2) {
TORCH_CHECK(
query.size(0) == positions.size(0) && (!key.has_value() || key->size(0) == positions.size(0)) &&
query.size(1) == positions.size(1) && (!key.has_value() || key->size(1) == positions.size(1)),
"query, key and positions must have the same batch_size and seq_len");
}
// Make sure head_size is valid for query and key
// hidden_size = num_heads * head_size
int query_hidden_size = query.numel() / num_tokens;
int key_hidden_size = key.has_value() ? key->numel() / num_tokens : 0;
TORCH_CHECK(query_hidden_size % head_size == 0);
TORCH_CHECK(key_hidden_size % head_size == 0);
// Make sure query and key have consistent number of heads
int num_heads = query_hidden_size / head_size;
int num_kv_heads = key.has_value() ? key_hidden_size / head_size : num_heads;
TORCH_CHECK(num_heads % num_kv_heads == 0);
int rot_dim = cos_sin_cache.size(1);
int seq_dim_idx = positions_ndim - 1;
int64_t query_stride = query.stride(seq_dim_idx);
int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;
// Determine head stride: for [*, heads, head_size] use stride of last dim;
// for flat [*, heads*head_size], heads blocks are contiguous of size
// head_size
int query_ndim = query.dim();
int64_t head_stride = (query_ndim == positions_ndim + 2) ? query.stride(-2) : head_size;
dim3 grid(num_tokens);
dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
DISPATCH_FLOAT_TYPES(query.scalar_type(), "rotary_embedding", [&] {
if (is_neox) {
rotary_embedding_kernel<scalar_t, true><<<grid, block, 0, stream>>>(
positions.data_ptr<int64_t>(),
query.data_ptr<scalar_t>(),
key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
cos_sin_cache.data_ptr<scalar_t>(),
rot_dim,
query_stride,
key_stride,
head_stride,
num_heads,
num_kv_heads,
head_size);
} else {
rotary_embedding_kernel<scalar_t, false><<<grid, block, 0, stream>>>(
positions.data_ptr<int64_t>(),
query.data_ptr<scalar_t>(),
key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
cos_sin_cache.data_ptr<scalar_t>(),
rot_dim,
query_stride,
key_stride,
head_stride,
num_heads,
num_kv_heads,
head_size);
}
});
}

467
third_party/sglang/sgl-kernel/csrc/elementwise/pos_enc.cuh vendored Normal file
View File

@@ -0,0 +1,467 @@
/*
* Copyright (c) 2023 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.
*/
#ifndef SGL_POS_ENC_CUH_
#define SGL_POS_ENC_CUH_
#include <flashinfer/pos_enc.cuh> // upstream
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
#endif // SGL_POS_ENC_CUH_

546
third_party/sglang/sgl-kernel/csrc/elementwise/topk.cu vendored Normal file
View File

@@ -0,0 +1,546 @@
/**
* @NOTE: This file is adapted from
* https://github.com/tile-ai/tilelang/blob/main/examples/deepseek_v32/topk_selector.py
* We:
* 1. adapt from tilelang to pure cuda
* 2. optimize the performance a little
* 3. fix the potential illegal memory access
*/
#include <ATen/core/TensorBase.h>
#include <ATen/core/TensorBody.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/macros/Macros.h>
#include <c10/util/Exception.h>
#include <cuda.h>
#include <cuda_fp16.h>
#include <cstddef>
#include <cstdint>
#include <optional>
namespace {
constexpr int TopK = 2048;
constexpr int kThreadsPerBlock = 1024;
#ifdef USE_ROCM
// On ROCm, the per-workgroup LDS budget depends on the target arch, so we inject a
// per-arch value from `setup_rocm.py` via `-DSGL_TOPK_DYNAMIC_SMEM_BYTES=...`.
#ifdef SGL_TOPK_DYNAMIC_SMEM_BYTES
constexpr size_t kSmem = static_cast<size_t>(SGL_TOPK_DYNAMIC_SMEM_BYTES);
#else
constexpr size_t kSmem = 48 * 1024; // bytes
#endif
#else
// Reduced from 128KB to 32KB to improve occupancy.
// Each radix pass needs at most ~TopK candidates in the threshold bin,
// so 4K entries per round (2 rounds = 8K entries = 32KB) is sufficient.
constexpr size_t kSmem = 8 * 1024 * sizeof(uint32_t); // 32KB (bytes)
#endif
struct FastTopKParams {
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;
};
// when length <= TopK, we can directly write the indices
__device__ void naive_topk_cuda(const float* __restrict__ score, int32_t* __restrict__ indice, int32_t length) {
const auto tid = threadIdx.x;
for (int i = tid; i < TopK; i += kThreadsPerBlock) {
indice[i] = (i < length) ? i : -1;
}
}
// keep the first `length` entries, set others to -1
__device__ void naive_topk_transform(
const float* __restrict__ score,
int32_t length,
int32_t* __restrict__ dst_page_table,
const int32_t* __restrict__ src_page_table) {
const auto tid = threadIdx.x;
for (auto i = tid; i < TopK; i += kThreadsPerBlock) {
dst_page_table[i] = (i < length) ? src_page_table[i] : -1;
}
}
// keep the first `length` entries, set others to -1
__device__ void naive_topk_transform_ragged(
const float* __restrict__ score, int32_t length, int32_t* __restrict__ topk_indices_ragged, int32_t offset) {
const auto tid = threadIdx.x;
for (auto i = tid; i < TopK; i += kThreadsPerBlock) {
topk_indices_ragged[i] = (i < length) ? static_cast<int32_t>(i) + offset : -1;
}
}
__device__ __forceinline__ auto convert_to_uint8(float x) -> uint8_t {
__half h = __float2half_rn(x);
uint16_t bits = __half_as_ushort(h);
uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits) : static_cast<uint16_t>(bits | 0x8000);
return static_cast<uint8_t>(key >> 8);
}
__device__ __forceinline__ auto convert_to_uint32(float x) -> uint32_t {
uint32_t bits = __float_as_uint(x);
return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
}
__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;
constexpr auto BLOCK_SIZE = 1024;
constexpr auto RADIX = 256;
constexpr auto SMEM_INPUT_SIZE = kSmem / (2 * sizeof(int));
alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128];
alignas(128) __shared__ int s_counter;
alignas(128) __shared__ int s_threshold_bin_id;
alignas(128) __shared__ int s_num_input[2];
auto& s_histogram = s_histogram_buf[0];
// allocate for two rounds
extern __shared__ int s_input_idx[][SMEM_INPUT_SIZE];
const int tx = threadIdx.x;
// stage 1: 8bit coarse histogram
if (tx < RADIX + 1) s_histogram[tx] = 0;
__syncthreads();
for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
const auto bin = convert_to_uint8(input[idx + row_start]);
::atomicAdd(&s_histogram[bin], 1);
}
__syncthreads();
const auto run_cumsum = [&] {
#pragma unroll 8
for (int i = 0; i < 8; ++i) {
static_assert(1 << 8 == RADIX);
if (C10_LIKELY(tx < RADIX)) {
const auto j = 1 << i;
const auto k = i & 1;
auto value = s_histogram_buf[k][tx];
if (tx < RADIX - j) {
value += s_histogram_buf[k][tx + j];
}
s_histogram_buf[k ^ 1][tx] = value;
}
__syncthreads();
}
};
run_cumsum();
if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) {
s_threshold_bin_id = tx;
s_num_input[0] = 0;
s_counter = 0;
}
__syncthreads();
const auto threshold_bin = s_threshold_bin_id;
topk -= s_histogram[threshold_bin + 1];
if (topk == 0) {
for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
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;
}
}
__syncthreads();
return;
} else {
__syncthreads();
if (tx < RADIX + 1) {
s_histogram[tx] = 0;
}
__syncthreads();
for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
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);
index[pos] = idx;
} else if (bin == threshold_bin) {
const auto pos = ::atomicAdd(&s_num_input[0], 1);
/// NOTE: (dark) fuse the histogram computation here
if (C10_LIKELY(pos < SMEM_INPUT_SIZE)) {
s_input_idx[0][pos] = idx;
const auto bin = convert_to_uint32(raw_input);
const auto sub_bin = (bin >> 24) & 0xFF;
::atomicAdd(&s_histogram[sub_bin], 1);
}
}
}
__syncthreads();
}
// stage 2: refine with 8bit radix passes
#pragma unroll 4
for (int round = 0; round < 4; ++round) {
__shared__ int s_last_remain;
const auto r_idx = round % 2;
// clip here to prevent overflow
const auto _raw_num_input = s_num_input[r_idx];
const auto num_input = (_raw_num_input < int(SMEM_INPUT_SIZE)) ? _raw_num_input : int(SMEM_INPUT_SIZE);
run_cumsum();
if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) {
s_threshold_bin_id = tx;
s_num_input[r_idx ^ 1] = 0;
s_last_remain = topk - s_histogram[tx + 1];
}
__syncthreads();
const auto threshold_bin = s_threshold_bin_id;
topk -= s_histogram[threshold_bin + 1];
if (topk == 0) {
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 + row_start]) >> offset) & 0xFF;
if (bin > threshold_bin) {
const auto pos = ::atomicAdd(&s_counter, 1);
index[pos] = idx;
}
}
__syncthreads();
break;
} else {
__syncthreads();
if (tx < RADIX + 1) {
s_histogram[tx] = 0;
}
__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 + row_start];
const auto offset = 24 - round * 8;
const auto bin = (convert_to_uint32(raw_input) >> offset) & 0xFF;
if (bin > threshold_bin) {
const auto pos = ::atomicAdd(&s_counter, 1);
index[pos] = idx;
} else if (bin == threshold_bin) {
if (round == 3) {
const auto pos = ::atomicAdd(&s_last_remain, -1);
if (pos > 0) {
index[TopK - pos] = idx;
}
} else {
const auto pos = ::atomicAdd(&s_num_input[r_idx ^ 1], 1);
if (C10_LIKELY(pos < SMEM_INPUT_SIZE)) {
/// NOTE: (dark) fuse the histogram computation here
s_input_idx[r_idx ^ 1][pos] = idx;
const auto bin = convert_to_uint32(raw_input);
const auto sub_bin = (bin >> (offset - 8)) & 0xFF;
::atomicAdd(&s_histogram[sub_bin], 1);
}
}
}
}
__syncthreads();
}
}
}
__global__ __launch_bounds__(kThreadsPerBlock) // topk
void topk_kernel(const FastTopKParams 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, row_start, length);
}
}
__global__ __launch_bounds__(kThreadsPerBlock) // decode
void topk_transform_decode_kernel(
const FastTopKParams params,
int32_t* __restrict__ dst_page_table,
const int32_t* __restrict__ src_page_table,
const int64_t src_stride) {
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;
const auto score = input + bid * input_stride;
if (length <= TopK) {
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, row_start, length);
// copy src[s_indices] to dst, we manually unroll here
static_assert(TopK % kThreadsPerBlock == 0);
static_assert(TopK / kThreadsPerBlock == 2);
const auto idx_0 = tid;
const auto pos_0 = s_indices[idx_0];
dst_page_entry[idx_0] = src_page_entry[pos_0];
const auto idx_1 = tid + kThreadsPerBlock;
const auto pos_1 = s_indices[idx_1];
dst_page_entry[idx_1] = src_page_entry[pos_1];
}
}
__global__ __launch_bounds__(kThreadsPerBlock) // prefill
void topk_transform_prefill_kernel(
const FastTopKParams params,
int32_t* __restrict__ dst_page_table,
const int32_t* __restrict__ src_page_table,
const int64_t src_stride,
const int32_t* __restrict__ cu_seqlens_q,
const int64_t prefill_bs) {
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;
/// NOTE: prefill bs is usually small, we can just use a simple loop here
/// We ensure that last cu_seqlens is equal to number of blocks launched
__shared__ const int32_t* s_src_page_entry;
if (C10_LIKELY(prefill_bs <= kThreadsPerBlock)) {
if (tid < prefill_bs) {
if (bid >= cu_seqlens_q[tid] && bid < cu_seqlens_q[tid + 1]) {
s_src_page_entry = src_page_table + tid * src_stride;
}
}
} else {
for (int64_t i = tid; i < prefill_bs; i += kThreadsPerBlock) {
if (bid >= cu_seqlens_q[i] && bid < cu_seqlens_q[i + 1]) {
s_src_page_entry = src_page_table + i * src_stride;
}
}
}
__syncthreads();
const auto src_page_entry = s_src_page_entry;
if (length <= TopK) {
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, row_start, length);
// copy src[s_indices] to dst, we manually unroll here
static_assert(TopK % kThreadsPerBlock == 0);
static_assert(TopK / kThreadsPerBlock == 2);
const auto idx_0 = tid;
const auto pos_0 = s_indices[idx_0];
dst_page_entry[idx_0] = src_page_entry[pos_0];
const auto idx_1 = tid + kThreadsPerBlock;
const auto pos_1 = s_indices[idx_1];
dst_page_entry[idx_1] = src_page_entry[pos_1];
}
}
__global__ __launch_bounds__(kThreadsPerBlock) // prefill, ragged kv
void topk_transform_prefill_ragged_kernel(
const FastTopKParams params,
int32_t* __restrict__ topk_indices_ragged,
const int32_t* __restrict__ topk_indices_offset) {
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;
const auto offset = topk_indices_offset[bid];
if (length <= TopK) {
return naive_topk_transform_ragged(score, length, dst_indices_entry, offset);
} else {
__shared__ int s_indices[TopK];
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);
const auto idx_0 = tid;
const auto pos_0 = s_indices[idx_0];
dst_indices_entry[idx_0] = pos_0 + offset;
const auto idx_1 = tid + kThreadsPerBlock;
const auto pos_1 = s_indices[idx_1];
dst_indices_entry[idx_1] = pos_1 + offset;
}
}
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;
if (indices_opt.has_value()) {
const auto& indices = indices_opt.value();
TORCH_CHECK(indices.dim() == 2 && indices.is_contiguous());
TORCH_CHECK(indices.size(0) == B);
TORCH_CHECK(indices.size(1) == TopK);
indices_data_ptr = indices.data_ptr<int32_t>();
}
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),
};
}
template <auto* f, size_t max_dynamic_smem>
void setup_kernel_smem_once() {
[[maybe_unused]]
static const auto result = [] {
#ifdef USE_ROCM
// hipify will turn cudaFuncSetAttribute -> hipFuncSetAttribute. On ROCm,
// hipFuncSetAttribute expects `const void*` and hipcc does not accept passing
// a function pointer directly, so cast explicitly.
return ::cudaFuncSetAttribute(
reinterpret_cast<const void*>(f), ::cudaFuncAttributeMaxDynamicSharedMemorySize, max_dynamic_smem);
#else
// CUDA: keep original behavior (no cast needed).
return ::cudaFuncSetAttribute(f, ::cudaFuncAttributeMaxDynamicSharedMemorySize, max_dynamic_smem);
#endif
}();
TORCH_CHECK(result == cudaSuccess, "set_up_kernel_once failed:", ::cudaGetErrorString(result));
}
} // namespace
#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, 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, 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)};
const auto block = dim3{kThreadsPerBlock};
setup_kernel_smem_once<topk_kernel, kSmem>();
topk_kernel<<<grid, block, kSmem, stream>>>(params);
const auto result = cudaGetLastError();
TORCH_CHECK(result == cudaSuccess, "topk kernel failed:", ::cudaGetErrorString(result));
}
void fast_topk_transform_interface(
const at::Tensor& score,
const at::Tensor& lengths,
at::Tensor& dst_page_table,
const at::Tensor& src_page_table,
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);
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);
TORCH_CHECK(cu_seqlens_q.dim() == 1 && cu_seqlens_q.is_contiguous());
const auto prefill_bs = cu_seqlens_q.size(0) - 1;
TORCH_CHECK(dst_page_table.size(0) == B);
TORCH_CHECK(dst_page_table.size(1) == TopK);
TORCH_CHECK(src_page_table.size(0) == prefill_bs);
TORCH_CHECK(prefill_bs <= B); // prefill_bs should be smaller than expanded bs
// launch kernel
const auto stream = at::cuda::getCurrentCUDAStream().stream();
const auto grid = dim3{static_cast<uint32_t>(B)};
const auto block = dim3{kThreadsPerBlock};
const auto src_stride = src_page_table.stride(0);
// dispatch to decode or prefill
// 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>>>(
params, dst_page_table.data_ptr<int32_t>(), src_page_table.data_ptr<int32_t>(), src_stride);
} else {
setup_kernel_smem_once<topk_transform_prefill_kernel, kSmem>();
topk_transform_prefill_kernel<<<grid, block, kSmem, stream>>>(
params,
dst_page_table.data_ptr<int32_t>(),
src_page_table.data_ptr<int32_t>(),
src_stride,
cu_seqlens_q.data_ptr<int32_t>(),
prefill_bs);
}
const auto result = cudaGetLastError();
TORCH_CHECK(result == cudaSuccess, "topk kernel failed:", ::cudaGetErrorString(result));
}
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,
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, 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);
TORCH_CHECK(topk_indices_ragged.size(0) == B);
TORCH_CHECK(topk_indices_ragged.size(1) == TopK);
TORCH_CHECK(topk_indices_offset.size(0) == B);
// launch kernel
const auto stream = at::cuda::getCurrentCUDAStream().stream();
const auto grid = dim3{static_cast<uint32_t>(B)};
const auto block = dim3{kThreadsPerBlock};
setup_kernel_smem_once<topk_transform_prefill_ragged_kernel, kSmem>();
topk_transform_prefill_ragged_kernel<<<grid, block, kSmem, stream>>>(
params, topk_indices_ragged.data_ptr<int32_t>(), topk_indices_offset.data_ptr<int32_t>());
const auto result = cudaGetLastError();
TORCH_CHECK(result == cudaSuccess, "topk kernel failed:", ::cudaGetErrorString(result));
}

72
third_party/sglang/sgl-kernel/csrc/elementwise/utils.cuh vendored Normal file
View File

@@ -0,0 +1,72 @@
// Adapted from https://github.com/deepseek-ai/DeepEP/blob/main/csrc/kernels/utils.cuh
#pragma once
#include <cuda_bf16.h>
#include <cuda_runtime.h>
#include <cstdint>
__forceinline__ __device__ int get_lane_id() {
int lane_id;
asm("mov.s32 %0, %laneid;" : "=r"(lane_id));
return lane_id;
}
int ceil_div(int a, int b) {
return (a + b - 1) / b;
}
__device__ __forceinline__ void st_na_global_v1(const int* ptr, int v) {
asm volatile("st.global.L1::no_allocate.s32 [%0], %1;" ::"l"(ptr), "r"(v) : "memory");
}
__device__ __forceinline__ void st_na_global_v2(const int2* ptr, const int2& v) {
asm volatile("st.global.L1::no_allocate.v2.s32 [%0], {%1, %2};" ::"l"(ptr), "r"(v.x), "r"(v.y) : "memory");
}
__device__ __forceinline__ void st_na_global_v4(const int4* ptr, const int4& v) {
asm volatile(
"st.global.L1::no_allocate.v4.s32 [%0], {%1, %2, %3, %4};" ::"l"(ptr), "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w)
: "memory");
}
__device__ __forceinline__ int ld_na_global_v1(const int* ptr) {
int r;
#ifdef USE_L2_HINT
asm volatile("ld.global.nc.L1::no_allocate.L2::128B.s32 %0, [%1];" : "=r"(r) : "l"(ptr));
#else
asm volatile("ld.global.nc.L1::no_allocate.s32 %0, [%1];" : "=r"(r) : "l"(ptr));
#endif
return r;
}
__device__ __forceinline__ int2 ld_na_global_v2(const int2* ptr) {
int2 r;
#ifdef USE_L2_HINT
asm volatile("ld.global.nc.L1::no_allocate.L2::128B.v2.s32 {%0, %1}, [%2];" : "=r"(r.x), "=r"(r.y) : "l"(ptr));
#else
asm volatile("ld.global.nc.L1::no_allocate.v2.s32 {%0, %1}, [%2];" : "=r"(r.x), "=r"(r.y) : "l"(ptr));
#endif
return r;
}
__device__ __forceinline__ int4 ld_na_global_v4(const int4* ptr) {
int4 r;
#ifdef USE_L2_HINT
asm volatile("ld.global.nc.L1::no_allocate.L2::128B.v4.s32 {%0, %1, %2, %3}, [%4];"
: "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w)
: "l"(ptr));
#else
asm volatile("ld.global.nc.L1::no_allocate.v4.s32 {%0, %1, %2, %3}, [%4];"
: "=r"(r.x), "=r"(r.y), "=r"(r.z), "=r"(r.w)
: "l"(ptr));
#endif
return r;
}
__device__ __forceinline__ void prefetch_L2(const void* p) {
#if defined(ENABLE_L2_PREFETCH)
asm volatile("prefetch.global.L2 [%0];" ::"l"(p));
#endif
}