diff --git a/sgl-kernel/csrc/moe/kimi_k2_moe_fused_gate.cu b/sgl-kernel/csrc/moe/kimi_k2_moe_fused_gate.cu index 7f719e52e..9afe4dfcd 100644 --- a/sgl-kernel/csrc/moe/kimi_k2_moe_fused_gate.cu +++ b/sgl-kernel/csrc/moe/kimi_k2_moe_fused_gate.cu @@ -4,21 +4,39 @@ #include -// Kimi K2 specific constants -static constexpr int WARP_SIZE = 32; -static constexpr int WARPS_PER_CTA = 6; -static constexpr int NUM_EXPERTS = 384; -static constexpr int VPT = 12; // 384 / 32 = 12 +// Kimi K2 MoE fused gate, supports NUM_EXPERTS in {256 (MiMo V2 Flash), 384 (Kimi K2)}. +// Routing (DeepSeek "noaux_tc" with num_expert_group = 1): +// 1. sigmoid(gate_logit) +// 2. add per-expert correction bias (ranking only) +// 3. pick top-k by biased score +// 4. weights = sigmoid (no bias) +// 5. optional renorm; routed_scaling_factor folded into renorm (no-op when not renormalizing) -// Small token optimization constants -static constexpr int SMALL_TOKEN_THRESHOLD = 512; -static constexpr int WARPS_PER_TOKEN_SMALL = 12; // Use 12 warps per token for small batches -static constexpr int THREADS_PER_BLOCK_SMALL = WARPS_PER_TOKEN_SMALL * WARP_SIZE; // 384 threads +__device__ __forceinline__ float sigmoid_accurate(float x) { + return 1.0f / (1.0f + expf(-x)); +} -// Vectorization constants (used by large token kernel) -static constexpr int VEC_SIZE = 4; // Use float4 for vectorized loads +template +struct GateConfig { + static_assert( + N == 256 || N == 384, + "kimi_k2_moe_fused_gate currently only supports " + "NUM_EXPERTS == 256 or 384"); + static constexpr int NUM_EXPERTS = N; + static constexpr int WARP_SIZE = 32; + static constexpr int WARPS_PER_CTA = 6; // only used by the large-token kernel + static constexpr int VPT = N / 32; // 8 (256) or 12 (384) + static constexpr int VEC_SIZE = 4; + static constexpr int VEC_PER_LANE = VPT / VEC_SIZE; // 2 or 3 + static constexpr int WARPS_PER_TOKEN_SMALL = N / 32; // 8 or 12 + static constexpr int THREADS_PER_BLOCK_SMALL = N; // 256 or 384 + static constexpr int SMALL_TOKEN_THRESHOLD = 512; + static constexpr int MAX_TOPK = 8; // must match TORCH_CHECK(topk <= 8) at the host launcher + static_assert(VPT % VEC_SIZE == 0, "VPT must be a multiple of VEC_SIZE for the float4 vec load"); +}; -// Small token optimized kernel: Each warp independently finds top-k, then merge, using warp-level topk +// Small-token kernel: 1 block per token, NUM_EXPERTS threads (1 thread = 1 expert). +template __global__ void kimi_k2_moe_fused_gate_kernel_small_token( float* input, float* bias, @@ -29,6 +47,12 @@ __global__ void kimi_k2_moe_fused_gate_kernel_small_token( bool renormalize, double routed_scaling_factor, bool apply_routed_scaling_factor_on_output) { + using Cfg = GateConfig; + constexpr int NUM_EXPERTS = Cfg::NUM_EXPERTS; + constexpr int WARP_SIZE = Cfg::WARP_SIZE; + constexpr int WARPS_PER_TOKEN_SMALL = Cfg::WARPS_PER_TOKEN_SMALL; + constexpr int MAX_TOPK = Cfg::MAX_TOPK; + int64_t row_idx = blockIdx.x; if (row_idx >= num_rows) return; @@ -36,38 +60,30 @@ __global__ void kimi_k2_moe_fused_gate_kernel_small_token( int warp_id = tid / WARP_SIZE; int lane_id = tid % WARP_SIZE; - // Shared memory: biased scores and original scores - __shared__ float shared_scores[NUM_EXPERTS]; + // Sigmoid weights (no bias) for final lookup, indexed by expert id. __shared__ float shared_original_scores[NUM_EXPERTS]; - // For storing selected top-k indices and values - __shared__ int selected_experts[8]; // Up to topk=6, I use 8 for alignment - __shared__ float selected_vals[8]; - // For warp-level reduction __shared__ float warp_maxs[WARPS_PER_TOKEN_SMALL]; __shared__ int warp_experts[WARPS_PER_TOKEN_SMALL]; + __shared__ int selected_experts[MAX_TOPK]; - // Load data: all 384 threads load one expert each - if (tid < NUM_EXPERTS) { - float input_val = input[row_idx * NUM_EXPERTS + tid]; - float bias_val = bias[tid]; - float sigmoid_val = 1.0f / (1.0f + expf(-input_val)); - float biased_val = sigmoid_val + bias_val; - shared_scores[tid] = biased_val; - shared_original_scores[tid] = sigmoid_val; - } + // Keep biased_val in register; mask the winner in-place each iteration to + // avoid round-tripping through shared memory. + float input_val = input[row_idx * NUM_EXPERTS + tid]; + float bias_val = bias[tid]; + float sigmoid_val = sigmoid_accurate(input_val); + float biased_val = sigmoid_val + bias_val; + shared_original_scores[tid] = sigmoid_val; __syncthreads(); - // Find top-k using iterative selection, each iteration finds the next maximum + // Lane 0 of warp 0 accumulates the renorm sum as it picks each winner, + // saving a second pass over selected_experts during writeback. + float sum_for_renorm = 0.0f; + for (int k = 0; k < topk; k++) { - // Each thread holds one expert's value - float my_val = (tid < NUM_EXPERTS) ? shared_scores[tid] : -FLT_MAX; - int my_expert = tid; - - // Use warp-level reduction first - float warp_max_val = my_val; - int warp_max_expert = my_expert; - + // Stage 1: per-warp argmax. + float warp_max_val = biased_val; + int warp_max_expert = tid; #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { float other_val = __shfl_down_sync(0xFFFFFFFF, warp_max_val, offset); @@ -77,20 +93,16 @@ __global__ void kimi_k2_moe_fused_gate_kernel_small_token( warp_max_expert = other_expert; } } - - // Warp leaders write to shared memory if (lane_id == 0) { warp_maxs[warp_id] = warp_max_val; warp_experts[warp_id] = warp_max_expert; } - __syncthreads(); - // Final reduction among warps (done by first warp) + // Stage 2: warp 0 merges warp-leaders into a single winner. if (warp_id == 0) { float final_max = (lane_id < WARPS_PER_TOKEN_SMALL) ? warp_maxs[lane_id] : -FLT_MAX; int final_expert = (lane_id < WARPS_PER_TOKEN_SMALL) ? warp_experts[lane_id] : -1; - #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { float other_val = __shfl_down_sync(0xFFFFFFFF, final_max, offset); @@ -100,59 +112,41 @@ __global__ void kimi_k2_moe_fused_gate_kernel_small_token( final_expert = other_expert; } } - if (lane_id == 0) { selected_experts[k] = final_expert; - selected_vals[k] = final_max; - } - } - - __syncthreads(); - - // Mark the selected expert as used for next iteration - // All threads can read from selected_experts[k] - int selected = selected_experts[k]; - if (tid == selected) { - shared_scores[tid] = -FLT_MAX; - } - - __syncthreads(); - } - - // Write output (done by thread 0) - if (tid == 0) { - for (int k = 0; k < topk; k++) { - int expert_id = selected_experts[k]; - if (expert_id >= 0 && expert_id < NUM_EXPERTS) { - output_ptr[row_idx * topk + k] = shared_original_scores[expert_id]; - indices_ptr[row_idx * topk + k] = expert_id; - } else { - output_ptr[row_idx * topk + k] = 0.0f; - indices_ptr[row_idx * topk + k] = 0; - } - } - - // Renormalization - if (renormalize) { - float sum = 0.0f; - for (int k = 0; k < topk; k++) { - sum += output_ptr[row_idx * topk + k]; - } - - if (sum > 0.0f) { - for (int k = 0; k < topk; k++) { - int64_t idx = row_idx * topk + k; - output_ptr[idx] /= sum; - if (apply_routed_scaling_factor_on_output) { - output_ptr[idx] *= static_cast(routed_scaling_factor); - } + if (renormalize && final_expert >= 0 && final_expert < NUM_EXPERTS) { + sum_for_renorm += shared_original_scores[final_expert]; } } } + __syncthreads(); + + int selected = selected_experts[k]; + if (tid == selected) biased_val = -FLT_MAX; + } + + // Lane 0 of warp 0 writes the output. sum_for_renorm was accumulated + // during the topk loop, so we just fold it into rcp. + if (warp_id == 0 && lane_id == 0) { + float rcp = 1.0f; + if (renormalize && sum_for_renorm > 0.0f) { + rcp = 1.0f / sum_for_renorm; + if (apply_routed_scaling_factor_on_output) { + rcp *= static_cast(routed_scaling_factor); + } + } + + for (int k = 0; k < topk; k++) { + int expert_id = selected_experts[k]; + bool valid = (expert_id >= 0 && expert_id < NUM_EXPERTS); + output_ptr[row_idx * topk + k] = valid ? shared_original_scores[expert_id] * rcp : 0.0f; + indices_ptr[row_idx * topk + k] = valid ? expert_id : 0; + } } } -// Large token kernel: Original implementation with vectorized loads +// Large-token kernel: 1 warp per token, WARPS_PER_CTA warps per block. +template __global__ void kimi_k2_moe_fused_gate_kernel( float* input, float* bias, @@ -163,6 +157,14 @@ __global__ void kimi_k2_moe_fused_gate_kernel( bool renormalize, double routed_scaling_factor, bool apply_routed_scaling_factor_on_output) { + using Cfg = GateConfig; + constexpr int NUM_EXPERTS = Cfg::NUM_EXPERTS; + constexpr int WARP_SIZE = Cfg::WARP_SIZE; + constexpr int WARPS_PER_CTA = Cfg::WARPS_PER_CTA; + constexpr int VEC_SIZE = Cfg::VEC_SIZE; + constexpr int VEC_PER_LANE = Cfg::VEC_PER_LANE; + constexpr int MAX_TOPK = Cfg::MAX_TOPK; + int64_t row_idx = blockIdx.x * WARPS_PER_CTA + threadIdx.y; if (row_idx >= num_rows) return; @@ -171,35 +173,42 @@ __global__ void kimi_k2_moe_fused_gate_kernel( __shared__ float shared_scores[NUM_EXPERTS * WARPS_PER_CTA]; __shared__ float shared_original_scores[NUM_EXPERTS * WARPS_PER_CTA]; - float* warp_scores = shared_scores + warp_id * NUM_EXPERTS; float* warp_original_scores = shared_original_scores + warp_id * NUM_EXPERTS; + float4* warp_scores_v4 = reinterpret_cast(warp_scores); + float4* warp_original_scores_v4 = reinterpret_cast(warp_original_scores); - // Vectorized loading: each lane loads multiple float4 chunks - // VPT = 12, so we load 12/4 = 3 float4 per lane - const int VEC_PER_LANE = VPT / VEC_SIZE; // 3 float4* input_vec = reinterpret_cast(input + row_idx * NUM_EXPERTS); float4* bias_vec = reinterpret_cast(bias); + // Lane-strided vec_idx (each lane k stores at vec_idx k, k+32, k+64, ...) so each + // iteration's STS.128 is lane-contiguous, avoiding shared-mem bank conflicts. #pragma unroll for (int i = 0; i < VEC_PER_LANE; i++) { - int vec_idx = lane_id * VEC_PER_LANE + i; + int vec_idx = lane_id + i * WARP_SIZE; float4 input_val = input_vec[vec_idx]; float4 bias_val = bias_vec[vec_idx]; + float4 sigmoid_v4; + float4 biased_v4; #pragma unroll for (int j = 0; j < VEC_SIZE; j++) { - int expert = vec_idx * VEC_SIZE + j; float inp = ((float*)&input_val)[j]; float b = ((float*)&bias_val)[j]; - float sigmoid_val = 1.0f / (1.0f + expf(-inp)); - float biased_val = sigmoid_val + b; - warp_scores[expert] = biased_val; - warp_original_scores[expert] = sigmoid_val; + float sigmoid_val = sigmoid_accurate(inp); + ((float*)&sigmoid_v4)[j] = sigmoid_val; + ((float*)&biased_v4)[j] = sigmoid_val + b; } + warp_original_scores_v4[vec_idx] = sigmoid_v4; + warp_scores_v4[vec_idx] = biased_v4; } - __syncthreads(); + __syncwarp(); + + // Lane 0 records the picked expert ids and accumulates the renorm sum as + // it goes; the global write is a single pass after the loop. + int top_indices[MAX_TOPK]; + float sum_for_renorm = 0.0f; for (int k = 0; k < topk; k++) { float max_val = -FLT_MAX; @@ -212,10 +221,11 @@ __global__ void kimi_k2_moe_fused_gate_kernel( } } - for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) { + // warp shfl reduce; tie-break by lower expert id +#pragma unroll + for (int offset = 16; offset > 0; offset /= 2) { float other_val = __shfl_down_sync(0xFFFFFFFF, max_val, offset); int other_expert = __shfl_down_sync(0xFFFFFFFF, max_expert, offset); - if (other_val > max_val || (other_val == max_val && other_expert < max_expert)) { max_val = other_val; max_expert = other_expert; @@ -223,37 +233,76 @@ __global__ void kimi_k2_moe_fused_gate_kernel( } if (lane_id == 0) { - int64_t output_idx = row_idx * topk + k; - if (max_expert != -1) { - output_ptr[output_idx] = warp_original_scores[max_expert]; - indices_ptr[output_idx] = max_expert; - warp_scores[max_expert] = -FLT_MAX; - } else { - output_ptr[output_idx] = 0.0f; - indices_ptr[output_idx] = 0; + bool valid = (max_expert >= 0 && max_expert < NUM_EXPERTS); + top_indices[k] = valid ? max_expert : -1; + if (renormalize && valid) { + sum_for_renorm += warp_original_scores[max_expert]; } + if (valid) warp_scores[max_expert] = -FLT_MAX; } - __syncwarp(); } - __syncthreads(); - - if (renormalize && lane_id == 0) { - float sum = 0.0f; - for (int k = 0; k < topk; k++) { - sum += output_ptr[row_idx * topk + k]; - } - - if (sum > 0.0f) { - for (int k = 0; k < topk; k++) { - int64_t idx = row_idx * topk + k; - output_ptr[idx] /= sum; - if (apply_routed_scaling_factor_on_output) { - output_ptr[idx] *= static_cast(routed_scaling_factor); - } + if (lane_id == 0) { + float rcp = 1.0f; + if (renormalize && sum_for_renorm > 0.0f) { + rcp = 1.0f / sum_for_renorm; + if (apply_routed_scaling_factor_on_output) { + rcp *= static_cast(routed_scaling_factor); } } + + for (int k = 0; k < topk; k++) { + int e = top_indices[k]; + bool valid = (e >= 0); + output_ptr[row_idx * topk + k] = valid ? warp_original_scores[e] * rcp : 0.0f; + indices_ptr[row_idx * topk + k] = valid ? e : 0; + } + } +} + +template +static void launch_for_n( + at::Tensor& input, + at::Tensor& bias, + at::Tensor& output, + at::Tensor& indices, + int64_t topk, + bool renormalize, + double routed_scaling_factor, + bool apply_routed_scaling_factor_on_output, + cudaStream_t stream) { + using Cfg = GateConfig; + int64_t num_rows = input.size(0); + bool use_small_token_kernel = num_rows <= Cfg::SMALL_TOKEN_THRESHOLD; + + if (use_small_token_kernel) { + dim3 grid(num_rows); + dim3 block(Cfg::THREADS_PER_BLOCK_SMALL); + kimi_k2_moe_fused_gate_kernel_small_token<<>>( + input.data_ptr(), + bias.data_ptr(), + output.data_ptr(), + indices.data_ptr(), + num_rows, + topk, + renormalize, + routed_scaling_factor, + apply_routed_scaling_factor_on_output); + } else { + int64_t num_blocks = (num_rows + Cfg::WARPS_PER_CTA - 1) / Cfg::WARPS_PER_CTA; + dim3 grid(num_blocks); + dim3 block(Cfg::WARP_SIZE, Cfg::WARPS_PER_CTA); + kimi_k2_moe_fused_gate_kernel<<>>( + input.data_ptr(), + bias.data_ptr(), + output.data_ptr(), + indices.data_ptr(), + num_rows, + topk, + renormalize, + routed_scaling_factor, + apply_routed_scaling_factor_on_output); } } @@ -267,9 +316,10 @@ std::vector kimi_k2_moe_fused_gate( int64_t num_rows = input.size(0); int32_t num_experts = input.size(1); - // Assert: Only support 384 experts - TORCH_CHECK(num_experts == 384, "kimi_k2_moe_fused_gate only supports 384 experts, but got ", num_experts); TORCH_CHECK(input.dtype() == bias.dtype(), "input and bias should have the same dtype"); + TORCH_CHECK(input.scalar_type() == at::kFloat, "kimi_k2_moe_fused_gate only supports float32 input"); + TORCH_CHECK(bias.scalar_type() == at::kFloat, "kimi_k2_moe_fused_gate only supports float32 bias"); + TORCH_CHECK(topk <= 8, "kimi_k2_moe_fused_gate only supports topk <= 8 (got ", topk, ")"); auto options = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); auto output = torch::empty({num_rows, topk}, options); @@ -277,42 +327,37 @@ std::vector kimi_k2_moe_fused_gate( const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - // Only support float32 - TORCH_CHECK(input.scalar_type() == at::kFloat, "kimi_k2_moe_fused_gate only supports float32 input"); - TORCH_CHECK(bias.scalar_type() == at::kFloat, "kimi_k2_moe_fused_gate only supports float32 bias"); - - bool use_small_token_kernel = num_rows <= SMALL_TOKEN_THRESHOLD; - - if (use_small_token_kernel) { - // Small token kernel: Each block handles 1 token with multiple warps collaborating - int64_t num_blocks = num_rows; - dim3 block_dim(THREADS_PER_BLOCK_SMALL); - - kimi_k2_moe_fused_gate_kernel_small_token<<>>( - input.data_ptr(), - bias.data_ptr(), - output.data_ptr(), - indices.data_ptr(), - num_rows, - topk, - renormalize, - routed_scaling_factor, - apply_routed_scaling_factor_on_output); - } else { - // Large token kernel: Original implementation - int64_t num_blocks = (num_rows + WARPS_PER_CTA - 1) / WARPS_PER_CTA; - dim3 block_dim(WARP_SIZE, WARPS_PER_CTA); - - kimi_k2_moe_fused_gate_kernel<<>>( - input.data_ptr(), - bias.data_ptr(), - output.data_ptr(), - indices.data_ptr(), - num_rows, - topk, - renormalize, - routed_scaling_factor, - apply_routed_scaling_factor_on_output); + switch (num_experts) { + case 256: + launch_for_n<256>( + input, + bias, + output, + indices, + topk, + renormalize, + routed_scaling_factor, + apply_routed_scaling_factor_on_output, + stream); + break; + case 384: + launch_for_n<384>( + input, + bias, + output, + indices, + topk, + renormalize, + routed_scaling_factor, + apply_routed_scaling_factor_on_output, + stream); + break; + default: + TORCH_CHECK( + false, + "kimi_k2_moe_fused_gate only supports num_experts in " + "{256, 384}, got ", + num_experts); } return {output, indices}; diff --git a/sgl-kernel/tests/test_kimi_k2_moe_fused_gate.py b/sgl-kernel/tests/test_kimi_k2_moe_fused_gate.py index b70dcd65b..1a42139e2 100644 --- a/sgl-kernel/tests/test_kimi_k2_moe_fused_gate.py +++ b/sgl-kernel/tests/test_kimi_k2_moe_fused_gate.py @@ -6,21 +6,26 @@ from sgl_kernel import kimi_k2_moe_fused_gate from sglang.srt.layers.moe.topk import kimi_k2_biased_topk_impl +# (num_experts, topk, routed_scaling_factor) +_CONFIGS = [ + (384, 6, 2.872), # Kimi K2 + (256, 8, 1.0), # MiMo V2.5 +] + @pytest.mark.parametrize( "seq_length", list(range(1, 10)) + [16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536], ) -@pytest.mark.parametrize("topk", [6]) # Kimi K2 uses topk=6 +@pytest.mark.parametrize("config", _CONFIGS, ids=["kimi384", "mimo256"]) @pytest.mark.parametrize("dtype", [torch.float32]) @pytest.mark.parametrize("apply_routed_scaling_factor_on_output", [False, True]) def test_kimi_k2_moe_fused_gate( - seq_length, topk, dtype, apply_routed_scaling_factor_on_output + seq_length, config, dtype, apply_routed_scaling_factor_on_output ): - num_experts = 384 # Kimi K2: only support 384 experts + num_experts, topk, routed_scaling_factor = config renormalize = True - routed_scaling_factor = 2.872 # Kimi K2's routed scaling factor torch.manual_seed(seq_length) tensor = torch.rand((seq_length, num_experts), dtype=dtype, device="cuda") @@ -65,13 +70,12 @@ def test_kimi_k2_moe_fused_gate( @pytest.mark.parametrize("seq_length", [1024, 4096]) -@pytest.mark.parametrize("num_experts", [384]) -@pytest.mark.parametrize("topk", [6]) -def test_kimi_k2_specific_case(seq_length, num_experts, topk): - """Test specifically for Kimi K2 configuration: 384 experts, topk=6""" +@pytest.mark.parametrize("config", _CONFIGS, ids=["kimi384", "mimo256"]) +def test_kimi_k2_specific_case(seq_length, config): + """Test specifically for supported configurations: 256 / 384 experts""" + num_experts, topk, routed_scaling_factor = config dtype = torch.float32 renormalize = True - routed_scaling_factor = 2.872 torch.manual_seed(42) tensor = torch.rand((seq_length, num_experts), dtype=dtype, device="cuda")