[JIT Kernel] Migrate moe_topk_softmax from AOT to JIT (#34509)

Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>

python/sglang/kernels/jit/csrc/moe/moe_topk_softmax.cuh

@@ -0,0 +1,685 @@
+// Adapt from https://github.com/vllm-project/vllm/blob/v0.7.3/csrc/moe/topk_softmax_kernels.cu
+// which is originally adapted from
+// https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu
+//
+// JIT port: the CUDA kernels are unchanged from the sgl-kernel implementation; only
+// the host launcher is adapted to the tvm-ffi TensorView API. The softmax workspace
+// (needed for non-power-of-two / >256 expert counts) is allocated by the Python
+// wrapper and passed in.
+#pragma once
+
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize, SymbolicDevice
+#include <sgl_kernel/utils.h>   // For RuntimeCheck
+
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel, fp16_t/bf16_t/fp32_t
+
+#include <cub/cub.cuh>
+#include <cub/util_type.cuh>
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/optional.h>
+
+// CUDA 12.9+ deprecated cub::Max/Min in favour of cuda::maximum/minimum
+#if CUDA_VERSION >= 12090
+#include <cuda/functional>
+#endif
+
+#include <cfloat>
+#include <cstdint>
+#include <type_traits>
+
+namespace sglang {
+
+// Namespaced so the macro cannot collide with a WARP_SIZE defined by another
+// translation unit pulled into the same JIT module.
+#define MOE_TOPK_SOFTMAX_WARP_SIZE 32
+static constexpr int WARP_SIZE = MOE_TOPK_SOFTMAX_WARP_SIZE;
+
+#define SGL_MAX(a, b) ((a) > (b) ? (a) : (b))
+#define SGL_MIN(a, b) ((a) < (b) ? (a) : (b))
+
+#define SGLANG_SHFL_XOR_SYNC_WIDTH(mask, var, lane_mask, width) __shfl_xor_sync(mask, var, lane_mask, width)
+
+// CUDA 13 (12.9+) deprecated cub::Max/Min in favor of cuda::maximum/minimum.
+#if CUDA_VERSION >= 12090
+using MaxReduceOp = cuda::maximum<>;
+using MinReduceOp = cuda::minimum<>;
+#else
+using MaxReduceOp = cub::Max;
+using MinReduceOp = cub::Min;
+#endif
+
+using cub_kvp = cub::KeyValuePair<int, float>;
+
+/// Aligned array type
+template <typename T, int N, int Alignment = sizeof(T) * N>
+class alignas(Alignment) AlignedArray {
+  T data[N];
+};
+
+template <typename T>
+__device__ float convert_to_float(T x) {
+  if constexpr (std::is_same_v<T, __half>) {
+    return __half2float(x);
+  } else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
+    return __bfloat162float(x);
+  } else if constexpr (std::is_same_v<T, float>) {
+    return x;
+  } else {
+    return static_cast<float>(x);
+  }
+}
+
+// ====================== Softmax things ===============================
+template <typename T, int TPB>
+__launch_bounds__(TPB) __global__ void moeSoftmax(
+    const T* input,
+    const bool* finished,
+    float* output,
+    const int num_cols,
+    const float moe_softcapping,
+    const float* correction_bias) {
+  using BlockReduce = cub::BlockReduce<float, TPB>;
+  __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+  __shared__ float normalizing_factor;
+  __shared__ float float_max;
+
+  const int thread_row_offset = blockIdx.x * num_cols;
+
+  float threadData(-FLT_MAX);
+
+  if ((finished != nullptr) && finished[blockIdx.x]) {
+    return;
+  }
+
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    float val = convert_to_float<T>(input[idx]);
+
+    if (moe_softcapping != 0.0f) {
+      val = tanhf(val / moe_softcapping) * moe_softcapping;
+    }
+
+    if (correction_bias != nullptr) {
+      val = val + correction_bias[ii];
+    }
+
+    output[idx] = val;
+    threadData = max(val, threadData);
+  }
+
+  const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, MaxReduceOp());
+
+  if (threadIdx.x == 0) {
+    float_max = maxElem;
+  }
+  __syncthreads();
+
+  threadData = 0;
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    threadData += exp((output[idx] - float_max));
+  }
+
+  const auto Z = BlockReduce(tmpStorage).Sum(threadData);
+
+  if (threadIdx.x == 0) {
+    normalizing_factor = 1.f / Z;
+  }
+  __syncthreads();
+
+  for (int ii = threadIdx.x; ii < num_cols; ii += TPB) {
+    const int idx = thread_row_offset + ii;
+    const float softmax_val = exp((output[idx] - float_max)) * normalizing_factor;
+    output[idx] = softmax_val;
+  }
+}
+
+namespace moe {
+struct TopKPair {
+  static const int PAIR = 2;
+  static const int MAX_INDEX = 0;
+  cub_kvp max;
+  cub_kvp secondMax;
+
+  __device__ TopKPair() {}
+  __device__ TopKPair(cub_kvp max, cub_kvp secondMax) : max(max), secondMax(secondMax) {}
+};
+
+struct TopKPairArgMax {
+  __device__ TopKPairArgMax() {}
+  __device__ __forceinline__ TopKPair operator()(const TopKPair& candidate1, const TopKPair& candidate2) const {
+    cub_kvp globalMax, globalSecondMax;
+
+    if (candidate1.max.value > candidate2.max.value) {
+      globalMax = candidate1.max;
+    } else {
+      globalMax = candidate2.max;
+    }
+
+    if (globalMax.key == candidate1.max.key) {
+      globalSecondMax = (candidate1.secondMax.value > candidate2.max.value) ? candidate1.secondMax : candidate2.max;
+    } else {
+      globalSecondMax = (candidate2.secondMax.value > candidate1.max.value) ? candidate2.secondMax : candidate1.max;
+    }
+    return TopKPair(globalMax, globalSecondMax);
+  }
+};
+}  // namespace moe
+
+template <int TPB>
+__launch_bounds__(TPB) __global__ void moeTopKFast(
+    float* inputs_after_softmax,
+    const bool* finished,
+    float* output,
+    int* indices,
+    const int num_experts,
+    const int k,
+    const int start_expert,
+    const int end_expert,
+    const bool renormalize) {
+  using namespace moe;
+  using BlockReduce = cub::BlockReduce<TopKPair, TPB>;
+  __shared__ typename BlockReduce::TempStorage tmpStorage;
+  TopKPair thread_pair;
+
+  const int block_row = blockIdx.x;
+
+  const bool row_is_active = finished ? !finished[block_row] : true;
+  const int thread_read_offset = blockIdx.x * num_experts;
+  float row_sum_for_renormalize = 0;
+  for (int k_idx = 0; k_idx < (k + TopKPair::PAIR - 1) / TopKPair::PAIR; ++k_idx) {
+    thread_pair.max.key = 0;
+    thread_pair.max.value = -1.f;
+    thread_pair.secondMax.key = 0;
+    thread_pair.secondMax.value = -1.f;
+
+    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;
+      inp_kvp.value = inputs_after_softmax[idx];
+      if (inp_kvp.value > thread_pair.max.value) {
+        thread_pair.secondMax = thread_pair.max;
+        thread_pair.max = inp_kvp;
+      } else if (inp_kvp.value > thread_pair.secondMax.value) {
+        thread_pair.secondMax = inp_kvp;
+      }
+    }
+
+    TopKPairArgMax reducer;
+    const TopKPair result_pair = BlockReduce(tmpStorage).Reduce(thread_pair, reducer);
+    if (threadIdx.x == 0) {
+#pragma unroll
+      for (int i = 0; i < TopKPair::PAIR; i++) {
+        if (k_idx * 2 + i >= k) break;
+        cub_kvp result = (i == TopKPair::MAX_INDEX) ? result_pair.max : result_pair.secondMax;
+        int expert = result.key;
+        bool node_uses_expert = expert >= start_expert && expert < end_expert;
+        bool should_process_row = row_is_active && node_uses_expert;
+        inputs_after_softmax[thread_read_offset + expert] = -1.f;
+        int idx = k * block_row + k_idx * 2 + i;
+        output[idx] = result.value;
+        indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
+        assert(indices[idx] >= 0);
+        row_sum_for_renormalize += result.value;
+      }
+    }
+    __syncthreads();
+  }
+
+  if (renormalize && threadIdx.x == 0) {
+    float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
+    for (int k_idx = 0; k_idx < k; ++k_idx) {
+      const int idx = k * block_row + k_idx;
+      output[idx] = output[idx] * row_sum_for_renormalize_inv;
+    }
+  }
+}
+
+template <int TPB>
+__launch_bounds__(TPB) __global__ void moeTopK(
+    float* inputs_after_softmax,
+    const bool* finished,
+    float* output,
+    int* indices,
+    const int num_experts,
+    const int k,
+    const int start_expert,
+    const int end_expert,
+    const bool renormalize) {
+  using BlockReduce = cub::BlockReduce<cub_kvp, TPB>;
+  __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+  cub_kvp thread_kvp;
+  cub::ArgMax arg_max;
+
+  const int block_row = blockIdx.x;
+
+  const bool row_is_active = finished ? !finished[block_row] : true;
+  const int thread_read_offset = blockIdx.x * num_experts;
+  float row_sum_for_renormalize = 0;
+  for (int k_idx = 0; k_idx < k; ++k_idx) {
+    thread_kvp.key = 0;
+    thread_kvp.value = -1.f;
+
+    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;
+      inp_kvp.value = inputs_after_softmax[idx];
+      thread_kvp = arg_max(inp_kvp, thread_kvp);
+    }
+
+    const cub_kvp result_kvp = BlockReduce(tmpStorage).Reduce(thread_kvp, arg_max);
+    if (threadIdx.x == 0) {
+      const int expert = result_kvp.key;
+      const bool node_uses_expert = expert >= start_expert && expert < end_expert;
+      const bool should_process_row = row_is_active && node_uses_expert;
+
+      const int idx = k * block_row + k_idx;
+      output[idx] = result_kvp.value;
+      indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
+      assert(indices[idx] >= 0);
+      row_sum_for_renormalize += result_kvp.value;
+      inputs_after_softmax[thread_read_offset + expert] = -1.f;
+    }
+    __syncthreads();
+  }
+
+  if (renormalize && threadIdx.x == 0) {
+    float row_sum_for_renormalize_inv = 1.f / row_sum_for_renormalize;
+    for (int k_idx = 0; k_idx < k; ++k_idx) {
+      const int idx = k * block_row + k_idx;
+      output[idx] = output[idx] * row_sum_for_renormalize_inv;
+    }
+  }
+}
+
+// ====================== TopK softmax things ===============================
+template <typename T, int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG>
+__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__ void topkGatingSoftmax(
+    const T* input,
+    const bool* finished,
+    float* output,
+    const int num_rows,
+    int* indices,
+    const int k,
+    const int start_expert,
+    const int end_expert,
+    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_assert(VPT % ELTS_PER_LDG == 0, "The elements per thread must be a multiple of the elements per ldg");
+  static_assert(WARP_SIZE % THREADS_PER_ROW == 0, "The threads per row must cleanly divide the threads per warp");
+  static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW), "THREADS_PER_ROW must be power of 2");
+  static_assert(THREADS_PER_ROW <= WARP_SIZE, "THREADS_PER_ROW can be at most warp size");
+
+  static constexpr int ELTS_PER_WARP = WARP_SIZE * 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_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0, "The elts per row must cleanly divide the total elt per warp");
+
+  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 bool row_is_active = finished ? !finished[thread_row] : true;
+
+  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, ELTS_PER_LDG>;
+
+  T row_chunk_temp[VPT];
+  AccessType* row_chunk_vec_ptr = reinterpret_cast<AccessType*>(&row_chunk_temp);
+  const AccessType* vec_thread_read_ptr = reinterpret_cast<const AccessType*>(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] = convert_to_float<T>(row_chunk_temp[ii]);
+  }
+
+  if (moe_softcapping != 0.0f || correction_bias != nullptr) {
+#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;
+      }
+
+      if (correction_bias != nullptr) {
+        const int group_id = ii / ELTS_PER_LDG;
+        const int local_id = ii % ELTS_PER_LDG;
+        const int expert_idx = first_elt_read_by_thread + group_id * THREADS_PER_ROW * ELTS_PER_LDG + local_id;
+        val = val + correction_bias[expert_idx];
+      }
+
+      row_chunk[ii] = val;
+    }
+  }
+
+  float thread_max = row_chunk[0];
+#pragma unroll
+  for (int ii = 1; ii < VPT; ++ii) {
+    thread_max = max(thread_max, row_chunk[ii]);
+  }
+
+#pragma unroll
+  for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+    thread_max = max(thread_max, SGLANG_SHFL_XOR_SYNC_WIDTH(0xffffffff, thread_max, mask, THREADS_PER_ROW));
+  }
+
+  float row_sum = 0;
+#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 += SGLANG_SHFL_XOR_SYNC_WIDTH(0xffffffff, row_sum, mask, THREADS_PER_ROW);
+  }
+
+  const float reciprocal_row_sum = 1.f / row_sum;
+
+#pragma unroll
+  for (int ii = 0; ii < VPT; ++ii) {
+    row_chunk[ii] = row_chunk[ii] * reciprocal_row_sum;
+  }
+
+  int start_col = first_elt_read_by_thread;
+  static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
+
+  float row_sum_for_renormalize = 0;
+
+  for (int k_idx = 0; k_idx < k; ++k_idx) {
+    float max_val = row_chunk[0];
+    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) {
+        float val = row_chunk[ldg * ELTS_PER_LDG + ii];
+
+        if (val > max_val) {
+          max_val = val;
+          expert = col + ii;
+        }
+      }
+    }
+
+#pragma unroll
+    for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2) {
+      float other_max = SGLANG_SHFL_XOR_SYNC_WIDTH(0xffffffff, max_val, mask, THREADS_PER_ROW);
+      int other_expert = SGLANG_SHFL_XOR_SYNC_WIDTH(0xffffffff, expert, mask, THREADS_PER_ROW);
+
+      if (other_max > max_val || (other_max == max_val && other_expert < expert)) {
+        max_val = other_max;
+        expert = other_expert;
+      }
+    }
+
+    if (thread_group_idx == 0) {
+      const bool node_uses_expert = expert >= start_expert && expert < end_expert;
+      const bool should_process_row = row_is_active && node_uses_expert;
+
+      const int idx = k * thread_row + k_idx;
+      output[idx] = max_val;
+      indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
+      row_sum_for_renormalize += max_val;
+    }
+
+    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] = -10000.f;
+      }
+    }
+  }
+
+  if (renormalize && thread_group_idx == 0) {
+    float row_sum_for_renormalize_inv = 1.f / 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] = output[idx] * row_sum_for_renormalize_inv;
+    }
+  }
+}
+
+namespace detail {
+template <typename T, int EXPERTS, int BYTES_PER_LDG>
+struct TopkConstants {
+  static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
+  static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE) == 0 || EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0, "");
+  static constexpr int VECs_PER_THREAD = SGL_MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE));
+  static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
+  static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
+  static constexpr int ROWS_PER_WARP = WARP_SIZE / THREADS_PER_ROW;
+};
+}  // namespace detail
+
+template <typename T, int EXPERTS, int WARPS_PER_TB>
+void topkGatingSoftmaxLauncherHelper(
+    const T* input,
+    const bool* finished,
+    float* output,
+    int* indices,
+    const int num_rows,
+    const int k,
+    const int start_expert,
+    const int end_expert,
+    const bool renormalize,
+    const float moe_softcapping,
+    const float* correction_bias,
+    cudaStream_t stream) {
+  static constexpr std::size_t MAX_BYTES_PER_LDG = 16;
+
+  static constexpr int BYTES_PER_LDG = SGL_MIN(MAX_BYTES_PER_LDG, sizeof(T) * EXPERTS);
+  using Constants = detail::TopkConstants<T, EXPERTS, BYTES_PER_LDG>;
+  static constexpr int VPT = Constants::VPT;
+  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(WARP_SIZE, WARPS_PER_TB);
+  host::LaunchKernel(dim3(num_blocks), block_dim, stream)(
+      topkGatingSoftmax<T, VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG>,
+      input,
+      finished,
+      output,
+      num_rows,
+      indices,
+      k,
+      start_expert,
+      end_expert,
+      renormalize,
+      moe_softcapping,
+      correction_bias);
+}
+
+#define LAUNCH_SOFTMAX(TYPE, NUM_EXPERTS, WARPS_PER_TB)             \
+  topkGatingSoftmaxLauncherHelper<TYPE, NUM_EXPERTS, WARPS_PER_TB>( \
+      gating_output,                                                \
+      nullptr,                                                      \
+      topk_weights,                                                 \
+      topk_indices,                                                 \
+      num_tokens,                                                   \
+      topk,                                                         \
+      0,                                                            \
+      num_experts,                                                  \
+      renormalize,                                                  \
+      moe_softcapping,                                              \
+      correction_bias,                                              \
+      stream);
+
+template <typename T>
+void topkGatingSoftmaxKernelLauncher(
+    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,
+    cudaStream_t stream) {
+  static constexpr int WARPS_PER_TB = 4;
+  switch (num_experts) {
+    case 1:
+      LAUNCH_SOFTMAX(T, 1, WARPS_PER_TB);
+      break;
+    case 2:
+      LAUNCH_SOFTMAX(T, 2, WARPS_PER_TB);
+      break;
+    case 4:
+      LAUNCH_SOFTMAX(T, 4, WARPS_PER_TB);
+      break;
+    case 8:
+      LAUNCH_SOFTMAX(T, 8, WARPS_PER_TB);
+      break;
+    case 16:
+      LAUNCH_SOFTMAX(T, 16, WARPS_PER_TB);
+      break;
+    case 32:
+      LAUNCH_SOFTMAX(T, 32, WARPS_PER_TB);
+      break;
+    case 64:
+      LAUNCH_SOFTMAX(T, 64, WARPS_PER_TB);
+      break;
+    case 128:
+      LAUNCH_SOFTMAX(T, 128, WARPS_PER_TB);
+      break;
+    case 256:
+      LAUNCH_SOFTMAX(T, 256, WARPS_PER_TB);
+      break;
+    case 512:
+      LAUNCH_SOFTMAX(T, 512, WARPS_PER_TB);
+      break;
+    default: {
+      host::RuntimeCheck(
+          softmax_workspace != nullptr,
+          "softmax_workspace must be provided for num_experts that are not a power of 2.");
+      static constexpr int TPB = 256;
+      host::LaunchKernel(dim3(num_tokens), dim3(TPB), stream)(
+          moeSoftmax<T, TPB>,
+          gating_output,
+          (const bool*)nullptr,
+          softmax_workspace,
+          num_experts,
+          moe_softcapping,
+          correction_bias);
+      if (topk == 1) {
+        host::LaunchKernel(dim3(num_tokens), dim3(TPB), stream)(
+            moeTopK<TPB>,
+            softmax_workspace,
+            (const bool*)nullptr,
+            topk_weights,
+            topk_indices,
+            num_experts,
+            topk,
+            0,
+            num_experts,
+            renormalize);
+      } else {
+        host::LaunchKernel(dim3(num_tokens), dim3(TPB), stream)(
+            moeTopKFast<TPB>,
+            softmax_workspace,
+            (const bool*)nullptr,
+            topk_weights,
+            topk_indices,
+            num_experts,
+            topk,
+            0,
+            num_experts,
+            renormalize);
+      }
+    }
+  }
+}
+
+template <typename T>
+void topk_softmax(
+    tvm::ffi::TensorView topk_weights,       // [num_tokens, topk]  (float32, out)
+    tvm::ffi::TensorView topk_indices,       // [num_tokens, topk]  (int32, out)
+    tvm::ffi::TensorView gating_output,      // [num_tokens, num_experts] (T, in)
+    tvm::ffi::TensorView softmax_workspace,  // [num_tokens * num_experts] or [0] (float32)
+    bool renormalize,
+    double moe_softcapping,
+    tvm::ffi::Optional<tvm::ffi::TensorView> correction_bias) {
+  using namespace host;
+
+  auto device = SymbolicDevice{};
+  auto N = SymbolicSize{"num_tokens"};
+  auto E = SymbolicSize{"num_experts"};
+  auto K = SymbolicSize{"topk"};
+  device.set_options<kDLCUDA>();
+
+  TensorMatcher({N, E}).with_dtype<T>().with_device(device).verify(gating_output);
+  TensorMatcher({N, K}).with_dtype<fp32_t>().with_device(device).verify(topk_weights);
+  TensorMatcher({N, K}).with_dtype<int32_t>().with_device(device).verify(topk_indices);
+
+  const int num_tokens = static_cast<int>(N.unwrap());
+  const int num_experts = static_cast<int>(E.unwrap());
+  const int topk = static_cast<int>(K.unwrap());
+
+  if (num_tokens == 0) return;
+
+  const float* bias_ptr = nullptr;
+  if (correction_bias.has_value()) {
+    const auto& bias = correction_bias.value();
+    RuntimeCheck(bias.dtype().code == kDLFloat && bias.dtype().bits == 32, "correction_bias must be float32");
+    bias_ptr = static_cast<const float*>(bias.data_ptr());
+  }
+
+  float* workspace_ptr = (softmax_workspace.numel() > 0) ? static_cast<float*>(softmax_workspace.data_ptr()) : nullptr;
+
+  const cudaStream_t stream = LaunchKernel::resolve_device(device.unwrap());
+  const float moe_softcapping_f = static_cast<float>(moe_softcapping);
+
+  topkGatingSoftmaxKernelLauncher<T>(
+      static_cast<const T*>(gating_output.data_ptr()),
+      static_cast<float*>(topk_weights.data_ptr()),
+      static_cast<int*>(topk_indices.data_ptr()),
+      workspace_ptr,
+      num_tokens,
+      num_experts,
+      topk,
+      renormalize,
+      moe_softcapping_f,
+      bias_ptr,
+      stream);
+}
+
+}  // namespace sglang

