[opt kimi k2 3/n] opt kimi_k2 moe_fused_gate kernel (#13374)

sgl-kernel/csrc/moe/kimi_k2_moe_fused_gate.cu

@@ -1,15 +1,9 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <cuda_runtime.h>
-#include <cutlass/array.h>
-#include <cutlass/cutlass.h>
-#include <cutlass/numeric_types.h>
 #include <torch/all.h>
 
 #include <cfloat>
 
-using bfloat16_t = cutlass::bfloat16_t;
-using float16_t = cutlass::half_t;
-
 // Kimi K2 specific constants
 static constexpr int WARP_SIZE = 32;
 static constexpr int WARPS_PER_CTA = 6;
@@ -21,21 +15,13 @@ 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
 
-template <typename T>
+// Vectorization constants (used by large token kernel)
-__device__ inline bool cmp_gt(const T& a, const T& b) {
+static constexpr int VEC_SIZE = 4;  // Use float4 for vectorized loads
-  return static_cast<float>(a) > static_cast<float>(b);
-}
 
-template <typename T>
+// Small token optimized kernel: Each warp independently finds top-k, then merge, using warp-level topk
-__device__ inline bool cmp_eq(const T& a, const T& b) {
-  return static_cast<float>(a) == static_cast<float>(b);
-}
-
-// Small token optimized kernel: Multiple warps collaborate on a single token
-template <typename T>
 __global__ void kimi_k2_moe_fused_gate_kernel_small_token(
-    T* input,
+    float* input,
-    T* bias,
+    float* bias,
     float* output_ptr,
     int32_t* indices_ptr,
     int64_t num_rows,
@@ -43,7 +29,6 @@ __global__ void kimi_k2_moe_fused_gate_kernel_small_token(
     bool renormalize,
     double routed_scaling_factor,
     bool apply_routed_scaling_factor_on_output) {
-  // Each block handles one token with WARPS_PER_TOKEN_SMALL warps collaborating
   int64_t row_idx = blockIdx.x;
   if (row_idx >= num_rows) return;
 
@@ -51,98 +36,101 @@ __global__ void kimi_k2_moe_fused_gate_kernel_small_token(
   int warp_id = tid / WARP_SIZE;
   int lane_id = tid % WARP_SIZE;
 
-  // Shared memory for all warps to collaborate
+  // Shared memory: biased scores and original scores
   __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];
 
-  // Each thread loads one expert (384 threads for 384 experts)
+  // Load data: all 384 threads load one expert each
   if (tid < NUM_EXPERTS) {
-    T input_val = input[row_idx * NUM_EXPERTS + tid];
+    float input_val = input[row_idx * NUM_EXPERTS + tid];
-    T bias_val = bias[tid];
+    float bias_val = bias[tid];
-    float sigmoid_val = 1.0f / (1.0f + expf(-static_cast<float>(input_val)));
+    float sigmoid_val = 1.0f / (1.0f + expf(-input_val));
-    float biased_val = sigmoid_val + static_cast<float>(bias_val);
+    float biased_val = sigmoid_val + bias_val;
     shared_scores[tid] = biased_val;
     shared_original_scores[tid] = sigmoid_val;
   }
 
   __syncthreads();
 
-  // Parallel TopK: Each warp processes a portion of experts
+  // Find top-k using iterative selection, each iteration finds the next maximum
-  // Use multiple warps to find top-k elements in parallel
-  int experts_per_warp = (NUM_EXPERTS + WARPS_PER_TOKEN_SMALL - 1) / WARPS_PER_TOKEN_SMALL;
-  int warp_start = warp_id * experts_per_warp;
-  int warp_end = min(warp_start + experts_per_warp, NUM_EXPERTS);
-
   for (int k = 0; k < topk; k++) {
-    float max_val = -FLT_MAX;
+    // Each thread holds one expert's value
-    int max_expert = -1;
+    float my_val = (tid < NUM_EXPERTS) ? shared_scores[tid] : -FLT_MAX;
+    int my_expert = tid;
 
-    // Each warp finds the max in its portion
+    // Use warp-level reduction first
-    for (int expert = warp_start + lane_id; expert < warp_end; expert += WARP_SIZE) {
+    float warp_max_val = my_val;
-      float val = shared_scores[expert];
+    int warp_max_expert = my_expert;
-      if (val > max_val) {
+
-        max_val = val;
+#pragma unroll
-        max_expert = expert;
+    for (int offset = 16; offset > 0; offset /= 2) {
+      float other_val = __shfl_down_sync(0xFFFFFFFF, warp_max_val, offset);
+      int other_expert = __shfl_down_sync(0xFFFFFFFF, warp_max_expert, offset);
+      if (other_val > warp_max_val) {
+        warp_max_val = other_val;
+        warp_max_expert = other_expert;
       }
     }
 
-    // Warp-level reduction to find warp's maximum
+    // Warp leaders write to shared memory
-    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+    if (lane_id == 0) {
-      float other_val = __shfl_down_sync(0xFFFFFFFF, max_val, offset);
+      warp_maxs[warp_id] = warp_max_val;
-      int other_expert = __shfl_down_sync(0xFFFFFFFF, max_expert, offset);
+      warp_experts[warp_id] = warp_max_expert;
-
-      if (other_val > max_val || (other_val == max_val && other_expert < max_expert)) {
-        max_val = other_val;
-        max_expert = other_expert;
-      }
     }
 
-    // Store warp results in shared memory
+    __syncthreads();
-    __shared__ float warp_max_vals[WARPS_PER_TOKEN_SMALL];
+
-    __shared__ int warp_max_experts[WARPS_PER_TOKEN_SMALL];
+    // Final reduction among warps (done by first warp)
+    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);
+        int other_expert = __shfl_down_sync(0xFFFFFFFF, final_expert, offset);
+        if (other_val > final_max) {
+          final_max = other_val;
+          final_expert = other_expert;
+        }
+      }
 
       if (lane_id == 0) {
-      warp_max_vals[warp_id] = max_val;
+        selected_experts[k] = final_expert;
-      warp_max_experts[warp_id] = max_expert;
+        selected_vals[k] = final_max;
+      }
     }
 
     __syncthreads();
 
-    // First warp reduces across all warp results
+    // Mark the selected expert as used for next iteration
-    if (warp_id == 0) {
+    // All threads can read from selected_experts[k]
-      float final_max_val = -FLT_MAX;
+    int selected = selected_experts[k];
-      int final_max_expert = -1;
+    if (tid == selected) {
-
+      shared_scores[tid] = -FLT_MAX;
-      if (lane_id < WARPS_PER_TOKEN_SMALL) {
-        final_max_val = warp_max_vals[lane_id];
-        final_max_expert = warp_max_experts[lane_id];
-      }
-
-      // Warp reduction
-      for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
-        float other_val = __shfl_down_sync(0xFFFFFFFF, final_max_val, offset);
-        int other_expert = __shfl_down_sync(0xFFFFFFFF, final_max_expert, offset);
-
-        if (other_val > final_max_val || (other_val == final_max_val && other_expert < final_max_expert)) {
-          final_max_val = other_val;
-          final_max_expert = other_expert;
-        }
-      }
-
-      // Lane 0 writes result and marks the expert as used
-      if (lane_id == 0 && final_max_expert != -1) {
-        int64_t output_idx = row_idx * topk + k;
-        output_ptr[output_idx] = shared_original_scores[final_max_expert];
-        indices_ptr[output_idx] = final_max_expert;
-        shared_scores[final_max_expert] = -FLT_MAX;
-      }
     }
 
     __syncthreads();
   }
 
-  // Renormalization (only first warp)
+  // Write output (done by thread 0)
-  if (renormalize && warp_id == 0 && lane_id == 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;
+      }
+    }
+
+    // Renormalization
+    if (renormalize) {
       float sum = 0.0f;
       for (int k = 0; k < topk; k++) {
         sum += output_ptr[row_idx * topk + k];
@@ -159,11 +147,12 @@ __global__ void kimi_k2_moe_fused_gate_kernel_small_token(
       }
     }
   }
+}
 
