From e72cf136932a7d8790ef4feeaa606932b6d62c5c Mon Sep 17 00:00:00 2001 From: Roger Young <42564206+rogeryoungh@users.noreply.github.com> Date: Thu, 20 Nov 2025 00:24:37 +0800 Subject: [PATCH] Support moe topk sigmoid kernel (#13049) Co-authored-by: xuebi --- sgl-kernel/CMakeLists.txt | 1 + .../benchmark/bench_moe_topk_sigmoid.py | 171 +++++ sgl-kernel/csrc/common_extension.cc | 5 + sgl-kernel/csrc/common_extension_rocm.cc | 5 + .../csrc/moe/moe_topk_sigmoid_kernels.cu | 592 ++++++++++++++++++ sgl-kernel/include/sgl_kernel_ops.h | 7 + sgl-kernel/python/sgl_kernel/__init__.py | 1 + sgl-kernel/python/sgl_kernel/moe.py | 26 + sgl-kernel/setup_rocm.py | 1 + sgl-kernel/tests/test_moe_topk_sigmoid.py | 183 ++++++ 10 files changed, 992 insertions(+) create mode 100644 sgl-kernel/benchmark/bench_moe_topk_sigmoid.py create mode 100644 sgl-kernel/csrc/moe/moe_topk_sigmoid_kernels.cu create mode 100644 sgl-kernel/tests/test_moe_topk_sigmoid.py diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index 949402531..fcb158924 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/sgl-kernel/CMakeLists.txt @@ -323,6 +323,7 @@ set(SOURCES "csrc/moe/moe_sum.cu" "csrc/moe/moe_sum_reduce.cu" "csrc/moe/moe_topk_softmax_kernels.cu" + "csrc/moe/moe_topk_sigmoid_kernels.cu" "csrc/moe/nvfp4_blockwise_moe.cu" "csrc/moe/fp8_blockwise_moe_kernel.cu" "csrc/moe/prepare_moe_input.cu" diff --git a/sgl-kernel/benchmark/bench_moe_topk_sigmoid.py b/sgl-kernel/benchmark/bench_moe_topk_sigmoid.py new file mode 100644 index 000000000..d34e68b98 --- /dev/null +++ b/sgl-kernel/benchmark/bench_moe_topk_sigmoid.py @@ -0,0 +1,171 @@ +import itertools +import os + +import pytest +import torch +import triton +from sgl_kernel import topk_sigmoid + +# CI environment detection +IS_CI = ( + os.getenv("CI", "false").lower() == "true" + or os.getenv("GITHUB_ACTIONS", "false").lower() == "true" +) + + +def torch_topk_sigmoid_native( + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + correction_bias: torch.Tensor = None, +): + scores = gating_output.sigmoid() + if correction_bias is not None: + n_routed_experts = gating_output.shape[-1] + scores_for_choice = scores.view( + -1, n_routed_experts + ) + correction_bias.unsqueeze(0) + _, topk_indices = torch.topk(scores_for_choice, k=topk, dim=-1) + topk_weights = scores.gather(1, topk_indices) + else: + topk_weights, topk_indices = torch.topk(scores, k=topk, dim=-1) + + if renormalize: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + + return topk_weights, topk_indices + + +def sglang_topk_sigmoid( + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + correction_bias: torch.Tensor = None, +): + num_tokens, num_experts = gating_output.shape + + topk_weights = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda") + topk_indices = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda") + + topk_sigmoid( + topk_weights, + topk_indices, + gating_output, + renormalize=renormalize, + correction_bias=correction_bias, + ) + + return topk_weights, topk_indices + + +def get_topk_sigmoid_input(num_tokens, num_experts): + gating_output = torch.randn( + (num_tokens, num_experts), dtype=torch.float32, device="cuda" + ) + correction_bias = torch.randn((num_experts), dtype=torch.float32, device="cuda") + return gating_output, correction_bias + + +def calculate_diff(num_tokens, num_experts, topk): + gating_output, correction_bias = get_topk_sigmoid_input(num_tokens, num_experts) + + weights_torch, indices_torch = torch_topk_sigmoid_native( + gating_output.clone(), + topk, + True, + correction_bias.clone(), + ) + weights_sglang, indices_sglang = sglang_topk_sigmoid( + gating_output.clone(), + topk, + True, + correction_bias.clone(), + ) + + weights_diff = torch.abs(weights_torch - weights_sglang).mean().item() + indices_match = torch.equal(indices_torch, indices_sglang) + + if ( + torch.allclose(weights_torch, weights_sglang, atol=1e-3, rtol=1e-3) + and indices_match + ): + print("✅ Torch and SGLang topk_sigmoid implementations match") + else: + print( + f"❌ Implementations differ: Weights diff={weights_diff}, Indices match={indices_match}" + ) + + +# CI environment uses simplified parameters +if IS_CI: + num_tokens_range = [128] # Single value for CI + num_experts_range = [32] # Single value for CI + topk_range = [2] # Single value for CI +else: + num_tokens_range = [128, 512, 1024, 2048, 4096, 8192, 16384, 32768] + num_experts_range = [32, 64, 128, 256, 12, 512] + topk_range = [1, 2, 4, 8] + +configs = list(itertools.product(num_tokens_range, num_experts_range, topk_range)) + + +# Filter providers based on vLLM availability +line_vals = ["sglang", "torch"] +line_names = ["SGLang", "Torch"] +styles = [("blue", "-"), ("green", "-")] + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["num_tokens", "num_experts", "topk"], + x_vals=configs, + line_arg="provider", + line_vals=line_vals, + line_names=line_names, + styles=styles, + ylabel="Latency (us)", + plot_name="topk-sigmoid-performance", + args={}, + ) +) +def benchmark(num_tokens, num_experts, topk, provider): + gating_output, correction_bias = get_topk_sigmoid_input(num_tokens, num_experts) + + if provider == "torch" or provider == "torch1": + + def fn(): + return torch_topk_sigmoid_native( + gating_output, + topk, + True, + correction_bias, + ) + + elif provider == "sglang" or provider == "sglang1": + + def fn(): + return sglang_topk_sigmoid(gating_output, topk, True, correction_bias) + + quantiles = [0.5, 0.2, 0.8] + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles) + + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + +if __name__ == "__main__": + # Simplify configs for CI environment + if IS_CI: + test_configs = [(20, 32, 2)] # Single config for CI + else: + test_configs = [ + (20, 256, 4), + (20, 256, 8), + (20, 12, 4), + (20, 12, 1), + (20, 512, 4), + (20, 512, 1), + ] + + for num_tokens, num_experts, topk in test_configs: + calculate_diff(num_tokens, num_experts, topk) + benchmark.run(print_data=True) diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index a40de3e24..6ca57579d 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/sgl-kernel/csrc/common_extension.cc @@ -230,6 +230,11 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "moe_softcapping, Tensor? correction_bias) -> ()"); m.impl("topk_softmax", torch::kCUDA, &topk_softmax); + m.def( + "topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor gating_output, bool renormalize, Tensor? " + "correction_bias) -> ()"); + m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid); + m.def("moe_sum_reduce(Tensor input, Tensor output, float routed_scaling_factor) -> ()"); m.impl("moe_sum_reduce", torch::kCUDA, &moe_sum_reduce); diff --git a/sgl-kernel/csrc/common_extension_rocm.cc b/sgl-kernel/csrc/common_extension_rocm.cc index 2868153d5..9981e4670 100644 --- a/sgl-kernel/csrc/common_extension_rocm.cc +++ b/sgl-kernel/csrc/common_extension_rocm.cc @@ -100,6 +100,11 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) { "moe_softcapping, Tensor? correction_bias) -> ()"); m.impl("topk_softmax", torch::kCUDA, &topk_softmax); + m.def( + "topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor gating_output, bool renormalize, Tensor? " + "correction_bias) -> ()"); + m.impl("topk_sigmoid", torch::kCUDA, &topk_sigmoid); + /* * From csrc/speculative */ diff --git a/sgl-kernel/csrc/moe/moe_topk_sigmoid_kernels.cu b/sgl-kernel/csrc/moe/moe_topk_sigmoid_kernels.cu new file mode 100644 index 000000000..81c584b1f --- /dev/null +++ b/sgl-kernel/csrc/moe/moe_topk_sigmoid_kernels.cu @@ -0,0 +1,592 @@ +// Adapt from https://github.com/vllm-project/vllm/blob/v0.7.3/csrc/moe/topk_softmax_kernels.cu +// which is originally adapted from +// https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu +/* 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 +#include +#include + +#ifndef USE_ROCM +#include +#include +#include +#else +#include +#include +#endif + +#include "utils.h" + +#define MAX(a, b) ((a) > (b) ? (a) : (b)) +#define MIN(a, b) ((a) < (b) ? (a) : (b)) + +// Define reduction operators based on CUDA version +// CUDA 13 (12.9+) deprecated cub::Max/Min in favor of cuda::maximum/minimum +#if CUDA_VERSION >= 12090 +using MaxReduceOp = cuda::maximum<>; +using MinReduceOp = cuda::minimum<>; +#else +using MaxReduceOp = cub::Max; +using MinReduceOp = cub::Min; +#endif + +/// Aligned array type +template < + typename T, + /// Number of elements in the array + int N, + /// Alignment requirement in bytes + int Alignment = sizeof(T) * N> +class alignas(Alignment) AlignedArray { + T data[N]; +}; + +// ========================== Util functions to convert types ========================== +template +__device__ float convert_to_float(T x) { + if constexpr (std::is_same_v) { + return __half2float(x); + } else if constexpr (std::is_same_v) { + return __bfloat162float(x); + } else if constexpr (std::is_same_v) { + return x; + } else { + return static_cast(x); + } +} + +// ====================== Sigmoid things =============================== +// We have our own implementation of sigmoid here so we can support transposing the output +// in the sigmoid kernel when we extend this module to support expert-choice routing. +template +__launch_bounds__(TPB) __global__ void moeSigmoid( + const T* input, const bool* finished, float* output, const int num_cols, const float* correction_bias) { + const int thread_row_offset = blockIdx.x * num_cols; + + // Don't touch finished rows. + if ((finished != nullptr) && finished[blockIdx.x]) { + 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; + float val = convert_to_float(input[idx]); + + val = 1.0f / (1.0f + expf(-val)); + + // Apply correction bias if provided + if (correction_bias != nullptr) { + val = val + correction_bias[ii]; + } + + output[idx] = val; // Store transformed value + } +} + +template +__launch_bounds__(TPB) __global__ void moeTopK( + const float* inputs_after_sigmoid, + const bool* finished, + float* output, + int* indices, + const int num_experts, + const int k, + const int start_expert, + const int end_expert, + const bool renormalize, + const float* correction_bias) { + using cub_kvp = cub::KeyValuePair; + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tmpStorage; + + cub_kvp thread_kvp; + cub::ArgMax arg_max; + + const int block_row = blockIdx.x; + + const bool row_is_active = finished ? !finished[block_row] : true; + const int thread_read_offset = blockIdx.x * num_experts; + float row_sum_for_renormalize = 0; + for (int k_idx = 0; k_idx < k; ++k_idx) { + thread_kvp.key = 0; + thread_kvp.value = -1.f; // This is OK because inputs are probabilities + + cub_kvp inp_kvp; + for (int expert = threadIdx.x; expert < num_experts; expert += TPB) { + const int idx = thread_read_offset + expert; + inp_kvp.key = expert; + inp_kvp.value = inputs_after_sigmoid[idx]; + + for (int prior_k = 0; prior_k < k_idx; ++prior_k) { + const int prior_winning_expert = indices[k * block_row + prior_k]; + + if (prior_winning_expert == expert) { + inp_kvp = thread_kvp; + } + } + + thread_kvp = arg_max(inp_kvp, thread_kvp); + } + + const cub_kvp result_kvp = BlockReduce(tmpStorage).Reduce(thread_kvp, arg_max); + if (threadIdx.x == 0) { + // Ignore experts the node isn't responsible for with expert parallelism + const int expert = result_kvp.key; + const bool node_uses_expert = expert >= start_expert && expert < end_expert; + const bool should_process_row = row_is_active && node_uses_expert; + + const int idx = k * block_row + k_idx; + float val = result_kvp.value; + if (correction_bias != nullptr) { + val -= correction_bias[expert]; + } + output[idx] = val; + indices[idx] = should_process_row ? (expert - start_expert) : num_experts; + assert(indices[idx] >= 0); + row_sum_for_renormalize += val; + } + __syncthreads(); + } + + if (renormalize && threadIdx.x == 0) { + float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize; + for (int k_idx = 0; k_idx < k; ++k_idx) { + const int idx = k * block_row + k_idx; + output[idx] = output[idx] * row_sum_for_renormalize_inv; + } + } +} + +// ====================== TopK sigmoid things =============================== + +/* + A Top-K gating sigmoid written to exploit when the number of experts in the MoE layers + are a small power of 2. This allows us to cleanly share the rows among the threads in + a single warp and eliminate communication between warps (so no need to use shared mem). + + It fuses the sigmoid, max and argmax into a single kernel. + + Limitations: + 1) This implementation is intended for when the number of experts is a small power of 2. + 2) This implementation assumes k is small, but will work for any k. +*/ + +template +__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__ void topkGatingSigmoid( + const T* input, + const bool* finished, + float* output, + const int num_rows, + int* indices, + const int k, + const int start_expert, + const int end_expert, + const bool renormalize, + 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"); + static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG), "BYTES_PER_LDG must be power of 2"); + static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16"); + + // Number of bytes each thread pulls in per load + static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T); + static constexpr int ELTS_PER_ROW = NUM_EXPERTS; + static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT; + static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG; + + // Restrictions based on previous section. + static_assert(VPT % ELTS_PER_LDG == 0, "The elements per thread must be a multiple of the elements per ldg"); + static_assert(WARP_SIZE % THREADS_PER_ROW == 0, "The threads per row must cleanly divide the threads per warp"); + static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW), "THREADS_PER_ROW must be power of 2"); + static_assert(THREADS_PER_ROW <= WARP_SIZE, "THREADS_PER_ROW can be at most warp size"); + + // We have NUM_EXPERTS elements per row. We specialize for small #experts + static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT; + static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW; + static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP; + + // Restrictions for previous section. + static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0, "The elts per row must cleanly divide the total elt per warp"); + + // ===================== From this point, we finally start computing run-time variables. ======================== + + // Compute CTA and warp rows. We pack multiple rows into a single warp, and a block contains WARPS_PER_CTA warps. + // This, each block processes a chunk of rows. We start by computing the start row for each block. + const int cta_base_row = blockIdx.x * ROWS_PER_CTA; + + // Now, using the base row per thread block, we compute the base row per warp. + const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP; + + // The threads in a warp are split into sub-groups that will work on a row. + // We compute row offset for each thread sub-group + const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW; + const int thread_row = warp_base_row + thread_row_in_warp; + + // Threads with indices out of bounds should early exit here. + if (thread_row >= num_rows) { + return; + } + const bool row_is_active = finished ? !finished[thread_row] : true; + + // We finally start setting up the read pointers for each thread. First, each thread jumps to the start of the + // row it will read. + const T* thread_row_ptr = input + thread_row * ELTS_PER_ROW; + + // Now, we compute the group each thread belong to in order to determine the first column to start loads. + const int thread_group_idx = threadIdx.x % THREADS_PER_ROW; + const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG; + const T* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread; + + // Determine the pointer type to use to read in the data depending on the BYTES_PER_LDG template param. In theory, + // this can support all powers of 2 up to 16. + // NOTE(woosuk): The original implementation uses CUTLASS aligned array here. + // We defined our own aligned array and use it here to avoid the dependency on CUTLASS. + using AccessType = AlignedArray; + + // Finally, we pull in the data from global mem + T row_chunk_temp[VPT]; + AccessType* row_chunk_vec_ptr = reinterpret_cast(&row_chunk_temp); + const AccessType* vec_thread_read_ptr = reinterpret_cast(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) { + float val = convert_to_float(row_chunk_temp[ii]); + val = 1.0f / (1.0f + expf(-val)); + // 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; + } + + // Now, row_chunk contains the sigmoid of the row chunk. Now, I want to find the topk elements in each row, along + // with the max index. + int start_col = first_elt_read_by_thread; + static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW; + + float row_sum_for_renormalize = 0; + + for (int k_idx = 0; k_idx < k; ++k_idx) { + // First, each thread does the local argmax + float max_val = row_chunk[0]; + int expert = start_col; +#pragma unroll + for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD; ++ldg, col += COLS_PER_GROUP_LDG) { +#pragma unroll + for (int ii = 0; ii < ELTS_PER_LDG; ++ii) { + float val = row_chunk[ldg * ELTS_PER_LDG + ii]; + + // No check on the experts here since columns with the smallest index are processed first and only + // updated if > (not >=) + if (val > max_val) { + max_val = val; + expert = col + ii; + } + } + } + +// Now, we perform the argmax reduce. We use the butterfly pattern so threads reach consensus about the max. +// This will be useful for K > 1 so that the threads can agree on "who" had the max value. That thread can +// then blank out their max with -inf and the warp can run more iterations... +#pragma unroll + for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { + float other_max = SGLANG_SHFL_XOR_SYNC_WIDTH(0xffffffff, max_val, mask, THREADS_PER_ROW); + int other_expert = SGLANG_SHFL_XOR_SYNC_WIDTH(0xffffffff, expert, mask, THREADS_PER_ROW); + + // We want lower indices to "win" in every thread so we break ties this way + if (other_max > max_val || (other_max == max_val && other_expert < expert)) { + max_val = other_max; + expert = other_expert; + } + } + + // Write the max for this k iteration to global memory. + if (thread_group_idx == 0) { + // Add a guard to ignore experts not included by this node + const bool node_uses_expert = expert >= start_expert && expert < end_expert; + const bool should_process_row = row_is_active && node_uses_expert; + + // The lead thread from each sub-group will write out the final results to global memory. (This will be a + // single) thread per row of the input/output matrices. + const int idx = k * thread_row + k_idx; + if (correction_bias != nullptr) { + max_val -= correction_bias[expert]; + } + output[idx] = max_val; + indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS; + row_sum_for_renormalize += max_val; + } + + // Finally, we clear the value in the thread with the current max if there is another iteration to run. + if (k_idx + 1 < k) { + const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG; + const int thread_to_clear_in_group = (expert / ELTS_PER_LDG) % THREADS_PER_ROW; + + // Only the thread in the group which produced the max will reset the "winning" value to -inf. + if (thread_group_idx == thread_to_clear_in_group) { + const int offset_for_expert = expert % ELTS_PER_LDG; + // Safe to set to any negative value since row_chunk values must be between 0 and 1. + row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] = -10000.f; + } + } + } + + // Fuse renormalization of topk_weights into this kernel + if (renormalize && thread_group_idx == 0) { + float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize; +#pragma unroll + for (int k_idx = 0; k_idx < k; ++k_idx) { + const int idx = k * thread_row + k_idx; + output[idx] = output[idx] * row_sum_for_renormalize_inv; + } + } +} + +namespace detail { +// Constructs some constants needed to partition the work across threads at compile time. +template +struct TopkConstants { + static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T); + static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE) == 0 || EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0, ""); + static constexpr int VECs_PER_THREAD = MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE)); + static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG; + static constexpr int THREADS_PER_ROW = EXPERTS / VPT; + static constexpr int ROWS_PER_WARP = WARP_SIZE / THREADS_PER_ROW; +}; +} // namespace detail + +template +void topkGatingSigmoidLauncherHelper( + const T* input, + const bool* finished, + float* output, + int* indices, + const int num_rows, + const int k, + const int start_expert, + const int end_expert, + const bool renormalize, + const float* correction_bias, + cudaStream_t stream) { + static constexpr std::size_t MAX_BYTES_PER_LDG = 16; + + static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(T) * EXPERTS); + using Constants = detail::TopkConstants; + static constexpr int VPT = Constants::VPT; + static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP; + const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP; + const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB; + + dim3 block_dim(WARP_SIZE, WARPS_PER_TB); + topkGatingSigmoid<<>>( + input, finished, output, num_rows, indices, k, start_expert, end_expert, renormalize, correction_bias); +} + +#define LAUNCH_SIGMOID(TYPE, NUM_EXPERTS, WARPS_PER_TB) \ + topkGatingSigmoidLauncherHelper( \ + gating_output, \ + nullptr, \ + topk_weights, \ + topk_indices, \ + num_tokens, \ + topk, \ + 0, \ + num_experts, \ + renormalize, \ + correction_bias, \ + stream); + +template +void topkGatingSigmoidKernelLauncher( + const T* gating_output, + float* topk_weights, + int* topk_indices, + float* sigmoid_workspace, + const int num_tokens, + const int num_experts, + const int topk, + const bool renormalize, + const float* correction_bias, + cudaStream_t stream) { + static constexpr int WARPS_PER_TB = 4; + switch (num_experts) { + case 1: + LAUNCH_SIGMOID(T, 1, WARPS_PER_TB); + break; + case 2: + LAUNCH_SIGMOID(T, 2, WARPS_PER_TB); + break; + case 4: + LAUNCH_SIGMOID(T, 4, WARPS_PER_TB); + break; + case 8: + LAUNCH_SIGMOID(T, 8, WARPS_PER_TB); + break; + case 16: + LAUNCH_SIGMOID(T, 16, WARPS_PER_TB); + break; + case 32: + LAUNCH_SIGMOID(T, 32, WARPS_PER_TB); + break; + case 64: + LAUNCH_SIGMOID(T, 64, WARPS_PER_TB); + break; + case 128: + LAUNCH_SIGMOID(T, 128, WARPS_PER_TB); + break; + case 256: + LAUNCH_SIGMOID(T, 256, WARPS_PER_TB); + break; + default: { + TORCH_CHECK( + sigmoid_workspace != nullptr, + "sigmoid_workspace must be provided for num_experts that are not a power of 2."); + static constexpr int TPB = 256; + moeSigmoid + <<>>(gating_output, nullptr, sigmoid_workspace, num_experts, correction_bias); + moeTopK<<>>( + sigmoid_workspace, + nullptr, + topk_weights, + topk_indices, + num_experts, + topk, + 0, + num_experts, + renormalize, + correction_bias); + } + } +} + +void topk_sigmoid( + 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 c10::optional& correction_bias) { + // Check data type + TORCH_CHECK( + gating_output.scalar_type() == at::ScalarType::Float || gating_output.scalar_type() == at::ScalarType::Half || + gating_output.scalar_type() == at::ScalarType::BFloat16, + "gating_output must be float32, float16, or bfloat16"); + + // Check dimensions + TORCH_CHECK(gating_output.dim() == 2, "gating_output must be 2D tensor [num_tokens, num_experts]"); + TORCH_CHECK(topk_weights.dim() == 2, "topk_weights must be 2D tensor [num_tokens, topk]"); + TORCH_CHECK(topk_indices.dim() == 2, "topk_indices must be 2D tensor [num_tokens, topk]"); + + // Check shapes + TORCH_CHECK( + gating_output.size(0) == topk_weights.size(0), + "First dimension of topk_weights must match num_tokens in gating_output"); + TORCH_CHECK( + gating_output.size(0) == topk_indices.size(0), + "First dimension of topk_indices must match num_tokens in gating_output"); + TORCH_CHECK( + topk_weights.size(-1) == topk_indices.size(-1), + "Second dimension of topk_indices must match topk in topk_weights"); + TORCH_CHECK(topk_weights.size(-1) <= gating_output.size(-1), "topk must be less than or equal to num_experts"); + + const int num_experts = static_cast(gating_output.size(-1)); + const int num_tokens = static_cast(gating_output.size(0)); + const int topk = static_cast(topk_weights.size(-1)); + + const bool is_pow_2 = (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0); + const bool needs_workspace = !is_pow_2 || num_experts > 256; + const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0; + + const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output)); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + torch::Tensor sigmoid_workspace = + 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(); + } + + if (dtype == at::ScalarType::Float) { + topkGatingSigmoidKernelLauncher( + gating_output.data_ptr(), + topk_weights.data_ptr(), + topk_indices.data_ptr(), + sigmoid_workspace.data_ptr(), + num_tokens, + num_experts, + topk, + renormalize, + bias_ptr, + stream); + } else if (dtype == at::ScalarType::Half) { + topkGatingSigmoidKernelLauncher<__half>( + reinterpret_cast(gating_output.data_ptr()), + topk_weights.data_ptr(), + topk_indices.data_ptr(), + sigmoid_workspace.data_ptr(), + num_tokens, + num_experts, + topk, + renormalize, + bias_ptr, + stream); + } else if (dtype == at::ScalarType::BFloat16) { + topkGatingSigmoidKernelLauncher<__nv_bfloat16>( + reinterpret_cast(gating_output.data_ptr()), + topk_weights.data_ptr(), + topk_indices.data_ptr(), + sigmoid_workspace.data_ptr(), + num_tokens, + num_experts, + topk, + renormalize, + bias_ptr, + stream); + } else { + TORCH_CHECK(false, "Unsupported gating_output dtype: ", dtype); + } +} diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index c7b6388c6..953682902 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -317,6 +317,13 @@ void topk_softmax( double moe_softcapping, const c10::optional& correction_bias); +void topk_sigmoid( + torch::Tensor& topk_weights, + torch::Tensor& topk_indices, + torch::Tensor& gating_output, + bool renormalize, + const c10::optional& correction_bias); + void moe_sum_reduce(at::Tensor& input, at::Tensor& output, double routed_scaling_factor); void moe_sum(torch::Tensor& input, torch::Tensor& output); diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py index ee1d1c5cd..0c056967b 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/sgl-kernel/python/sgl_kernel/__init__.py @@ -91,6 +91,7 @@ from sgl_kernel.moe import ( moe_sum, moe_sum_reduce, prepare_moe_input, + topk_sigmoid, topk_softmax, ) from sgl_kernel.quantization import ( diff --git a/sgl-kernel/python/sgl_kernel/moe.py b/sgl-kernel/python/sgl_kernel/moe.py index 4f849451b..e14eebcb2 100755 --- a/sgl-kernel/python/sgl_kernel/moe.py +++ b/sgl-kernel/python/sgl_kernel/moe.py @@ -54,6 +54,32 @@ def topk_softmax( ) +def topk_sigmoid( + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + gating_output: torch.Tensor, + renormalize: bool = False, + correction_bias: Optional[torch.Tensor] = None, +) -> None: + """ + Compute top-k sigmoid for MoE routing. + + Args: + topk_weights: Output tensor for top-k weights [num_tokens, topk] + topk_ids: Output tensor for top-k expert indices [num_tokens, topk] + gating_output: Gating logits [num_tokens, num_experts] + renormalize: Whether to renormalize the top-k weights + correction_bias: Per-expert bias correction [num_experts], must be float32 if provided + """ + torch.ops.sgl_kernel.topk_sigmoid.default( + topk_weights, + topk_ids, + gating_output, + renormalize, + correction_bias, + ) + + def moe_sum_reduce( input_tensor, output_tensor, diff --git a/sgl-kernel/setup_rocm.py b/sgl-kernel/setup_rocm.py index 6e3466ec3..8cdb7c695 100644 --- a/sgl-kernel/setup_rocm.py +++ b/sgl-kernel/setup_rocm.py @@ -48,6 +48,7 @@ sources = [ "csrc/grammar/apply_token_bitmask_inplace_cuda.cu", "csrc/moe/moe_align_kernel.cu", "csrc/moe/moe_topk_softmax_kernels.cu", + "csrc/moe/moe_topk_sigmoid_kernels.cu", "csrc/speculative/eagle_utils.cu", "csrc/kvcacheio/transfer.cu", ] diff --git a/sgl-kernel/tests/test_moe_topk_sigmoid.py b/sgl-kernel/tests/test_moe_topk_sigmoid.py new file mode 100644 index 000000000..45b8222a9 --- /dev/null +++ b/sgl-kernel/tests/test_moe_topk_sigmoid.py @@ -0,0 +1,183 @@ +import itertools + +import pytest +import torch +from sgl_kernel import topk_sigmoid + + +@pytest.mark.parametrize( + "num_tokens, num_experts, topk", + list( + itertools.product( + [1, 16, 128, 512, 1024, 2048], # num_tokens + [4, 8, 16, 32, 64, 128, 256], # num_experts + [1, 2, 4], # topk + ) + ), +) +def test_topk_sigmoid(num_tokens, num_experts, topk): + gating_output = torch.randn( + (num_tokens, num_experts), dtype=torch.float32, device="cuda" + ) + + topk_weights = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda") + topk_indices = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda") + + topk_sigmoid( + topk_weights, + topk_indices, + gating_output, + ) + + # Native torch implementation + sigmoid_output = torch.sigmoid(gating_output) + topk_weights_ref, topk_indices_ref = torch.topk(sigmoid_output, topk, dim=-1) + + # Verify the top-k weights and indices match the torch native ones + assert torch.allclose( + topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 + ), f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}" + + assert torch.allclose( + topk_indices_ref.int(), topk_indices, atol=0, rtol=0 + ), f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}" + + +@pytest.mark.parametrize( + "num_tokens, num_experts, topk, dtype", + list( + itertools.product( + [1, 16, 128, 512, 1024, 2048], # num_tokens + [4, 8, 16, 32, 64, 128, 256], # num_experts + [1, 2, 4], # topk + [torch.float16, torch.bfloat16, torch.float32], # dtype + ) + ), +) +def test_topk_sigmoid_dtype_regression(num_tokens, num_experts, topk, dtype): + gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda") + + topk_weights = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda") + topk_indices = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda") + + topk_sigmoid( + topk_weights, + topk_indices, + gating_output, + ) + + topk_weights_ref = torch.empty( + (num_tokens, topk), dtype=torch.float32, device="cuda" + ) + topk_indices_ref = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda") + + topk_sigmoid( + topk_weights_ref, + topk_indices_ref, + gating_output.float(), + ) + + assert torch.allclose( + topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 + ), f"Weights mismatch: SGLang old interface={topk_weights_ref} vs SGLang new interface={topk_weights}" + + assert torch.allclose( + topk_indices_ref.int(), topk_indices, atol=0, rtol=0 + ), f"Indices mismatch: SGLang old interface={topk_indices_ref}, SGLang new interface={topk_indices}" + + +@pytest.mark.parametrize( + "num_tokens, num_experts, topk", + list( + itertools.product( + [1, 16, 128, 512, 1024, 2048], # num_tokens + [4, 8, 16, 32, 64, 128, 256], # num_experts + [1, 2, 4], # topk + ) + ), +) +def test_topk_sigmoid_renormalize(num_tokens, num_experts, topk): + gating_output = torch.randn( + (num_tokens, num_experts), dtype=torch.bfloat16, device="cuda" + ) + + topk_weights = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda") + topk_indices = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda") + + topk_sigmoid( + topk_weights, + topk_indices, + gating_output, + renormalize=True, + ) + + topk_weights_ref = torch.empty( + (num_tokens, topk), dtype=torch.float32, device="cuda" + ) + topk_indices_ref = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda") + token_expert_indices_ref = torch.empty( + (num_tokens, topk), dtype=torch.int32, device="cuda" + ) + + topk_sigmoid( + topk_weights_ref, + topk_indices_ref, + gating_output, + ) + topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True) + + assert torch.allclose( + topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 + ), f"Weights mismatch: SGLang w/o fused renormalize={topk_weights_ref} vs SGLang w/ fused renormalize={topk_weights}" + + assert torch.allclose( + topk_indices_ref.int(), topk_indices, atol=0, rtol=0 + ), f"Indices mismatch: SGLang w/o fused renormalize={topk_indices_ref}, SGLang w/ fused renormalize={topk_indices}" + + +@pytest.mark.parametrize( + "num_tokens, num_experts, topk", + list( + itertools.product( + [1, 16, 128, 512, 1024, 2048], # num_tokens + [4, 8, 16, 32, 48, 64, 128, 256], # num_experts + [1, 2, 4], # topk + ) + ), +) +def test_topk_sigmoid_renormalize_correction_bias(num_tokens, num_experts, topk): + gating_output = torch.randn( + (num_tokens, num_experts), dtype=torch.float32, device="cuda" + ) + correction_bias = torch.randn((num_experts), dtype=torch.float32, device="cuda") + + topk_weights = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda") + topk_indices = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda") + + topk_sigmoid( + topk_weights, + topk_indices, + gating_output, + renormalize=True, + correction_bias=correction_bias, + ) + + # Native torch implementation + sigmoid_output = torch.sigmoid(gating_output) + sigmoid_scores = sigmoid_output.view(-1, num_experts) + correction_bias.unsqueeze(0) + _, topk_indices_ref = torch.topk(sigmoid_scores, k=topk, dim=-1) + topk_weights_ref = sigmoid_output.gather(1, topk_indices_ref) + topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True) + + # Verify the top-k weights and indices match the torch native ones + assert torch.allclose( + topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3 + ), f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}" + + assert torch.allclose( + topk_indices_ref.int(), topk_indices, atol=0, rtol=0 + ), f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}" + + +if __name__ == "__main__": + pytest.main([__file__])