python/sglang/kernels/ops/moe/__init__.py

@@ -17,6 +17,7 @@ if TYPE_CHECKING:
     import torch
 
 _CUDA = frozenset({CapabilityRequirement.CUDA})
+_HIP = frozenset({CapabilityRequirement.HIP})
 
 register_kernel(
     KernelSpec(
@@ -48,11 +49,25 @@ register_kernel(
         op="moe.topk_softmax",
         backend=KernelBackend.AOT,
         target="sgl_kernel:topk_softmax",
+        capabilities=_HIP,
         format_signature=FormatSignature(
             in_place=True,
             description="top-k softmax routing weights/ids",
         ),
-        description="MoE top-k softmax (sgl_kernel wheel).",
+        description="MoE top-k softmax (sgl_kernel ROCm wheel).",
+    )
+)
+register_kernel(
+    KernelSpec(
+        op="moe.topk_softmax",
+        backend=KernelBackend.JIT,
+        target="sglang.kernels.ops.moe.moe_topk_softmax:topk_softmax",
+        capabilities=_CUDA,
+        format_signature=FormatSignature(
+            in_place=True,
+            description="top-k softmax routing weights/ids",
+        ),
+        description="MoE top-k softmax (sglang.kernels.jit).",
     )
 )
 
@@ -103,7 +118,7 @@ def topk_softmax(
     correction_bias: Optional[torch.Tensor] = None,
 ) -> None:
     """Compute top-k softmax routing weights/ids for MoE."""
-    return get_kernel("moe.topk_softmax", KernelBackend.AOT)(
+    return get_kernel("moe.topk_softmax")(
         topk_weights,
         topk_ids,
         gating_output,

python/sglang/kernels/ops/moe/moe_topk_softmax.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
+from sglang.srt.utils.custom_op import register_custom_op
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_moe_topk_softmax_module(dtype: torch.dtype) -> Module:
+    args = make_cpp_args(dtype)
+    return load_jit(
+        "moe_topk_softmax",
+        *args,
+        cuda_files=["moe/moe_topk_softmax.cuh"],
+        cuda_wrappers=[("topk_softmax", f"topk_softmax<{args}>")],
+        extra_cuda_cflags=["--use_fast_math"],
+    )
+
+
+@register_custom_op(
+    op_name="moe_topk_softmax_out",
+    mutates_args=["topk_weights", "topk_ids", "workspace"],
+)
+def moe_topk_softmax_out(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    gating_output: torch.Tensor,
+    workspace: torch.Tensor,
+    renormalize: bool,
+    moe_softcapping: float,
+    correction_bias: Optional[torch.Tensor],
+) -> None:
+    """
+    Fused softmax top-k MoE gate (destination-passing style).
+
+    Args:
+        topk_weights:  [num_tokens, topk], float32, pre-allocated output
+        topk_ids:      [num_tokens, topk], int32,   pre-allocated output
+        gating_output: [num_tokens, num_experts], fp32/fp16/bf16
+        workspace:     [num_tokens * num_experts] float32 scratch (may be size 1
+                       when num_experts is a supported power-of-2 <= 256)
+        renormalize:   whether to renormalize weights to sum to 1 per row
+        moe_softcapping: tanh softcapping value applied to the logits, 0 disables
+        correction_bias: [num_experts] float32 per-expert bias, or None
+    """
+    module = _jit_moe_topk_softmax_module(gating_output.dtype)
+    module.topk_softmax(
+        topk_weights,
+        topk_ids,
+        gating_output,
+        workspace,
+        renormalize,
+        moe_softcapping,
+        correction_bias,
+    )
+
+
+def topk_softmax(
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    gating_output: torch.Tensor,
+    renormalize: bool = False,
+    moe_softcapping: float = 0.0,
+    correction_bias: Optional[torch.Tensor] = None,
+) -> None:
+    """
+    Fused softmax top-k MoE gate with the same call signature as
+    ``sgl_kernel.topk_softmax`` (destination-passing, in-place).
+
+    Args:
+        topk_weights:  [num_tokens, topk] float32, written in-place
+        topk_ids:      [num_tokens, topk] int32,   written in-place
+        gating_output: [num_tokens, num_experts] fp32/fp16/bf16
+        renormalize:   whether to renormalize weights to sum to 1 per row
+        moe_softcapping: tanh softcapping value applied to the logits, 0 disables
+        correction_bias: [num_experts] float32 per-expert bias, or None
+    """
+    num_tokens = gating_output.shape[0]
+    num_experts = gating_output.shape[1]
+
+    # The warp-specialized fast path covers power-of-two expert counts up to 512
+    # and needs no scratch; everything else goes through the two-pass
+    # softmax + top-k path, which writes probabilities to the workspace first.
+    # This threshold must stay in sync with the `case 512:` in the .cuh
+    # dispatcher and with the AOT host entry in
+    # kernels/aot/csrc/moe/moe_topk_softmax_kernels.cu.
+    is_pow2 = num_experts != 0 and (num_experts & (num_experts - 1)) == 0
+    needs_workspace = not is_pow2 or num_experts > 512
+    workspace_size = num_tokens * num_experts if needs_workspace else 1
+    workspace = torch.empty(
+        workspace_size, dtype=torch.float32, device=gating_output.device
+    )
+
+    moe_topk_softmax_out(
+        topk_weights,
+        topk_ids,
+        gating_output,
+        workspace,
+        renormalize,
+        moe_softcapping,
+        correction_bias,
+    )

test/registered/kernels/benchmark/moe/bench_moe_topk_softmax.py

@@ -0,0 +1,65 @@
+import torch
+from sgl_kernel import topk_softmax as aot_topk_softmax
+
+from sglang.kernels.jit.benchmark import marker
+from sglang.kernels.jit.benchmark.utils import create_random
+from sglang.kernels.ops.moe.moe_topk_softmax import topk_softmax as jit_topk_softmax
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(
+    est_time=20, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)
+
+
+def _jit(topk_weights, topk_ids, gating_output):
+    jit_topk_softmax(topk_weights, topk_ids, gating_output)
+
+
+def _aot(topk_weights, topk_ids, gating_output):
+    aot_topk_softmax(
+        topk_weights=topk_weights, topk_ids=topk_ids, gating_output=gating_output
+    )
+
+
+def _torch(topk_weights, topk_ids, gating_output):
+    probs = torch.softmax(gating_output.float(), dim=-1)
+    return probs.topk(topk_weights.shape[-1], dim=-1)
+
+
+FN_MAP = {
+    "jit": _jit,
+    "aot": _aot,
+    "torch": _torch,
+}
+
+
+# 32/128/256/512 take the warp-specialized power-of-two path; 12/160 fall back to
+# the two-pass path through the softmax workspace. 512 is kept because it sits on
+# the boundary between the two.
+@marker.parametrize("num_tokens", [128, 512, 1024, 4096, 8192, 32768], [512, 4096])
+@marker.parametrize("num_experts", [32, 128, 256, 512, 12, 160], [256, 160])
+@marker.parametrize("topk", [1, 2, 4, 8], [2])
+@marker.benchmark("impl", ["jit", "aot", "torch"])
+def benchmark(num_tokens: int, num_experts: int, topk: int, impl: str):
+    if topk > num_experts:
+        marker.skip("topk must be <= num_experts")
+
+    gating_output = create_random(num_tokens, num_experts, dtype=torch.float32)
+    topk_weights = torch.empty((num_tokens, topk), dtype=torch.float32, device="cuda")
+    topk_ids = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda")
+
+    return marker.do_bench(
+        FN_MAP[impl],
+        input_args=(topk_weights, topk_ids, gating_output),
+        # Only the gating logits are read, so they are the only arg worth
+        # rotating to defeat the L2 cache; the two outputs are written every
+        # iteration.
+        graph_clone_args=(2,),
+        # Routing is latency-bound at these sizes, so an achieved-bandwidth
+        # number is not meaningful; report latency only.
+        disable_log_bandwidth=True,
+    )
+
+
+if __name__ == "__main__":
+    benchmark.run()

test/registered/kernels/ops/moe/test_moe_topk_softmax.py

@@ -0,0 +1,200 @@
+"""Correctness tests for the JIT :func:`topk_softmax` MoE router.
+
+The JIT kernel is a host-side port of the AOT ``sgl_kernel.topk_softmax``: the
+device code is unchanged, only the host launcher moves to the tvm-ffi
+``TensorView`` API and the softmax workspace is allocated by the Python wrapper.
+
+We validate it two ways:
+
+* against a definition-based torch reference (documents the math), and
+* against the AOT kernel it replaces, when ``sgl_kernel`` is importable.
+
+Both expert-count regimes are covered: the warp-specialized fast path
+(power-of-two ``num_experts`` <= 512, no scratch) and the two-pass
+softmax + top-k path that everything else falls back to. 512 and 1024 are both
+in the matrix so the boundary between them is pinned on either side.
+
+Against the torch reference, index comparisons are tie-robust: rather than
+requiring identical index tensors, we check that the probability sitting at each
+returned index matches the returned weight, so an arbitrary but valid tie-break
+is accepted. Against the AOT kernel the comparison is exact, because the device
+code is the same on both sides.
+"""
+
+from __future__ import annotations
+
+import sys
+from typing import Optional
+
+import pytest
+import torch
+
+from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.moe.moe_topk_softmax import topk_softmax
+from sglang.test.ci.ci_register import register_cuda_ci
+
+# CI runs the trimmed matrix (17 cases, one dtype), but on a cold runner the
+# single JIT compile is ~29s of the ~31s total -- the case count is nearly free.
+register_cuda_ci(est_time=40, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+DEVICE = "cuda"
+
+try:
+    from sgl_kernel import topk_softmax as aot_topk_softmax
+
+    AOT_AVAILABLE = True
+except Exception:  # pragma: no cover - depends on the installed wheel
+    aot_topk_softmax = None
+    AOT_AVAILABLE = False
+
+DTYPES = get_ci_test_range(
+    full_range=[torch.float32, torch.float16, torch.bfloat16],
+    ci_range=[torch.bfloat16],
+)
+# 8/128/256/512 exercise the warp-specialized power-of-two fast path;
+# 6/160/1024 exercise the workspace (two-pass) path. 512 and 1024 sit either
+# side of the boundary between them, so both must stay in the list.
+NUM_EXPERTS = get_ci_test_range(
+    full_range=[8, 128, 256, 512, 6, 160, 1024],
+    ci_range=[8, 160],
+)
+TOPKS = get_ci_test_range(full_range=[1, 2, 4, 8], ci_range=[2])
+SOFTCAPS = get_ci_test_range(full_range=[0.0, 30.0], ci_range=[0.0])
+
+
+def _reference_probs(
+    gating_output: torch.Tensor,
+    moe_softcapping: float,
+    correction_bias: Optional[torch.Tensor],
+) -> torch.Tensor:
+    """Definition-based reference for the routing probabilities."""
+    logits = gating_output.float()
+    if moe_softcapping:
+        logits = torch.tanh(logits / moe_softcapping) * moe_softcapping
+    if correction_bias is not None:
+        logits = logits + correction_bias.float()
+    return torch.softmax(logits, dim=-1)
+
+
+def _run_jit(gating_output: torch.Tensor, topk: int, renormalize: bool, softcap, bias):
+    num_tokens = gating_output.shape[0]
+    topk_weights = torch.empty(
+        (num_tokens, topk), dtype=torch.float32, device=gating_output.device
+    )
+    topk_ids = torch.empty(
+        (num_tokens, topk), dtype=torch.int32, device=gating_output.device
+    )
+    topk_softmax(topk_weights, topk_ids, gating_output, renormalize, softcap, bias)
+    return topk_weights, topk_ids
+
+
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("num_experts", NUM_EXPERTS)
+@pytest.mark.parametrize("topk", TOPKS)
+@pytest.mark.parametrize("moe_softcapping", SOFTCAPS)
+@pytest.mark.parametrize("use_bias", [False, True])
+@pytest.mark.parametrize("renormalize", [False, True])
+def test_topk_softmax_vs_torch(
+    dtype, num_experts, topk, moe_softcapping, use_bias, renormalize
+):
+    if topk > num_experts:
+        pytest.skip("topk must be <= num_experts")
+
+    num_tokens = 200
+    gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device=DEVICE)
+    correction_bias = (
+        torch.randn(num_experts, dtype=torch.float32, device=DEVICE)
+        if use_bias
+        else None
+    )
+
+    weights, ids = _run_jit(
+        gating_output, topk, renormalize, moe_softcapping, correction_bias
+    )
+
+    probs = _reference_probs(gating_output, moe_softcapping, correction_bias)
+    ref_weights, _ = probs.topk(topk, dim=-1)
+    expected = (
+        ref_weights / ref_weights.sum(-1, keepdim=True) if renormalize else ref_weights
+    )
+
+    tol = 1e-3 if dtype == torch.float32 else 2e-2
+    torch.testing.assert_close(weights, expected, rtol=tol, atol=tol)
+
+    # Tie-robust index check: the probability at each returned index must equal
+    # the returned weight (undoing renormalization first).
+    gathered = torch.gather(probs, 1, ids.long())
+    unnormalized = (
+        weights * ref_weights.sum(-1, keepdim=True) if renormalize else weights
+    )
+    torch.testing.assert_close(gathered, unnormalized, rtol=tol, atol=tol)
+
+    # Indices must be distinct within a row.
+    sorted_ids, _ = ids.sort(dim=-1)
+    assert (sorted_ids[:, 1:] != sorted_ids[:, :-1]).all(), "duplicate expert ids"
+    assert ((ids >= 0) & (ids < num_experts)).all(), "expert id out of range"
+
+
+@pytest.mark.skipif(not AOT_AVAILABLE, reason="sgl_kernel (AOT) is not importable")
+@pytest.mark.parametrize("dtype", DTYPES)
+@pytest.mark.parametrize("num_experts", NUM_EXPERTS)
+@pytest.mark.parametrize("topk", TOPKS)
+@pytest.mark.parametrize("renormalize", [False, True])
+def test_topk_softmax_matches_aot(dtype, num_experts, topk, renormalize):
+    """The JIT port must be numerically identical to the AOT kernel."""
+    if topk > num_experts:
+        pytest.skip("topk must be <= num_experts")
+
+    num_tokens = 512
+    gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device=DEVICE)
+
+    jit_weights, jit_ids = _run_jit(gating_output, topk, renormalize, 0.0, None)
+
+    aot_weights = torch.empty((num_tokens, topk), dtype=torch.float32, device=DEVICE)
+    aot_ids = torch.empty((num_tokens, topk), dtype=torch.int32, device=DEVICE)
+    aot_topk_softmax(aot_weights, aot_ids, gating_output, renormalize, 0.0, None)
+
+    # The device code is unchanged from the AOT kernel and both dispatchers
+    # agree on which path each expert count takes, so this is bit-identical --
+    # on the warp-specialized path and on the two-pass path alike. Keeping the
+    # tolerance at exactly zero is deliberate: it is what caught the dispatcher
+    # falling out of sync with the AOT one at num_experts == 512.
+    assert torch.equal(jit_ids, aot_ids)
+    torch.testing.assert_close(jit_weights, aot_weights, rtol=0, atol=0)
+
+
+@pytest.mark.parametrize("num_experts", [8, 160])
+def test_topk_softmax_single_token(num_experts):
+    gating_output = torch.randn((1, num_experts), dtype=torch.bfloat16, device=DEVICE)
+    weights, ids = _run_jit(gating_output, 2, True, 0.0, None)
+    torch.testing.assert_close(
+        weights.sum(-1), torch.ones(1, device=DEVICE), rtol=1e-2, atol=1e-2
+    )
+    assert ids.shape == (1, 2)
+
+
+@pytest.mark.parametrize("num_experts", [8, 160])
+def test_topk_softmax_full_topk(num_experts):
+    """topk == num_experts: weights must be a permutation of the full softmax."""
+    gating_output = torch.randn((16, num_experts), dtype=torch.float32, device=DEVICE)
+    weights, ids = _run_jit(gating_output, num_experts, False, 0.0, None)
+    probs = _reference_probs(gating_output, 0.0, None)
+    torch.testing.assert_close(
+        weights.sort(dim=-1).values, probs.sort(dim=-1).values, rtol=1e-3, atol=1e-3
+    )
+    assert (
+        ids.sort(dim=-1)
+        .values.eq(torch.arange(num_experts, device=DEVICE, dtype=torch.int32))
+        .all()
+    )
+
+
+def test_topk_softmax_zero_tokens():
+    """An empty batch must be a no-op rather than a launch failure."""
+    gating_output = torch.randn((0, 8), dtype=torch.bfloat16, device=DEVICE)
+    weights, ids = _run_jit(gating_output, 2, False, 0.0, None)
+    assert weights.shape == (0, 2) and ids.shape == (0, 2)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))