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

sgl-kernel/csrc/moe/kimi_k2_moe_fused_gate.cu

@@ -4,21 +4,39 @@
 
 #include <cfloat>
 
-// Kimi K2 specific constants
+// Kimi K2 MoE fused gate, supports NUM_EXPERTS in {256 (MiMo V2 Flash), 384 (Kimi K2)}.
-static constexpr int WARP_SIZE = 32;
+// Routing (DeepSeek "noaux_tc" with num_expert_group = 1):
-static constexpr int WARPS_PER_CTA = 6;
+//   1. sigmoid(gate_logit)
-static constexpr int NUM_EXPERTS = 384;
+//   2. add per-expert correction bias (ranking only)
-static constexpr int VPT = 12;  // 384 / 32 = 12
+//   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
+__device__ __forceinline__ float sigmoid_accurate(float x) {
-static constexpr int SMALL_TOKEN_THRESHOLD = 512;
+  return 1.0f / (1.0f + expf(-x));
-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
 
-// Vectorization constants (used by large token kernel)
+template <int N>
-static constexpr int VEC_SIZE = 4;  // Use float4 for vectorized loads
+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
+  // Sigmoid weights (no bias) for final lookup, indexed by expert id.
-  __shared__ float shared_scores[NUM_EXPERTS];
   __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
+  // Keep biased_val in register; mask the winner in-place each iteration to
-  if (tid < NUM_EXPERTS) {
+  // avoid round-tripping through shared memory.
-    float input_val = input[row_idx * NUM_EXPERTS + tid];
+  float input_val = input[row_idx * NUM_EXPERTS + tid];
-    float bias_val = bias[tid];
+  float bias_val = bias[tid];
-    float sigmoid_val = 1.0f / (1.0f + expf(-input_val));
+  float sigmoid_val = sigmoid_accurate(input_val);
-    float biased_val = sigmoid_val + bias_val;
+  float biased_val = sigmoid_val + bias_val;
-    shared_scores[tid] = biased_val;
+  shared_original_scores[tid] = sigmoid_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
+    // Stage 1: per-warp argmax.
-    float my_val = (tid < NUM_EXPERTS) ? shared_scores[tid] : -FLT_MAX;
+    float warp_max_val = biased_val;
-    int my_expert = tid;
+    int warp_max_expert = tid;
-
-    // Use warp-level reduction first
-    float warp_max_val = my_val;
-    int warp_max_expert = my_expert;
-
 #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;
+        if (renormalize && final_expert >= 0 && final_expert < NUM_EXPERTS) {
-      }
+          sum_for_renorm += shared_original_scores[final_expert];
-    }
-
-    __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);
-          }
         }
       }
     }
+    __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 sigmoid_val = sigmoid_accurate(inp);
-      float biased_val = sigmoid_val + b;
+      ((float*)&sigmoid_v4)[j] = sigmoid_val;
-      warp_scores[expert] = biased_val;
+      ((float*)&biased_v4)[j] = sigmoid_val + b;
-      warp_original_scores[expert] = sigmoid_val;
     }
+    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;
+      bool valid = (max_expert >= 0 && max_expert < NUM_EXPERTS);
-      if (max_expert != -1) {
+      top_indices[k] = valid ? max_expert : -1;
-        output_ptr[output_idx] = warp_original_scores[max_expert];
+      if (renormalize && valid) {
-        indices_ptr[output_idx] = max_expert;
+        sum_for_renorm += warp_original_scores[max_expert];
-        warp_scores[max_expert] = -FLT_MAX;
-      } else {
-        output_ptr[output_idx] = 0.0f;
-        indices_ptr[output_idx] = 0;
       }
+      if (valid) warp_scores[max_expert] = -FLT_MAX;
     }
-
     __syncwarp();
   }
 
-  __syncthreads();
+  if (lane_id == 0) {
-
+    float rcp = 1.0f;
-  if (renormalize && lane_id == 0) {
+    if (renormalize && sum_for_renorm > 0.0f) {
-    float sum = 0.0f;
+      rcp = 1.0f / sum_for_renorm;
-    for (int k = 0; k < topk; k++) {
+      if (apply_routed_scaling_factor_on_output) {
-      sum += output_ptr[row_idx * topk + k];
+        rcp *= static_cast<float>(routed_scaling_factor);
-    }
-
-    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);
-        }
       }
     }
+
+    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
+  switch (num_experts) {
-  TORCH_CHECK(input.scalar_type() == at::kFloat, "kimi_k2_moe_fused_gate only supports float32 input");
+    case 256:
-  TORCH_CHECK(bias.scalar_type() == at::kFloat, "kimi_k2_moe_fused_gate only supports float32 bias");
+      launch_for_n<256>(
-
+          input,
-  bool use_small_token_kernel = num_rows <= SMALL_TOKEN_THRESHOLD;
+          bias,
-
+          output,
-  if (use_small_token_kernel) {
+          indices,
-    // Small token kernel: Each block handles 1 token with multiple warps collaborating
+          topk,
-    int64_t num_blocks = num_rows;
+          renormalize,
-    dim3 block_dim(THREADS_PER_BLOCK_SMALL);
+          routed_scaling_factor,
-
+          apply_routed_scaling_factor_on_output,
-    kimi_k2_moe_fused_gate_kernel_small_token<<<num_blocks, block_dim, 0, stream>>>(
+          stream);
-        input.data_ptr<float>(),
+      break;
-        bias.data_ptr<float>(),
+    case 384:
-        output.data_ptr<float>(),
+      launch_for_n<384>(
-        indices.data_ptr<int32_t>(),
+          input,
-        num_rows,
+          bias,
-        topk,
+          output,
-        renormalize,
+          indices,
-        routed_scaling_factor,
+          topk,
-        apply_routed_scaling_factor_on_output);
+          renormalize,
-  } else {
+          routed_scaling_factor,
-    // Large token kernel: Original implementation
+          apply_routed_scaling_factor_on_output,
-    int64_t num_blocks = (num_rows + WARPS_PER_CTA - 1) / WARPS_PER_CTA;
+          stream);
-    dim3 block_dim(WARP_SIZE, WARPS_PER_CTA);
+      break;
-
+    default:
-    kimi_k2_moe_fused_gate_kernel<<<num_blocks, block_dim, 0, stream>>>(
+      TORCH_CHECK(
-        input.data_ptr<float>(),
+          false,
-        bias.data_ptr<float>(),
+          "kimi_k2_moe_fused_gate only supports num_experts in "
-        output.data_ptr<float>(),
+          "{256, 384}, got ",
-        indices.data_ptr<int32_t>(),
+          num_experts);
-        num_rows,
-        topk,
-        renormalize,
-        routed_scaling_factor,
-        apply_routed_scaling_factor_on_output);
   }
 
   return {output, indices};

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("config", _CONFIGS, ids=["kimi384", "mimo256"])
-@pytest.mark.parametrize("topk", [6])
+def test_kimi_k2_specific_case(seq_length, config):
-def test_kimi_k2_specific_case(seq_length, num_experts, topk):
+    """Test specifically for supported configurations: 256 / 384 experts"""
-    """Test specifically for Kimi K2 configuration: 384 experts, topk=6"""
+    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")