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

Clean up sgl kernel (#12413)

Browse Source 
Co-authored-by: Byron Hsu <byronhsu1230@gmail.com>
This commit is contained in:
Lianmin Zheng
2025-10-31 01:13:34 -07:00
committed by GitHub
co-authored by Byron Hsu
parent 2e48584b62
commit c0652d907b
10 changed files with 514 additions and 379 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/moe/moe_topk_softmax_kernels.cu
+129 -16
View File
@@ -73,8 +73,13 @@ __device__ float convert_to_float(T x) {
// We have our own implementation of softmax here so we can support transposing the output
// in the softmax kernel when we extend this module to support expert-choice routing.
template <typename T, int TPB>
__launch_bounds__(TPB) __global__
void moeSoftmax(const T* input, const bool* finished, float* output, const int num_cols) {
__launch_bounds__(TPB) __global__ void moeSoftmax(
const T* input,
const bool* finished,
float* output,
const int num_cols,
const float moe_softcapping,
const float* correction_bias) {
using BlockReduce = cub::BlockReduce<float, TPB>;
__shared__ typename BlockReduce::TempStorage tmpStorage;
@@ -90,9 +95,23 @@ __launch_bounds__(TPB) __global__
return;
}
// First pass: Apply transformation, find max, and write transformed values to output
for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
const int idx = thread_row_offset + ii;
threadData = max(convert_to_float<T>(input[idx]), threadData);
float val = convert_to_float<T>(input[idx]);
// Apply tanh softcapping if enabled
if (moe_softcapping != 0.0f) {
val = tanhf(val / moe_softcapping) * moe_softcapping;
}
// Apply correction bias if provided
if (correction_bias != nullptr) {
val = val + correction_bias[ii];
}
output[idx] = val; // Store transformed value
threadData = max(val, threadData);
}
const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, MaxReduceOp());
@@ -102,11 +121,11 @@ __launch_bounds__(TPB) __global__
}
__syncthreads();
// Second pass: Compute sum using transformed values from output
threadData = 0;
for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
const int idx = thread_row_offset + ii;
threadData += exp((convert_to_float<T>(input[idx]) - float_max));
threadData += exp((output[idx] - float_max));
}
const auto Z = BlockReduce(tmpStorage).Sum(threadData);
@@ -116,10 +135,11 @@ __launch_bounds__(TPB) __global__
}
__syncthreads();
// Third pass: Compute final softmax using transformed values from output
for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
const int idx = thread_row_offset + ii;
const float val = exp((convert_to_float<T>(input[idx]) - float_max)) * normalizing_factor;
output[idx] = val;
const float softmax_val = exp((output[idx] - float_max)) * normalizing_factor;
output[idx] = softmax_val;
}
}
@@ -216,7 +236,9 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__ void topkGatingSoftmax(
const int k,
const int start_expert,
const int end_expert,
const bool renormalize) {
const bool renormalize,
const float moe_softcapping,
const float* correction_bias) {
// We begin by enforcing compile time assertions and setting up compile time constants.
static_assert(VPT == (VPT & -VPT), "VPT must be power of 2");
static_assert(NUM_EXPERTS == (NUM_EXPERTS & -NUM_EXPERTS), "NUM_EXPERTS must be power of 2");
@@ -283,16 +305,48 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__ void topkGatingSoftmax(
AccessType* row_chunk_vec_ptr = reinterpret_cast<AccessType*>(&row_chunk_temp);
const AccessType* vec_thread_read_ptr = reinterpret_cast<const AccessType*>(thread_read_ptr);
#pragma unroll
// Note(Byron): interleaved loads to achieve better memory coalescing
// | thread[0] | thread[1] | thread[2] | thread[3] | thread[0] | thread[1] | thread[2] | thread[3] | ...
for (int ii = 0; ii < LDG_PER_THREAD; ++ii) {
row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
}
float row_chunk[VPT];
#pragma unroll
// Note(Byron): upcast logits to float32
for (int ii = 0; ii < VPT; ++ii) {
row_chunk[ii] = convert_to_float<T>(row_chunk_temp[ii]);
}
// Apply tanh softcapping and correction bias
if (moe_softcapping != 0.0f || correction_bias != nullptr) {
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
// Apply tanh softcapping if enabled
if (moe_softcapping != 0.0f) {
val = tanhf(val / moe_softcapping) * moe_softcapping;
}
// Apply correction bias if provided
if (correction_bias != nullptr) {
/*
LDG is interleaved
|thread0 LDG| |thread1 LDG| |thread0 LDG| |thread1 LDG|
|--------- group0 --------| |----------group1 --------|
^ local2
*/
const int group_id = ii / ELTS_PER_LDG;
const int local_id = ii % ELTS_PER_LDG;
const int expert_idx = first_elt_read_by_thread + group_id * THREADS_PER_ROW * ELTS_PER_LDG + local_id;
val = val + correction_bias[expert_idx];
}
row_chunk[ii] = val;
}
}
// First, we perform a max reduce within the thread. We can do the max in fp16 safely (I think) and just
// convert to float afterwards for the exp + sum reduction.
float thread_max = row_chunk[0];
@@ -301,9 +355,15 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__ void topkGatingSoftmax(
thread_max = max(thread_max, row_chunk[ii]);
}
/*********************************/
/********* Softmax Begin *********/
/*********************************/
// Now, we find the max within the thread group and distribute among the threads. We use a butterfly reduce.
// lane id: 0-31 within a warp
#pragma unroll
for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
// butterfly reduce with (lane id ^ mask)
thread_max = max(thread_max, SGLANG_SHFL_XOR_SYNC_WIDTH(0xffffffff, thread_max, mask, THREADS_PER_ROW));
}
@@ -333,6 +393,9 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__ void topkGatingSoftmax(
for (int ii = 0; ii < VPT; ++ii) {
row_chunk[ii] = row_chunk[ii] * reciprocal_row_sum;
}
/*******************************/
/********* Softmax End *********/
/*******************************/
// Now, softmax_res contains the softmax of the row chunk. Now, I want to find the topk elements in each row, along
// with the max index.
@@ -438,6 +501,8 @@ void topkGatingSoftmaxLauncherHelper(
const int start_expert,
const int end_expert,
const bool renormalize,
const float moe_softcapping,
const float* correction_bias,
cudaStream_t stream) {
static constexpr std::size_t MAX_BYTES_PER_LDG = 16;
@@ -450,12 +515,33 @@ void topkGatingSoftmaxLauncherHelper(
dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
topkGatingSoftmax<T, VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG><<<num_blocks, block_dim, 0, stream>>>(
input, finished, output, num_rows, indices, k, start_expert, end_expert, renormalize);
input,
finished,
output,
num_rows,
indices,
k,
start_expert,
end_expert,
renormalize,
moe_softcapping,
correction_bias);
}
#define LAUNCH_SOFTMAX(TYPE, NUM_EXPERTS, WARPS_PER_TB) \
topkGatingSoftmaxLauncherHelper<TYPE, NUM_EXPERTS, WARPS_PER_TB>( \
gating_output, nullptr, topk_weights, topk_indices, num_tokens, topk, 0, num_experts, renormalize, stream);
gating_output, \
nullptr, \
topk_weights, \
topk_indices, \
num_tokens, \
topk, \
0, \
num_experts, \
renormalize, \
moe_softcapping, \
correction_bias, \
stream);
template <typename T>
void topkGatingSoftmaxKernelLauncher(
@@ -467,6 +553,8 @@ void topkGatingSoftmaxKernelLauncher(
const int num_experts,
const int topk,
const bool renormalize,
const float moe_softcapping,
const float* correction_bias,
cudaStream_t stream) {
static constexpr int WARPS_PER_TB = 4;
switch (num_experts) {
@@ -502,7 +590,8 @@ void topkGatingSoftmaxKernelLauncher(
softmax_workspace != nullptr,
"softmax_workspace must be provided for num_experts that are not a power of 2.");
static constexpr int TPB = 256;
moeSoftmax<T, TPB><<<num_tokens, TPB, 0, stream>>>(gating_output, nullptr, softmax_workspace, num_experts);
moeSoftmax<T, TPB><<<num_tokens, TPB, 0, stream>>>(
gating_output, nullptr, softmax_workspace, num_experts, moe_softcapping, correction_bias);
moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
softmax_workspace, nullptr, topk_weights, topk_indices, num_experts, topk, 0, num_experts, renormalize);
}
@@ -510,11 +599,12 @@ void topkGatingSoftmaxKernelLauncher(
}
void topk_softmax(
torch::Tensor& topk_weights, // [num_tokens, topk]
torch::Tensor& topk_indices, // [num_tokens, topk]
torch::Tensor& gating_output,
const bool renormalize) // [num_tokens, num_experts]
{
torch::Tensor& topk_weights, // [num_tokens, topk]
torch::Tensor& topk_indices, // [num_tokens, topk]
torch::Tensor& gating_output, // [num_tokens, num_experts]
const bool renormalize,
const double moe_softcapping,
const c10::optional<torch::Tensor>& correction_bias) {
// Check data type
TORCH_CHECK(
gating_output.scalar_type() == at::ScalarType::Float || gating_output.scalar_type() == at::ScalarType::Half ||
@@ -552,6 +642,23 @@ void topk_softmax(
torch::empty({workspace_size}, gating_output.options().dtype(at::ScalarType::Float));
const at::ScalarType dtype = gating_output.scalar_type();
// Validate correction_bias if provided - must always be float32
const float* bias_ptr = nullptr;
if (correction_bias.has_value()) {
const torch::Tensor& bias_tensor = correction_bias.value();
TORCH_CHECK(bias_tensor.dim() == 1, "correction_bias must be 1D tensor [num_experts]");
TORCH_CHECK(bias_tensor.size(0) == num_experts, "correction_bias size must match num_experts");
TORCH_CHECK(
bias_tensor.scalar_type() == at::ScalarType::Float,
"correction_bias must be float32, got ",
bias_tensor.scalar_type());
bias_ptr = bias_tensor.data_ptr<float>();
}
// Cast moe_softcapping from double to float for CUDA kernels
const float moe_softcapping_f = static_cast<float>(moe_softcapping);
if (dtype == at::ScalarType::Float) {
topkGatingSoftmaxKernelLauncher<float>(
gating_output.data_ptr<float>(),
@@ -562,6 +669,8 @@ void topk_softmax(
num_experts,
topk,
renormalize,
moe_softcapping_f,
bias_ptr,
stream);
} else if (dtype == at::ScalarType::Half) {
topkGatingSoftmaxKernelLauncher<__half>(
@@ -573,6 +682,8 @@ void topk_softmax(
num_experts,
topk,
renormalize,
moe_softcapping_f,
bias_ptr,
stream);
} else if (dtype == at::ScalarType::BFloat16) {
topkGatingSoftmaxKernelLauncher<__nv_bfloat16>(
@@ -584,6 +695,8 @@ void topk_softmax(
num_experts,
topk,
renormalize,
moe_softcapping_f,
bias_ptr,
stream);
} else {
TORCH_CHECK(false, "Unsupported gating_output dtype: ", dtype);