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

[MoE] Extend kimi_k2_moe_fused_gate to support 256 experts (MiMo V2 Flash) (#26303)

Browse Source
This commit is contained in:
Mengxuan Xiong
2026-06-01 16:03:58 +08:00
committed by GitHub
parent f60710a1d7
commit ff642ed936
2 changed files with 219 additions and 170 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/moe/kimi_k2_moe_fused_gate.cu
+206 -161
View File
@@ -4,21 +4,39 @@
#include <cfloat>
// 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 <int N>
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 <int N>
__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<N>;
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<float>(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<float>(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 <int N>
__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<N>;
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<float4*>(warp_scores);
float4* warp_original_scores_v4 = reinterpret_cast<float4*>(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<float4*>(input + row_idx * NUM_EXPERTS);
float4* bias_vec = reinterpret_cast<float4*>(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<float>(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<float>(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 <int N>
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<N>;
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<N><<<grid, block, 0, stream>>>(
input.data_ptr<float>(),
bias.data_ptr<float>(),
output.data_ptr<float>(),
indices.data_ptr<int32_t>(),
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<N><<<grid, block, 0, stream>>>(
input.data_ptr<float>(),
bias.data_ptr<float>(),
output.data_ptr<float>(),
indices.data_ptr<int32_t>(),
num_rows,
topk,
renormalize,
routed_scaling_factor,
apply_routed_scaling_factor_on_output);
}
}
@@ -267,9 +316,10 @@ std::vector<at::Tensor> 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<at::Tensor> 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<<<num_blocks, block_dim, 0, stream>>>(
input.data_ptr<float>(),
bias.data_ptr<float>(),
output.data_ptr<float>(),
indices.data_ptr<int32_t>(),
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<<<num_blocks, block_dim, 0, stream>>>(
input.data_ptr<float>(),
bias.data_ptr<float>(),
output.data_ptr<float>(),
indices.data_ptr<int32_t>(),
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};
sgl-kernel/tests/test_kimi_k2_moe_fused_gate.py
+13 -9
View File
@@ -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")