-template <typename T>
+// Large token kernel: Original implementation with vectorized loads
 __global__ void kimi_k2_moe_fused_gate_kernel(
-    T* input,
+    float* input,
-    T* bias,
+    float* bias,
     float* output_ptr,
     int32_t* indices_ptr,
     int64_t num_rows,
@@ -183,14 +172,29 @@ __global__ void kimi_k2_moe_fused_gate_kernel(
   float* warp_scores = shared_scores + warp_id * NUM_EXPERTS;
   float* warp_original_scores = shared_original_scores + warp_id * NUM_EXPERTS;
 
-  for (int expert = lane_id; expert < NUM_EXPERTS; expert += WARP_SIZE) {
+  // Vectorized loading: each lane loads multiple float4 chunks
-    T input_val = input[row_idx * NUM_EXPERTS + expert];
+  // VPT = 12, so we load 12/4 = 3 float4 per lane
-    T bias_val = bias[expert];
+  const int VEC_PER_LANE = VPT / VEC_SIZE;  // 3
-    float sigmoid_val = 1.0f / (1.0f + expf(-static_cast<float>(input_val)));
+  float4* input_vec = reinterpret_cast<float4*>(input + row_idx * NUM_EXPERTS);
-    float biased_val = sigmoid_val + static_cast<float>(bias_val);
+  float4* bias_vec = reinterpret_cast<float4*>(bias);
+
+#pragma unroll
+  for (int i = 0; i < VEC_PER_LANE; i++) {
+    int vec_idx = lane_id * VEC_PER_LANE + i;
+    float4 input_val = input_vec[vec_idx];
+    float4 bias_val = bias_vec[vec_idx];
+
+#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;
     }
+  }
 
   __syncthreads();
 
@@ -265,6 +269,10 @@ 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) {
@@ -272,30 +280,7 @@ std::vector<at::Tensor> kimi_k2_moe_fused_gate(
     int64_t num_blocks = num_rows;
     dim3 block_dim(THREADS_PER_BLOCK_SMALL);
 
-    if (input.scalar_type() == at::kBFloat16) {
+    kimi_k2_moe_fused_gate_kernel_small_token<<<num_blocks, block_dim, 0, stream>>>(
-      kimi_k2_moe_fused_gate_kernel_small_token<bfloat16_t><<<num_blocks, block_dim, 0, stream>>>(
-          reinterpret_cast<bfloat16_t*>(input.data_ptr()),
-          reinterpret_cast<bfloat16_t*>(bias.data_ptr()),
-          output.data_ptr<float>(),
-          indices.data_ptr<int32_t>(),
-          num_rows,
-          topk,
-          renormalize,
-          routed_scaling_factor,
-          apply_routed_scaling_factor_on_output);
-    } else if (input.scalar_type() == at::kHalf) {
-      kimi_k2_moe_fused_gate_kernel_small_token<float16_t><<<num_blocks, block_dim, 0, stream>>>(
-          reinterpret_cast<float16_t*>(input.data_ptr()),
-          reinterpret_cast<float16_t*>(bias.data_ptr()),
-          output.data_ptr<float>(),
-          indices.data_ptr<int32_t>(),
-          num_rows,
-          topk,
-          renormalize,
-          routed_scaling_factor,
-          apply_routed_scaling_factor_on_output);
-    } else if (input.scalar_type() == at::kFloat) {
-      kimi_k2_moe_fused_gate_kernel_small_token<float><<<num_blocks, block_dim, 0, stream>>>(
         input.data_ptr<float>(),
         bias.data_ptr<float>(),
         output.data_ptr<float>(),
@@ -306,36 +291,11 @@ std::vector<at::Tensor> kimi_k2_moe_fused_gate(
         routed_scaling_factor,
         apply_routed_scaling_factor_on_output);
   } else {
-      TORCH_CHECK(false, "Unsupported data type for kimi_k2_moe_fused_gate");
+    // Large token kernel: Original implementation
-    }
-  } else {
     int64_t num_blocks = (num_rows + WARPS_PER_CTA - 1) / WARPS_PER_CTA;
     dim3 block_dim(WARP_SIZE, WARPS_PER_CTA);
 
-    if (input.scalar_type() == at::kBFloat16) {
+    kimi_k2_moe_fused_gate_kernel<<<num_blocks, block_dim, 0, stream>>>(
-      kimi_k2_moe_fused_gate_kernel<bfloat16_t><<<num_blocks, block_dim, 0, stream>>>(
-          reinterpret_cast<bfloat16_t*>(input.data_ptr()),
-          reinterpret_cast<bfloat16_t*>(bias.data_ptr()),
-          output.data_ptr<float>(),
-          indices.data_ptr<int32_t>(),
-          num_rows,
-          topk,
-          renormalize,
-          routed_scaling_factor,
-          apply_routed_scaling_factor_on_output);
-    } else if (input.scalar_type() == at::kHalf) {
-      kimi_k2_moe_fused_gate_kernel<float16_t><<<num_blocks, block_dim, 0, stream>>>(
-          reinterpret_cast<float16_t*>(input.data_ptr()),
-          reinterpret_cast<float16_t*>(bias.data_ptr()),
-          output.data_ptr<float>(),
-          indices.data_ptr<int32_t>(),
-          num_rows,
-          topk,
-          renormalize,
-          routed_scaling_factor,
-          apply_routed_scaling_factor_on_output);
-    } else if (input.scalar_type() == at::kFloat) {
-      kimi_k2_moe_fused_gate_kernel<float><<<num_blocks, block_dim, 0, stream>>>(
         input.data_ptr<float>(),
         bias.data_ptr<float>(),
         output.data_ptr<float>(),
@@ -345,9 +305,6 @@ std::vector<at::Tensor> kimi_k2_moe_fused_gate(
         renormalize,
         routed_scaling_factor,
         apply_routed_scaling_factor_on_output);
-    } else {
-      TORCH_CHECK(false, "Unsupported data type for kimi_k2_moe_fused_gate");
-    }
   }
 
   return {output, indices};