diff --git a/src/base/moe_topk_softmax.h b/src/base/moe_topk_softmax.h new file mode 100644 index 000000000..166f94385 --- /dev/null +++ b/src/base/moe_topk_softmax.h @@ -0,0 +1,73 @@ +#ifndef INFINI_OPS_BASE_MOE_TOPK_SOFTMAX_H_ +#define INFINI_OPS_BASE_MOE_TOPK_SOFTMAX_H_ + +#include + +#include "operator.h" + +namespace infini::ops { + +class MoeTopkSoftmax : public Operator { + public: + MoeTopkSoftmax(Tensor topk_weights, Tensor topk_indices, Tensor gating_output, + Tensor correction_bias, bool renormalize, + float moe_softcapping) + : topk_weights_shape_{topk_weights.shape()}, + topk_indices_shape_{topk_indices.shape()}, + gating_output_shape_{gating_output.shape()}, + correction_bias_shape_{correction_bias.shape()}, + renormalize_{renormalize}, + moe_softcapping_{moe_softcapping}, + num_tokens_{gating_output.size(0)}, + num_experts_{gating_output.size(1)}, + topk_{static_cast(topk_weights.size(1))}, + has_correction_bias_{correction_bias.numel() > 0} { + assert(gating_output.ndim() == 2 && + "`MoeTopkSoftmax` gating_output must be a 2D tensor"); + assert(topk_weights.ndim() == 2 && + "`MoeTopkSoftmax` topk_weights must be a 2D tensor"); + assert(topk_indices.ndim() == 2 && + "`MoeTopkSoftmax` topk_indices must be a 2D tensor"); + assert(topk_weights.dtype() == DataType::kFloat32 && + "`MoeTopkSoftmax` topk_weights must be float32"); + assert(topk_indices.dtype() == DataType::kInt32 && + "`MoeTopkSoftmax` topk_indices must be int32"); + assert(gating_output.dtype() == DataType::kFloat16 || + gating_output.dtype() == DataType::kBFloat16 || + gating_output.dtype() == DataType::kFloat32 && + "`MoeTopkSoftmax` gating_output must be fp16/bf16/fp32"); + assert(topk_ > 0 && topk_ <= static_cast(num_experts_) && + "`MoeTopkSoftmax` topk must be in (0, num_experts]"); + assert(topk_weights.size(0) == static_cast(num_tokens_) && + topk_indices.size(0) == static_cast(num_tokens_) && + topk_indices.size(1) == static_cast(topk_) && + "`MoeTopkSoftmax` output shapes must match"); + assert( + correction_bias.numel() == 0 || + (correction_bias.ndim() == 1 && + correction_bias.numel() == static_cast(num_experts_) && + correction_bias.dtype() == DataType::kFloat32) && + "`MoeTopkSoftmax` correction_bias must be empty or (num_experts,) " + "float32"); + } + + virtual void operator()(Tensor topk_weights, Tensor topk_indices, + Tensor gating_output, Tensor correction_bias, + bool renormalize, float moe_softcapping) const = 0; + + protected: + Tensor::Shape topk_weights_shape_; + Tensor::Shape topk_indices_shape_; + Tensor::Shape gating_output_shape_; + Tensor::Shape correction_bias_shape_; + bool renormalize_{false}; + float moe_softcapping_{0.0f}; + Tensor::Size num_tokens_{0}; + Tensor::Size num_experts_{0}; + int64_t topk_{0}; + bool has_correction_bias_{false}; +}; + +} // namespace infini::ops + +#endif diff --git a/src/native/cuda/metax/ops/moe_topk_softmax/kernel.h b/src/native/cuda/metax/ops/moe_topk_softmax/kernel.h new file mode 100644 index 000000000..bad6229f6 --- /dev/null +++ b/src/native/cuda/metax/ops/moe_topk_softmax/kernel.h @@ -0,0 +1,21 @@ +#ifndef INFINI_OPS_METAX_MOE_TOPK_SOFTMAX_KERNEL_H_ +#define INFINI_OPS_METAX_MOE_TOPK_SOFTMAX_KERNEL_H_ + +#include + +#include "native/cuda/metax/caster.cuh" +#include "native/cuda/metax/runtime_.h" +#include "native/cuda/ops/moe_topk_softmax/kernel.h" + +namespace infini::ops { + +template <> +class Operator + : public CudaMoeTopkSoftmax> { + public: + using CudaMoeTopkSoftmax>::CudaMoeTopkSoftmax; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_METAX_MOE_TOPK_SOFTMAX_KERNEL_H_ \ No newline at end of file diff --git a/src/native/cuda/ops/moe_topk_softmax/kernel.cuh b/src/native/cuda/ops/moe_topk_softmax/kernel.cuh new file mode 100644 index 000000000..42a53736a --- /dev/null +++ b/src/native/cuda/ops/moe_topk_softmax/kernel.cuh @@ -0,0 +1,601 @@ +/* + * Portions of the CUDA kernels in this file are adapted from SGLang: + * /sgl-kernel/csrc/moe/moe_topk_softmax_kernels.cu + * + * Copyright 2025 SGLang Team. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0. + */ + +#ifndef INFINI_OPS_CUDA_MOE_TOPK_SOFTMAX_KERNEL_CUH_ +#define INFINI_OPS_CUDA_MOE_TOPK_SOFTMAX_KERNEL_CUH_ + +#include +#include +#include +#include + +#include "native/cuda/caster.cuh" + +namespace infini::ops { + +namespace { + +constexpr int kWarpSize = 32; + +// Vectorized memory access helper used by the power-of-two expert kernels. +template +class alignas(Alignment) AlignedArray { + T data[N]; +}; + +template +__device__ float ConvertToFloat(T x) { + return Caster::template Cast(x); +} + +struct MoeMaxReduceOp { + __device__ __host__ float operator()(float a, float b) const { + return a > b ? a : b; + } +}; + +} // namespace + +namespace moe { + +using cub_kvp = cub::KeyValuePair; + +struct MoeTopKPair { + static const int PAIR = 2; + static const int MAX_INDEX = 0; + cub_kvp max; + cub_kvp second_max; + + __device__ MoeTopKPair() {} + __device__ MoeTopKPair(cub_kvp max, cub_kvp second_max) + : max(max), second_max(second_max) {} +}; + +// Reduces a (max, second-max) KVP pair into a single pair. +struct MoeTopKPairArgMax { + __device__ MoeTopKPairArgMax() {} + __device__ __forceinline__ MoeTopKPair operator()( + const MoeTopKPair& candidate1, const MoeTopKPair& candidate2) const { + cub_kvp global_max, global_second_max; + if (candidate1.max.value > candidate2.max.value) { + global_max = candidate1.max; + } else { + global_max = candidate2.max; + } + if (global_max.key == candidate1.max.key) { + global_second_max = (candidate1.second_max.value > candidate2.max.value) + ? candidate1.second_max + : candidate2.max; + } else { + global_second_max = (candidate2.second_max.value > candidate1.max.value) + ? candidate2.second_max + : candidate1.max; + } + return MoeTopKPair(global_max, global_second_max); + } +}; + +} // namespace moe + +// Generic softmax over a single row of `num_cols` experts. Writes the +// (softcapped) distribution into `output`. +template +__launch_bounds__(TPB) __global__ + void MoeSoftmaxKernel(const T* input, float* output, const int num_cols, + const float moe_softcapping) { + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tmp_storage; + __shared__ float normalizing_factor; + __shared__ float float_max; + + const int thread_row_offset = blockIdx.x * num_cols; + + float thread_data = -FLT_MAX; + for (int ii = threadIdx.x; ii < num_cols; ii += TPB) { + const int idx = thread_row_offset + ii; + float val = ConvertToFloat(input[idx]); + if (moe_softcapping != 0.0f) { + val = tanhf(val / moe_softcapping) * moe_softcapping; + } + output[idx] = val; + thread_data = fmaxf(val, thread_data); + } + + const float max_elem = + BlockReduce(tmp_storage).Reduce(thread_data, MoeMaxReduceOp()); + if (threadIdx.x == 0) { + float_max = max_elem; + } + __syncthreads(); + + thread_data = 0.0f; + for (int ii = threadIdx.x; ii < num_cols; ii += TPB) { + const int idx = thread_row_offset + ii; + thread_data += expf(output[idx] - float_max); + } + const float z = BlockReduce(tmp_storage).Sum(thread_data); + if (threadIdx.x == 0) { + normalizing_factor = 1.0f / z; + } + __syncthreads(); + + for (int ii = threadIdx.x; ii < num_cols; ii += TPB) { + const int idx = thread_row_offset + ii; + output[idx] = expf(output[idx] - float_max) * normalizing_factor; + } +} + +// Two experts per pass over the block. Used together with the templated +// power-of-two kernel as the fast selection path. +template +__launch_bounds__(TPB) __global__ + void MoeTopKFastKernel(float* inputs_after_softmax, float* output, + int* indices, const int num_experts, const int k, + const bool renormalize, + const float* correction_bias) { + using namespace moe; + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tmp_storage; + MoeTopKPair thread_pair; + + const int block_row = blockIdx.x; + const int thread_read_offset = blockIdx.x * num_experts; + float row_sum_for_renormalize = 0.0f; + + for (int k_idx = 0; k_idx < (k + MoeTopKPair::PAIR - 1) / MoeTopKPair::PAIR; + ++k_idx) { + thread_pair.max.key = 0; + thread_pair.max.value = -1.0f; + thread_pair.second_max.key = 0; + thread_pair.second_max.value = -1.0f; + + 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; + const float prob = inputs_after_softmax[idx]; + inp_kvp.value = + correction_bias == nullptr ? prob : prob + correction_bias[expert]; + if (inp_kvp.value > thread_pair.max.value) { + thread_pair.second_max = thread_pair.max; + thread_pair.max = inp_kvp; + } else if (inp_kvp.value > thread_pair.second_max.value) { + thread_pair.second_max = inp_kvp; + } + } + + MoeTopKPairArgMax reducer; + const MoeTopKPair result_pair = + BlockReduce(tmp_storage).Reduce(thread_pair, reducer); + if (threadIdx.x == 0) { +#pragma unroll + for (int i = 0; i < MoeTopKPair::PAIR; ++i) { + if (k_idx * 2 + i >= k) { + break; + } + cub_kvp result = (i == MoeTopKPair::MAX_INDEX) ? result_pair.max + : result_pair.second_max; + int expert = result.key; + const float prob = inputs_after_softmax[thread_read_offset + expert]; + inputs_after_softmax[thread_read_offset + expert] = -FLT_MAX; + int idx = k * block_row + k_idx * 2 + i; + output[idx] = prob; + indices[idx] = expert; + row_sum_for_renormalize += prob; + } + } + __syncthreads(); + } + + if (renormalize && threadIdx.x == 0) { + const float inv = 1.0f / row_sum_for_renormalize; + for (int k_idx = 0; k_idx < k; ++k_idx) { + const int idx = k * block_row + k_idx; + output[idx] *= inv; + } + } +} + +// One expert per pass over the block. Slower than MoeTopKFastKernel but +// applies to arbitrary `k`. +template +__launch_bounds__(TPB) __global__ + void MoeTopKKernel(float* inputs_after_softmax, float* output, int* indices, + const int num_experts, const int k, + const bool renormalize, const float* correction_bias) { + using cub_kvp = cub::KeyValuePair; + using BlockReduce = cub::BlockReduce; + __shared__ typename BlockReduce::TempStorage tmp_storage; + cub_kvp thread_kvp; + cub::ArgMax arg_max; + + const int block_row = blockIdx.x; + const int thread_read_offset = blockIdx.x * num_experts; + float row_sum_for_renormalize = 0.0f; + + for (int k_idx = 0; k_idx < k; ++k_idx) { + thread_kvp.key = 0; + thread_kvp.value = -1.0f; + 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; + const float prob = inputs_after_softmax[idx]; + inp_kvp.value = + correction_bias == nullptr ? prob : prob + correction_bias[expert]; + thread_kvp = arg_max(inp_kvp, thread_kvp); + } + + const cub_kvp result_kvp = + BlockReduce(tmp_storage).Reduce(thread_kvp, arg_max); + if (threadIdx.x == 0) { + const int expert = result_kvp.key; + const int idx = k * block_row + k_idx; + const float prob = inputs_after_softmax[thread_read_offset + expert]; + output[idx] = prob; + indices[idx] = expert; + row_sum_for_renormalize += prob; + inputs_after_softmax[thread_read_offset + expert] = -FLT_MAX; + } + __syncthreads(); + } + + if (renormalize && threadIdx.x == 0) { + const float inv = 1.0f / row_sum_for_renormalize; + for (int k_idx = 0; k_idx < k; ++k_idx) { + const int idx = k * block_row + k_idx; + output[idx] *= inv; + } + } +} + +// Vectorized gating-softmax + top-k for power-of-two expert counts. Each row +// is processed cooperatively by `THREADS_PER_ROW` lanes, producing both the +// softmax probabilities and the (optionally biased) top-k selection. +template +__launch_bounds__(WARPS_PER_CTA* kWarpSize) __global__ + void TopkGatingSoftmaxKernel(const T* input, float* output, + const int num_rows, int* indices, const int k, + 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 constexpr int ELTS_PER_WARP = kWarpSize * 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 constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW; + + static_assert(VPT % ELTS_PER_LDG == 0, + "VPT must be a multiple of elements per load"); + static_assert(kWarpSize % THREADS_PER_ROW == 0, + "threads per row must divide warp size"); + static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW), + "THREADS_PER_ROW must be power of 2"); + static_assert(THREADS_PER_ROW <= kWarpSize, + "THREADS_PER_ROW can be at most warp size"); + static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0, + "row elements must divide warp elements"); + + 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 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]; + auto* row_chunk_vec_ptr = reinterpret_cast(&row_chunk_temp); + const auto* 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] = ConvertToFloat(row_chunk_temp[ii]); + } + + if (moe_softcapping != 0.0f) { +#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; + } + row_chunk[ii] = val; + } + } + + float thread_max = row_chunk[0]; +#pragma unroll + for (int ii = 1; ii < VPT; ++ii) { + thread_max = fmaxf(thread_max, row_chunk[ii]); + } +#pragma unroll + for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { + thread_max = fmaxf(thread_max, __shfl_xor_sync(0xffffffff, thread_max, mask, + THREADS_PER_ROW)); + } + + float row_sum = 0.0f; +#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 += __shfl_xor_sync(0xffffffff, row_sum, mask, THREADS_PER_ROW); + } + const float reciprocal_row_sum = 1.0f / row_sum; +#pragma unroll + for (int ii = 0; ii < VPT; ++ii) { + row_chunk[ii] *= reciprocal_row_sum; + } + + const int start_col = first_elt_read_by_thread; + float row_sum_for_renormalize = 0.0f; + for (int k_idx = 0; k_idx < k; ++k_idx) { + float max_prob = row_chunk[0]; + float max_choice = correction_bias == nullptr + ? max_prob + : max_prob + correction_bias[start_col]; + 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) { + const int expert_idx = col + ii; + float prob = row_chunk[ldg * ELTS_PER_LDG + ii]; + float choice = correction_bias == nullptr + ? prob + : prob + correction_bias[expert_idx]; + if (choice > max_choice) { + max_choice = choice; + max_prob = prob; + expert = expert_idx; + } + } + } + +#pragma unroll + for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) { + float other_choice = + __shfl_xor_sync(0xffffffff, max_choice, mask, THREADS_PER_ROW); + float other_prob = + __shfl_xor_sync(0xffffffff, max_prob, mask, THREADS_PER_ROW); + int other_expert = + __shfl_xor_sync(0xffffffff, expert, mask, THREADS_PER_ROW); + if (other_choice > max_choice || + (other_choice == max_choice && other_expert < expert)) { + max_choice = other_choice; + max_prob = other_prob; + expert = other_expert; + } + } + + if (thread_group_idx == 0) { + const int idx = k * thread_row + k_idx; + output[idx] = max_prob; + indices[idx] = expert; + row_sum_for_renormalize += max_prob; + } + + 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] = + -FLT_MAX; + } + } + } + + if (renormalize && thread_group_idx == 0) { + const float inv = 1.0f / 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] *= inv; + } + } +} + +namespace detail { + +template +struct MoeTopkConstants { + static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T); + static_assert(EXPERTS / (ELTS_PER_LDG * kWarpSize) == 0 || + EXPERTS % (ELTS_PER_LDG * kWarpSize) == 0, + ""); + static constexpr int VECS_PER_THREAD = (EXPERTS / + (ELTS_PER_LDG * kWarpSize)) > 1 + ? (EXPERTS / + (ELTS_PER_LDG * kWarpSize)) + : 1; + static constexpr int VPT = VECS_PER_THREAD * ELTS_PER_LDG; + static constexpr int THREADS_PER_ROW = EXPERTS / VPT; + static constexpr int ROWS_PER_WARP = kWarpSize / THREADS_PER_ROW; +}; + +} // namespace detail + +namespace { + +// Launches the vectorized power-of-two kernel for a template-computed +// configuration. `StreamT` is the backend-specific stream type. +template +void LaunchTopkGatingSoftmaxHelper(const T* input, float* output, int* indices, + const int num_rows, const int k, + const bool renormalize, + const float moe_softcapping, + const float* correction_bias, + StreamT stream) { + static constexpr int MAX_BYTES_PER_LDG = 16; + static constexpr int BYTES_PER_LDG = + (MAX_BYTES_PER_LDG < static_cast(sizeof(T)) * EXPERTS) + ? MAX_BYTES_PER_LDG + : static_cast(sizeof(T)) * EXPERTS; + using Constants = detail::MoeTopkConstants; + static constexpr int VPT = Constants::VPT; + + if constexpr (EXPERTS > kWarpSize * VPT) { + return; + } else { + 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(kWarpSize, WARPS_PER_TB); + TopkGatingSoftmaxKernel + <<>>(input, output, num_rows, indices, + k, renormalize, moe_softcapping, + correction_bias); + } +} + +template +void LaunchMoeTopkGeneric(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, StreamT stream) { + assert(softmax_workspace != nullptr && + "moe_topk_softmax generic path requires a workspace"); + MoeSoftmaxKernel<<>>( + gating_output, softmax_workspace, num_experts, moe_softcapping); + if (topk == 1) { + MoeTopKKernel<<>>( + softmax_workspace, topk_weights, topk_indices, num_experts, topk, + renormalize, correction_bias); + } else { + MoeTopKFastKernel<<>>( + softmax_workspace, topk_weights, topk_indices, num_experts, topk, + renormalize, correction_bias); + } +} + +} // namespace + +// Reports whether the generic fallback path requires a device workspace buffer +// sized `num_tokens * num_experts * sizeof(float)`. Zero for the power-of-two +// fast path. +inline bool MoeTopkSoftmaxNeedsWorkspace(std::size_t num_experts) { + const bool is_pow_2 = + num_experts != 0 && ((num_experts & (num_experts - 1)) == 0); + return !is_pow_2 || num_experts > 512; +} + +// Host entry point dispatching to either the power-of-two vectorized kernel +// or the generic softmax + top-k fallback. `workspace` is only required when +// `MoeTopkSoftmaxNeedsWorkspace(num_experts)` is true. +template +void LaunchMoeTopkSoftmax(const T* gating_output, float* topk_weights, + int* topk_indices, float* workspace, + const int num_tokens, const int num_experts, + const int topk, const bool renormalize, + const float moe_softcapping, + const float* correction_bias, StreamT stream) { + static constexpr int WARPS_PER_TB = 4; + if (!MoeTopkSoftmaxNeedsWorkspace(num_experts)) { + switch (num_experts) { + case 1: + LaunchTopkGatingSoftmaxHelper( + gating_output, topk_weights, topk_indices, num_tokens, topk, + renormalize, moe_softcapping, correction_bias, stream); + break; + case 2: + LaunchTopkGatingSoftmaxHelper( + gating_output, topk_weights, topk_indices, num_tokens, topk, + renormalize, moe_softcapping, correction_bias, stream); + break; + case 4: + LaunchTopkGatingSoftmaxHelper( + gating_output, topk_weights, topk_indices, num_tokens, topk, + renormalize, moe_softcapping, correction_bias, stream); + break; + case 8: + LaunchTopkGatingSoftmaxHelper( + gating_output, topk_weights, topk_indices, num_tokens, topk, + renormalize, moe_softcapping, correction_bias, stream); + break; + case 16: + LaunchTopkGatingSoftmaxHelper( + gating_output, topk_weights, topk_indices, num_tokens, topk, + renormalize, moe_softcapping, correction_bias, stream); + break; + case 32: + LaunchTopkGatingSoftmaxHelper( + gating_output, topk_weights, topk_indices, num_tokens, topk, + renormalize, moe_softcapping, correction_bias, stream); + break; + case 64: + LaunchTopkGatingSoftmaxHelper( + gating_output, topk_weights, topk_indices, num_tokens, topk, + renormalize, moe_softcapping, correction_bias, stream); + break; + case 128: + LaunchTopkGatingSoftmaxHelper( + gating_output, topk_weights, topk_indices, num_tokens, topk, + renormalize, moe_softcapping, correction_bias, stream); + break; + case 256: + LaunchTopkGatingSoftmaxHelper( + gating_output, topk_weights, topk_indices, num_tokens, topk, + renormalize, moe_softcapping, correction_bias, stream); + break; + case 512: + LaunchTopkGatingSoftmaxHelper( + gating_output, topk_weights, topk_indices, num_tokens, topk, + renormalize, moe_softcapping, correction_bias, stream); + break; + default: + LaunchMoeTopkGeneric( + gating_output, topk_weights, topk_indices, workspace, num_tokens, + num_experts, topk, renormalize, moe_softcapping, correction_bias, + stream); + break; + } + } else { + LaunchMoeTopkGeneric( + gating_output, topk_weights, topk_indices, workspace, num_tokens, + num_experts, topk, renormalize, moe_softcapping, correction_bias, + stream); + } +} + +} // namespace infini::ops + +#endif // INFINI_OPS_CUDA_MOE_TOPK_SOFTMAX_KERNEL_CUH_ \ No newline at end of file diff --git a/src/native/cuda/ops/moe_topk_softmax/kernel.h b/src/native/cuda/ops/moe_topk_softmax/kernel.h new file mode 100644 index 000000000..1532180ad --- /dev/null +++ b/src/native/cuda/ops/moe_topk_softmax/kernel.h @@ -0,0 +1,72 @@ +#ifndef INFINI_OPS_CUDA_MOE_TOPK_SOFTMAX_KERNEL_H_ +#define INFINI_OPS_CUDA_MOE_TOPK_SOFTMAX_KERNEL_H_ + +#include +#include + +#include "base/moe_topk_softmax.h" +#include "data_type.h" +#include "dispatcher.h" +#include "native/cuda/ops/moe_topk_softmax/kernel.cuh" +#include "native/cuda/runtime_utils.h" + +namespace infini::ops { + +// MoE top-k gating with softmax normalization. For power-of-two expert +// counts (up to 512) a fused, vectorized kernel performs softmax and top-k +// selection in a single pass; otherwise the generic softmax + top-k path is +// used and requires a workspace of num_tokens * num_experts floats. +template +class CudaMoeTopkSoftmax : public MoeTopkSoftmax { + public: + CudaMoeTopkSoftmax(Tensor topk_weights, Tensor topk_indices, + Tensor gating_output, Tensor correction_bias, + bool renormalize, float moe_softcapping) + : MoeTopkSoftmax(topk_weights, topk_indices, gating_output, + correction_bias, renormalize, moe_softcapping), + gating_output_type_{gating_output.dtype()} {} + + std::size_t workspace_size_in_bytes() const override { + if (MoeTopkSoftmaxNeedsWorkspace(num_experts_)) { + return num_tokens_ * num_experts_ * sizeof(float); + } + return 0; + } + + void operator()(Tensor topk_weights, Tensor topk_indices, + Tensor gating_output, Tensor correction_bias, + bool renormalize, float moe_softcapping) const override { + if (num_tokens_ == 0) { + return; + } + + auto stream = static_cast(stream_ ? stream_ : 0); + const float* correction_bias_ptr = + has_correction_bias_ + ? reinterpret_cast(correction_bias.data()) + : nullptr; + + DispatchFunc( + gating_output_type_, + [&](auto type_tag) { + using T = typename decltype(type_tag)::type; + LaunchMoeTopkSoftmax( + reinterpret_cast(gating_output.data()), + reinterpret_cast(topk_weights.data()), + reinterpret_cast(topk_indices.data()), + reinterpret_cast(workspace_), + static_cast(num_tokens_), static_cast(num_experts_), + static_cast(topk_), renormalize, moe_softcapping, + correction_bias_ptr, stream); + }, + "CudaMoeTopkSoftmax::operator()"); + } + + private: + DataType gating_output_type_; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_CUDA_MOE_TOPK_SOFTMAX_KERNEL_H_ \ No newline at end of file diff --git a/tests/test_moe_topk_softmax.py b/tests/test_moe_topk_softmax.py new file mode 100644 index 000000000..4755a3b42 --- /dev/null +++ b/tests/test_moe_topk_softmax.py @@ -0,0 +1,102 @@ +import sys, os +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import infini.ops +import pytest +import torch +import torch.nn.functional as F + +from tests.utils import empty_strided, get_stream + + +def _ref_moe_topk_softmax(gating_output, correction_bias, topk, renormalize, moe_softcapping): + """Reference implementation using PyTorch.""" + # Convert to float32 for computation + gating_fp32 = gating_output.float() + + # Apply softcapping + if moe_softcapping != 0.0: + gating_fp32 = torch.tanh(gating_fp32 / moe_softcapping) * moe_softcapping + + # Softmax (bias is applied *after* softmax in the kernel) + probs = F.softmax(gating_fp32, dim=-1) + + # Top-k with optional post-softmax bias correction. + # The SGLang kernel adds correction_bias to probs for selection, + # but returns the original softmax probs as weights. + if correction_bias is not None: + selection_scores = probs + correction_bias.unsqueeze(0) + else: + selection_scores = probs + topk_weights, topk_indices = torch.topk(selection_scores, topk, dim=-1) + + # Gather the original softmax probs at the selected indices + topk_weights = torch.gather(probs, dim=-1, index=topk_indices) + + # Renormalize if requested + if renormalize: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + + return topk_weights, topk_indices.to(torch.int32) + + +@pytest.mark.parametrize( + "num_tokens, num_experts, topk, renormalize, moe_softcapping, has_bias", + ( + (4, 8, 2, True, 0.0, False), + (8, 16, 4, False, 50.0, False), + (8, 16, 4, True, 50.0, True), + (1, 4, 1, True, 0.0, False), + (16, 64, 6, True, 30.0, True), + ), +) +@pytest.mark.parametrize( + ("dtype", "rtol", "atol"), + ( + (torch.float32, 1e-4, 1e-4), + (torch.float16, 1e-2, 1e-2), + (torch.bfloat16, 1e-1, 5e-2), + ), +) +def test_moe_topk_softmax( + num_tokens, + num_experts, + topk, + renormalize, + moe_softcapping, + has_bias, + dtype, + device, + rtol, + atol, +): + gating_output = torch.randn(num_tokens, num_experts, dtype=dtype, device=device) + + topk_weights = empty_strided( + (num_tokens, topk), None, dtype=torch.float32, device=device + ) + topk_indices = empty_strided( + (num_tokens, topk), None, dtype=torch.int32, device=device + ) + + if has_bias: + correction_bias = torch.randn(num_experts, dtype=torch.float32, device=device) + else: + correction_bias = torch.empty(0, dtype=torch.float32, device=device) + + infini.ops.moe_topk_softmax( + topk_weights, + topk_indices, + gating_output, + correction_bias, + renormalize, + moe_softcapping, + stream=get_stream(device), + ) + + ref_weights, ref_indices = _ref_moe_topk_softmax( + gating_output, correction_bias if has_bias else None, topk, renormalize, moe_softcapping + ) + + torch.testing.assert_close(topk_weights, ref_weights, rtol=rtol, atol=atol) + torch.testing.assert_close(topk_indices, ref_indices, rtol=0, atol=0)