From 56a759cffc2ef381cda9202634f55d3c3ebc4ad0 Mon Sep 17 00:00:00 2001 From: Chenzhou Li <67032602+VoidC-minor@users.noreply.github.com> Date: Sun, 16 Aug 2026 15:02:57 +0800 Subject: [PATCH] [JIT Kernel] Migrate moe_topk_softmax from AOT to JIT (#34509) Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> --- .../kernels/jit/csrc/moe/moe_topk_softmax.cuh | 685 ++++++++++++++++++ python/sglang/kernels/ops/moe/__init__.py | 19 +- .../kernels/ops/moe/moe_topk_softmax.py | 108 +++ .../benchmark/moe/bench_moe_topk_softmax.py | 65 ++ .../kernels/ops/moe/test_moe_topk_softmax.py | 200 +++++ 5 files changed, 1075 insertions(+), 2 deletions(-) create mode 100644 python/sglang/kernels/jit/csrc/moe/moe_topk_softmax.cuh create mode 100644 python/sglang/kernels/ops/moe/moe_topk_softmax.py create mode 100644 test/registered/kernels/benchmark/moe/bench_moe_topk_softmax.py create mode 100644 test/registered/kernels/ops/moe/test_moe_topk_softmax.py diff --git a/python/sglang/kernels/jit/csrc/moe/moe_topk_softmax.cuh b/python/sglang/kernels/jit/csrc/moe/moe_topk_softmax.cuh new file mode 100644 index 000000000..59bb78df5 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/moe/moe_topk_softmax.cuh @@ -0,0 +1,685 @@ +// 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 +// +// JIT port: the CUDA kernels are unchanged from the sgl-kernel implementation; only +// the host launcher is adapted to the tvm-ffi TensorView API. The softmax workspace +// (needed for non-power-of-two / >256 expert counts) is allocated by the Python +// wrapper and passed in. +#pragma once + +#include // For TensorMatcher, SymbolicSize, SymbolicDevice +#include // For RuntimeCheck + +#include // For LaunchKernel, fp16_t/bf16_t/fp32_t + +#include +#include +#include +#include +#include + +// CUDA 12.9+ deprecated cub::Max/Min in favour of cuda::maximum/minimum +#if CUDA_VERSION >= 12090 +#include +#endif + +#include +#include +#include + +namespace sglang { + +// Namespaced so the macro cannot collide with a WARP_SIZE defined by another +// translation unit pulled into the same JIT module. +#define MOE_TOPK_SOFTMAX_WARP_SIZE 32 +static constexpr int WARP_SIZE = MOE_TOPK_SOFTMAX_WARP_SIZE; + +#define SGL_MAX(a, b) ((a) > (b) ? (a) : (b)) +#define SGL_MIN(a, b) ((a) < (b) ? (a) : (b)) + +#define SGLANG_SHFL_XOR_SYNC_WIDTH(mask, var, lane_mask, width) __shfl_xor_sync(mask, var, lane_mask, width) + +// 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 + +using cub_kvp = cub::KeyValuePair; + +/// Aligned array type +template +class alignas(Alignment) AlignedArray { + T data[N]; +}; + +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); + } +} + +// ====================== Softmax things =============================== +template +__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; + __shared__ typename BlockReduce::TempStorage tmpStorage; + + __shared__ float normalizing_factor; + __shared__ float float_max; + + const int thread_row_offset = blockIdx.x * num_cols; + + float threadData(-FLT_MAX); + + if ((finished != nullptr) && finished[blockIdx.x]) { + return; + } + + for (int ii = threadIdx.x; ii < num_cols; ii += TPB) { + const int idx = thread_row_offset + ii; + float val = convert_to_float(input[idx]); + + if (moe_softcapping != 0.0f) { + val = tanhf(val / moe_softcapping) * moe_softcapping; + } + + if (correction_bias != nullptr) { + val = val + correction_bias[ii]; + } + + output[idx] = val; + threadData = max(val, threadData); + } + + const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, MaxReduceOp()); + + if (threadIdx.x == 0) { + float_max = maxElem; + } + __syncthreads(); + + threadData = 0; + for (int ii = threadIdx.x; ii < num_cols; ii += TPB) { + const int idx = thread_row_offset + ii; + threadData += exp((output[idx] - float_max)); + } + + const auto Z = BlockReduce(tmpStorage).Sum(threadData); + + if (threadIdx.x == 0) { + normalizing_factor = 1.f / Z; + } + __syncthreads(); + + for (int ii = threadIdx.x; ii < num_cols; ii += TPB) { + const int idx = thread_row_offset + ii; + const float softmax_val = exp((output[idx] - float_max)) * normalizing_factor; + output[idx] = softmax_val; + } +} + +namespace moe { +struct TopKPair { + static const int PAIR = 2; + static const int MAX_INDEX = 0; + cub_kvp max; + cub_kvp secondMax; + + __device__ TopKPair() {} + __device__ TopKPair(cub_kvp max, cub_kvp secondMax) : max(max), secondMax(secondMax) {} +}; + +struct TopKPairArgMax { + __device__ TopKPairArgMax() {} + __device__ __forceinline__ TopKPair operator()(const TopKPair& candidate1, const TopKPair& candidate2) const { + cub_kvp globalMax, globalSecondMax; + + if (candidate1.max.value > candidate2.max.value) { + globalMax = candidate1.max; + } else { + globalMax = candidate2.max; + } + + if (globalMax.key == candidate1.max.key) { + globalSecondMax = (candidate1.secondMax.value > candidate2.max.value) ? candidate1.secondMax : candidate2.max; + } else { + globalSecondMax = (candidate2.secondMax.value > candidate1.max.value) ? candidate2.secondMax : candidate1.max; + } + return TopKPair(globalMax, globalSecondMax); + } +}; +} // namespace moe + +template +__launch_bounds__(TPB) __global__ void moeTopKFast( + float* inputs_after_softmax, + 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) { + using namespace moe; + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tmpStorage; + TopKPair thread_pair; + + 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 + TopKPair::PAIR - 1) / TopKPair::PAIR; ++k_idx) { + thread_pair.max.key = 0; + thread_pair.max.value = -1.f; + thread_pair.secondMax.key = 0; + thread_pair.secondMax.value = -1.f; + + 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_softmax[idx]; + if (inp_kvp.value > thread_pair.max.value) { + thread_pair.secondMax = thread_pair.max; + thread_pair.max = inp_kvp; + } else if (inp_kvp.value > thread_pair.secondMax.value) { + thread_pair.secondMax = inp_kvp; + } + } + + TopKPairArgMax reducer; + const TopKPair result_pair = BlockReduce(tmpStorage).Reduce(thread_pair, reducer); + if (threadIdx.x == 0) { +#pragma unroll + for (int i = 0; i < TopKPair::PAIR; i++) { + if (k_idx * 2 + i >= k) break; + cub_kvp result = (i == TopKPair::MAX_INDEX) ? result_pair.max : result_pair.secondMax; + int expert = result.key; + bool node_uses_expert = expert >= start_expert && expert < end_expert; + bool should_process_row = row_is_active && node_uses_expert; + inputs_after_softmax[thread_read_offset + expert] = -1.f; + int idx = k * block_row + k_idx * 2 + i; + output[idx] = result.value; + indices[idx] = should_process_row ? (expert - start_expert) : num_experts; + assert(indices[idx] >= 0); + row_sum_for_renormalize += result.value; + } + } + __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; + } + } +} + +template +__launch_bounds__(TPB) __global__ void moeTopK( + float* inputs_after_softmax, + 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) { + 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; + + 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_softmax[idx]; + thread_kvp = arg_max(inp_kvp, thread_kvp); + } + + const cub_kvp result_kvp = BlockReduce(tmpStorage).Reduce(thread_kvp, arg_max); + if (threadIdx.x == 0) { + 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; + output[idx] = result_kvp.value; + indices[idx] = should_process_row ? (expert - start_expert) : num_experts; + assert(indices[idx] >= 0); + row_sum_for_renormalize += result_kvp.value; + inputs_after_softmax[thread_read_offset + expert] = -1.f; + } + __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 softmax things =============================== +template +__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__ void topkGatingSoftmax( + 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 moe_softcapping, + const float* correction_bias) { + 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"); + + 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; + + 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"); + + 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; + + static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0, "The elts per row must cleanly divide the total elt per warp"); + + const int cta_base_row = blockIdx.x * ROWS_PER_CTA; + const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP; + const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW; + const int thread_row = warp_base_row + thread_row_in_warp; + + if (thread_row >= num_rows) { + return; + } + const bool row_is_active = finished ? !finished[thread_row] : true; + + const T* thread_row_ptr = input + thread_row * ELTS_PER_ROW; + 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; + + using AccessType = AlignedArray; + + 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 + 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 + for (int ii = 0; ii < VPT; ++ii) { + row_chunk[ii] = convert_to_float(row_chunk_temp[ii]); + } + + if (moe_softcapping != 0.0f || correction_bias != nullptr) { +#pragma unroll + for (int ii = 0; ii < VPT; ++ii) { + float val = row_chunk[ii]; + + if (moe_softcapping != 0.0f) { + val = tanhf(val / moe_softcapping) * moe_softcapping; + } + + if (correction_bias != nullptr) { + 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; + } + } + + float thread_max = row_chunk[0]; +#pragma unroll + for (int ii = 1; ii < VPT; ++ii) { + thread_max = max(thread_max, row_chunk[ii]); + } + +#pragma unroll + for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { + thread_max = max(thread_max, SGLANG_SHFL_XOR_SYNC_WIDTH(0xffffffff, thread_max, mask, THREADS_PER_ROW)); + } + + float row_sum = 0; +#pragma unroll + for (int ii = 0; ii < VPT; ++ii) { + row_chunk[ii] = expf(row_chunk[ii] - thread_max); + row_sum += row_chunk[ii]; + } + +#pragma unroll + for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { + row_sum += SGLANG_SHFL_XOR_SYNC_WIDTH(0xffffffff, row_sum, mask, THREADS_PER_ROW); + } + + const float reciprocal_row_sum = 1.f / row_sum; + +#pragma unroll + for (int ii = 0; ii < VPT; ++ii) { + row_chunk[ii] = row_chunk[ii] * reciprocal_row_sum; + } + + 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) { + 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]; + + if (val > max_val) { + max_val = val; + expert = col + ii; + } + } + } + +#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); + + if (other_max > max_val || (other_max == max_val && other_expert < expert)) { + max_val = other_max; + expert = other_expert; + } + } + + if (thread_group_idx == 0) { + 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 * thread_row + k_idx; + output[idx] = max_val; + indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS; + row_sum_for_renormalize += max_val; + } + + 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; + + if (thread_group_idx == thread_to_clear_in_group) { + const int offset_for_expert = expert % ELTS_PER_LDG; + row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] = -10000.f; + } + } + } + + 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 { +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 = SGL_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 topkGatingSoftmaxLauncherHelper( + 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 moe_softcapping, + const float* correction_bias, + cudaStream_t stream) { + static constexpr std::size_t MAX_BYTES_PER_LDG = 16; + + static constexpr int BYTES_PER_LDG = SGL_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); + host::LaunchKernel(dim3(num_blocks), block_dim, stream)( + topkGatingSoftmax, + 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( \ + gating_output, \ + nullptr, \ + topk_weights, \ + topk_indices, \ + num_tokens, \ + topk, \ + 0, \ + num_experts, \ + renormalize, \ + moe_softcapping, \ + correction_bias, \ + stream); + +template +void topkGatingSoftmaxKernelLauncher( + const T* gating_output, + float* topk_weights, + int* topk_indices, + float* softmax_workspace, + const int num_tokens, + 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) { + case 1: + LAUNCH_SOFTMAX(T, 1, WARPS_PER_TB); + break; + case 2: + LAUNCH_SOFTMAX(T, 2, WARPS_PER_TB); + break; + case 4: + LAUNCH_SOFTMAX(T, 4, WARPS_PER_TB); + break; + case 8: + LAUNCH_SOFTMAX(T, 8, WARPS_PER_TB); + break; + case 16: + LAUNCH_SOFTMAX(T, 16, WARPS_PER_TB); + break; + case 32: + LAUNCH_SOFTMAX(T, 32, WARPS_PER_TB); + break; + case 64: + LAUNCH_SOFTMAX(T, 64, WARPS_PER_TB); + break; + case 128: + LAUNCH_SOFTMAX(T, 128, WARPS_PER_TB); + break; + case 256: + LAUNCH_SOFTMAX(T, 256, WARPS_PER_TB); + break; + case 512: + LAUNCH_SOFTMAX(T, 512, WARPS_PER_TB); + break; + default: { + host::RuntimeCheck( + softmax_workspace != nullptr, + "softmax_workspace must be provided for num_experts that are not a power of 2."); + static constexpr int TPB = 256; + host::LaunchKernel(dim3(num_tokens), dim3(TPB), stream)( + moeSoftmax, + gating_output, + (const bool*)nullptr, + softmax_workspace, + num_experts, + moe_softcapping, + correction_bias); + if (topk == 1) { + host::LaunchKernel(dim3(num_tokens), dim3(TPB), stream)( + moeTopK, + softmax_workspace, + (const bool*)nullptr, + topk_weights, + topk_indices, + num_experts, + topk, + 0, + num_experts, + renormalize); + } else { + host::LaunchKernel(dim3(num_tokens), dim3(TPB), stream)( + moeTopKFast, + softmax_workspace, + (const bool*)nullptr, + topk_weights, + topk_indices, + num_experts, + topk, + 0, + num_experts, + renormalize); + } + } + } +} + +template +void topk_softmax( + tvm::ffi::TensorView topk_weights, // [num_tokens, topk] (float32, out) + tvm::ffi::TensorView topk_indices, // [num_tokens, topk] (int32, out) + tvm::ffi::TensorView gating_output, // [num_tokens, num_experts] (T, in) + tvm::ffi::TensorView softmax_workspace, // [num_tokens * num_experts] or [0] (float32) + bool renormalize, + double moe_softcapping, + tvm::ffi::Optional correction_bias) { + using namespace host; + + auto device = SymbolicDevice{}; + auto N = SymbolicSize{"num_tokens"}; + auto E = SymbolicSize{"num_experts"}; + auto K = SymbolicSize{"topk"}; + device.set_options(); + + TensorMatcher({N, E}).with_dtype().with_device(device).verify(gating_output); + TensorMatcher({N, K}).with_dtype().with_device(device).verify(topk_weights); + TensorMatcher({N, K}).with_dtype().with_device(device).verify(topk_indices); + + const int num_tokens = static_cast(N.unwrap()); + const int num_experts = static_cast(E.unwrap()); + const int topk = static_cast(K.unwrap()); + + if (num_tokens == 0) return; + + const float* bias_ptr = nullptr; + if (correction_bias.has_value()) { + const auto& bias = correction_bias.value(); + RuntimeCheck(bias.dtype().code == kDLFloat && bias.dtype().bits == 32, "correction_bias must be float32"); + bias_ptr = static_cast(bias.data_ptr()); + } + + float* workspace_ptr = (softmax_workspace.numel() > 0) ? static_cast(softmax_workspace.data_ptr()) : nullptr; + + const cudaStream_t stream = LaunchKernel::resolve_device(device.unwrap()); + const float moe_softcapping_f = static_cast(moe_softcapping); + + topkGatingSoftmaxKernelLauncher( + static_cast(gating_output.data_ptr()), + static_cast(topk_weights.data_ptr()), + static_cast(topk_indices.data_ptr()), + workspace_ptr, + num_tokens, + num_experts, + topk, + renormalize, + moe_softcapping_f, + bias_ptr, + stream); +} + +} // namespace sglang diff --git a/python/sglang/kernels/ops/moe/__init__.py b/python/sglang/kernels/ops/moe/__init__.py index 00359e8fd..ad7a62ad2 100644 --- a/python/sglang/kernels/ops/moe/__init__.py +++ b/python/sglang/kernels/ops/moe/__init__.py @@ -17,6 +17,7 @@ if TYPE_CHECKING: import torch _CUDA = frozenset({CapabilityRequirement.CUDA}) +_HIP = frozenset({CapabilityRequirement.HIP}) register_kernel( KernelSpec( @@ -48,11 +49,25 @@ register_kernel( op="moe.topk_softmax", backend=KernelBackend.AOT, target="sgl_kernel:topk_softmax", + capabilities=_HIP, format_signature=FormatSignature( in_place=True, description="top-k softmax routing weights/ids", ), - description="MoE top-k softmax (sgl_kernel wheel).", + description="MoE top-k softmax (sgl_kernel ROCm wheel).", + ) +) +register_kernel( + KernelSpec( + op="moe.topk_softmax", + backend=KernelBackend.JIT, + target="sglang.kernels.ops.moe.moe_topk_softmax:topk_softmax", + capabilities=_CUDA, + format_signature=FormatSignature( + in_place=True, + description="top-k softmax routing weights/ids", + ), + description="MoE top-k softmax (sglang.kernels.jit).", ) ) @@ -103,7 +118,7 @@ def topk_softmax( correction_bias: Optional[torch.Tensor] = None, ) -> None: """Compute top-k softmax routing weights/ids for MoE.""" - return get_kernel("moe.topk_softmax", KernelBackend.AOT)( + return get_kernel("moe.topk_softmax")( topk_weights, topk_ids, gating_output, diff --git a/python/sglang/kernels/ops/moe/moe_topk_softmax.py b/python/sglang/kernels/ops/moe/moe_topk_softmax.py new file mode 100644 index 000000000..93bf1f77e --- /dev/null +++ b/python/sglang/kernels/ops/moe/moe_topk_softmax.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args +from sglang.srt.utils.custom_op import register_custom_op + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +@cache_once +def _jit_moe_topk_softmax_module(dtype: torch.dtype) -> Module: + args = make_cpp_args(dtype) + return load_jit( + "moe_topk_softmax", + *args, + cuda_files=["moe/moe_topk_softmax.cuh"], + cuda_wrappers=[("topk_softmax", f"topk_softmax<{args}>")], + extra_cuda_cflags=["--use_fast_math"], + ) + + +@register_custom_op( + op_name="moe_topk_softmax_out", + mutates_args=["topk_weights", "topk_ids", "workspace"], +) +def moe_topk_softmax_out( + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + gating_output: torch.Tensor, + workspace: torch.Tensor, + renormalize: bool, + moe_softcapping: float, + correction_bias: Optional[torch.Tensor], +) -> None: + """ + Fused softmax top-k MoE gate (destination-passing style). + + Args: + topk_weights: [num_tokens, topk], float32, pre-allocated output + topk_ids: [num_tokens, topk], int32, pre-allocated output + gating_output: [num_tokens, num_experts], fp32/fp16/bf16 + workspace: [num_tokens * num_experts] float32 scratch (may be size 1 + when num_experts is a supported power-of-2 <= 256) + renormalize: whether to renormalize weights to sum to 1 per row + moe_softcapping: tanh softcapping value applied to the logits, 0 disables + correction_bias: [num_experts] float32 per-expert bias, or None + """ + module = _jit_moe_topk_softmax_module(gating_output.dtype) + module.topk_softmax( + topk_weights, + topk_ids, + gating_output, + workspace, + renormalize, + moe_softcapping, + correction_bias, + ) + + +def topk_softmax( + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + gating_output: torch.Tensor, + renormalize: bool = False, + moe_softcapping: float = 0.0, + correction_bias: Optional[torch.Tensor] = None, +) -> None: + """ + Fused softmax top-k MoE gate with the same call signature as + ``sgl_kernel.topk_softmax`` (destination-passing, in-place). + + Args: + topk_weights: [num_tokens, topk] float32, written in-place + topk_ids: [num_tokens, topk] int32, written in-place + gating_output: [num_tokens, num_experts] fp32/fp16/bf16 + renormalize: whether to renormalize weights to sum to 1 per row + moe_softcapping: tanh softcapping value applied to the logits, 0 disables + correction_bias: [num_experts] float32 per-expert bias, or None + """ + num_tokens = gating_output.shape[0] + num_experts = gating_output.shape[1] + + # The warp-specialized fast path covers power-of-two expert counts up to 512 + # and needs no scratch; everything else goes through the two-pass + # softmax + top-k path, which writes probabilities to the workspace first. + # This threshold must stay in sync with the `case 512:` in the .cuh + # dispatcher and with the AOT host entry in + # kernels/aot/csrc/moe/moe_topk_softmax_kernels.cu. + is_pow2 = num_experts != 0 and (num_experts & (num_experts - 1)) == 0 + needs_workspace = not is_pow2 or num_experts > 512 + workspace_size = num_tokens * num_experts if needs_workspace else 1 + workspace = torch.empty( + workspace_size, dtype=torch.float32, device=gating_output.device + ) + + moe_topk_softmax_out( + topk_weights, + topk_ids, + gating_output, + workspace, + renormalize, + moe_softcapping, + correction_bias, + ) diff --git a/test/registered/kernels/benchmark/moe/bench_moe_topk_softmax.py b/test/registered/kernels/benchmark/moe/bench_moe_topk_softmax.py new file mode 100644 index 000000000..f43b38ba3 --- /dev/null +++ b/test/registered/kernels/benchmark/moe/bench_moe_topk_softmax.py @@ -0,0 +1,65 @@ +import torch +from sgl_kernel import topk_softmax as aot_topk_softmax + +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.jit.benchmark.utils import create_random +from sglang.kernels.ops.moe.moe_topk_softmax import topk_softmax as jit_topk_softmax +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=20, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" +) + + +def _jit(topk_weights, topk_ids, gating_output): + jit_topk_softmax(topk_weights, topk_ids, gating_output) + + +def _aot(topk_weights, topk_ids, gating_output): + aot_topk_softmax( + topk_weights=topk_weights, topk_ids=topk_ids, gating_output=gating_output + ) + + +def _torch(topk_weights, topk_ids, gating_output): + probs = torch.softmax(gating_output.float(), dim=-1) + return probs.topk(topk_weights.shape[-1], dim=-1) + + +FN_MAP = { + "jit": _jit, + "aot": _aot, + "torch": _torch, +} + + +# 32/128/256/512 take the warp-specialized power-of-two path; 12/160 fall back to +# the two-pass path through the softmax workspace. 512 is kept because it sits on +# the boundary between the two. +@marker.parametrize("num_tokens", [128, 512, 1024, 4096, 8192, 32768], [512, 4096]) +@marker.parametrize("num_experts", [32, 128, 256, 512, 12, 160], [256, 160]) +@marker.parametrize("topk", [1, 2, 4, 8], [2]) +@marker.benchmark("impl", ["jit", "aot", "torch"]) +def benchmark(num_tokens: int, num_experts: int, topk: int, impl: str): + if topk > num_experts: + marker.skip("topk must be <= num_experts") + + gating_output = create_random(num_tokens, num_experts, dtype=torch.float32) + topk_weights = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda") + topk_ids = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda") + + return marker.do_bench( + FN_MAP[impl], + input_args=(topk_weights, topk_ids, gating_output), + # Only the gating logits are read, so they are the only arg worth + # rotating to defeat the L2 cache; the two outputs are written every + # iteration. + graph_clone_args=(2,), + # Routing is latency-bound at these sizes, so an achieved-bandwidth + # number is not meaningful; report latency only. + disable_log_bandwidth=True, + ) + + +if __name__ == "__main__": + benchmark.run() diff --git a/test/registered/kernels/ops/moe/test_moe_topk_softmax.py b/test/registered/kernels/ops/moe/test_moe_topk_softmax.py new file mode 100644 index 000000000..e90b49ff2 --- /dev/null +++ b/test/registered/kernels/ops/moe/test_moe_topk_softmax.py @@ -0,0 +1,200 @@ +"""Correctness tests for the JIT :func:`topk_softmax` MoE router. + +The JIT kernel is a host-side port of the AOT ``sgl_kernel.topk_softmax``: the +device code is unchanged, only the host launcher moves to the tvm-ffi +``TensorView`` API and the softmax workspace is allocated by the Python wrapper. + +We validate it two ways: + +* against a definition-based torch reference (documents the math), and +* against the AOT kernel it replaces, when ``sgl_kernel`` is importable. + +Both expert-count regimes are covered: the warp-specialized fast path +(power-of-two ``num_experts`` <= 512, no scratch) and the two-pass +softmax + top-k path that everything else falls back to. 512 and 1024 are both +in the matrix so the boundary between them is pinned on either side. + +Against the torch reference, index comparisons are tie-robust: rather than +requiring identical index tensors, we check that the probability sitting at each +returned index matches the returned weight, so an arbitrary but valid tie-break +is accepted. Against the AOT kernel the comparison is exact, because the device +code is the same on both sides. +""" + +from __future__ import annotations + +import sys +from typing import Optional + +import pytest +import torch + +from sglang.kernels.jit.utils import get_ci_test_range +from sglang.kernels.ops.moe.moe_topk_softmax import topk_softmax +from sglang.test.ci.ci_register import register_cuda_ci + +# CI runs the trimmed matrix (17 cases, one dtype), but on a cold runner the +# single JIT compile is ~29s of the ~31s total -- the case count is nearly free. +register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +DEVICE = "cuda" + +try: + from sgl_kernel import topk_softmax as aot_topk_softmax + + AOT_AVAILABLE = True +except Exception: # pragma: no cover - depends on the installed wheel + aot_topk_softmax = None + AOT_AVAILABLE = False + +DTYPES = get_ci_test_range( + full_range=[torch.float32, torch.float16, torch.bfloat16], + ci_range=[torch.bfloat16], +) +# 8/128/256/512 exercise the warp-specialized power-of-two fast path; +# 6/160/1024 exercise the workspace (two-pass) path. 512 and 1024 sit either +# side of the boundary between them, so both must stay in the list. +NUM_EXPERTS = get_ci_test_range( + full_range=[8, 128, 256, 512, 6, 160, 1024], + ci_range=[8, 160], +) +TOPKS = get_ci_test_range(full_range=[1, 2, 4, 8], ci_range=[2]) +SOFTCAPS = get_ci_test_range(full_range=[0.0, 30.0], ci_range=[0.0]) + + +def _reference_probs( + gating_output: torch.Tensor, + moe_softcapping: float, + correction_bias: Optional[torch.Tensor], +) -> torch.Tensor: + """Definition-based reference for the routing probabilities.""" + logits = gating_output.float() + if moe_softcapping: + logits = torch.tanh(logits / moe_softcapping) * moe_softcapping + if correction_bias is not None: + logits = logits + correction_bias.float() + return torch.softmax(logits, dim=-1) + + +def _run_jit(gating_output: torch.Tensor, topk: int, renormalize: bool, softcap, bias): + num_tokens = gating_output.shape[0] + topk_weights = torch.empty( + (num_tokens, topk), dtype=torch.float32, device=gating_output.device + ) + topk_ids = torch.empty( + (num_tokens, topk), dtype=torch.int32, device=gating_output.device + ) + topk_softmax(topk_weights, topk_ids, gating_output, renormalize, softcap, bias) + return topk_weights, topk_ids + + +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("num_experts", NUM_EXPERTS) +@pytest.mark.parametrize("topk", TOPKS) +@pytest.mark.parametrize("moe_softcapping", SOFTCAPS) +@pytest.mark.parametrize("use_bias", [False, True]) +@pytest.mark.parametrize("renormalize", [False, True]) +def test_topk_softmax_vs_torch( + dtype, num_experts, topk, moe_softcapping, use_bias, renormalize +): + if topk > num_experts: + pytest.skip("topk must be <= num_experts") + + num_tokens = 200 + gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device=DEVICE) + correction_bias = ( + torch.randn(num_experts, dtype=torch.float32, device=DEVICE) + if use_bias + else None + ) + + weights, ids = _run_jit( + gating_output, topk, renormalize, moe_softcapping, correction_bias + ) + + probs = _reference_probs(gating_output, moe_softcapping, correction_bias) + ref_weights, _ = probs.topk(topk, dim=-1) + expected = ( + ref_weights / ref_weights.sum(-1, keepdim=True) if renormalize else ref_weights + ) + + tol = 1e-3 if dtype == torch.float32 else 2e-2 + torch.testing.assert_close(weights, expected, rtol=tol, atol=tol) + + # Tie-robust index check: the probability at each returned index must equal + # the returned weight (undoing renormalization first). + gathered = torch.gather(probs, 1, ids.long()) + unnormalized = ( + weights * ref_weights.sum(-1, keepdim=True) if renormalize else weights + ) + torch.testing.assert_close(gathered, unnormalized, rtol=tol, atol=tol) + + # Indices must be distinct within a row. + sorted_ids, _ = ids.sort(dim=-1) + assert (sorted_ids[:, 1:] != sorted_ids[:, :-1]).all(), "duplicate expert ids" + assert ((ids >= 0) & (ids < num_experts)).all(), "expert id out of range" + + +@pytest.mark.skipif(not AOT_AVAILABLE, reason="sgl_kernel (AOT) is not importable") +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("num_experts", NUM_EXPERTS) +@pytest.mark.parametrize("topk", TOPKS) +@pytest.mark.parametrize("renormalize", [False, True]) +def test_topk_softmax_matches_aot(dtype, num_experts, topk, renormalize): + """The JIT port must be numerically identical to the AOT kernel.""" + if topk > num_experts: + pytest.skip("topk must be <= num_experts") + + num_tokens = 512 + gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device=DEVICE) + + jit_weights, jit_ids = _run_jit(gating_output, topk, renormalize, 0.0, None) + + aot_weights = torch.empty((num_tokens, topk), dtype=torch.float32, device=DEVICE) + aot_ids = torch.empty((num_tokens, topk), dtype=torch.int32, device=DEVICE) + aot_topk_softmax(aot_weights, aot_ids, gating_output, renormalize, 0.0, None) + + # The device code is unchanged from the AOT kernel and both dispatchers + # agree on which path each expert count takes, so this is bit-identical -- + # on the warp-specialized path and on the two-pass path alike. Keeping the + # tolerance at exactly zero is deliberate: it is what caught the dispatcher + # falling out of sync with the AOT one at num_experts == 512. + assert torch.equal(jit_ids, aot_ids) + torch.testing.assert_close(jit_weights, aot_weights, rtol=0, atol=0) + + +@pytest.mark.parametrize("num_experts", [8, 160]) +def test_topk_softmax_single_token(num_experts): + gating_output = torch.randn((1, num_experts), dtype=torch.bfloat16, device=DEVICE) + weights, ids = _run_jit(gating_output, 2, True, 0.0, None) + torch.testing.assert_close( + weights.sum(-1), torch.ones(1, device=DEVICE), rtol=1e-2, atol=1e-2 + ) + assert ids.shape == (1, 2) + + +@pytest.mark.parametrize("num_experts", [8, 160]) +def test_topk_softmax_full_topk(num_experts): + """topk == num_experts: weights must be a permutation of the full softmax.""" + gating_output = torch.randn((16, num_experts), dtype=torch.float32, device=DEVICE) + weights, ids = _run_jit(gating_output, num_experts, False, 0.0, None) + probs = _reference_probs(gating_output, 0.0, None) + torch.testing.assert_close( + weights.sort(dim=-1).values, probs.sort(dim=-1).values, rtol=1e-3, atol=1e-3 + ) + assert ( + ids.sort(dim=-1) + .values.eq(torch.arange(num_experts, device=DEVICE, dtype=torch.int32)) + .all() + ) + + +def test_topk_softmax_zero_tokens(): + """An empty batch must be a no-op rather than a launch failure.""" + gating_output = torch.randn((0, 8), dtype=torch.bfloat16, device=DEVICE) + weights, ids = _run_jit(gating_output, 2, False, 0.0, None) + assert weights.shape == (0, 2) and ids.shape == (0, 2) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"]))