chore: import upstream snapshot with attribution

python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh

@@ -0,0 +1,461 @@
+#pragma once
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/utils.cuh>
+
+#include <cuda/ptx>
+#include <tvm/ffi/container/tensor.h>
+
+#include "cute/tensor.hpp"
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+namespace expert_specialization {
+
+using namespace cute;
+
+constexpr uint32_t THREAD_BLOCK_SIZE = 128;
+constexpr uint32_t WARP_SIZE = 32;
+constexpr int BLOCK_M = 128;
+constexpr int BLOCK_K = 128;
+using ThrLayout = Layout<Shape<_16, _8>, Stride<_8, _1>>;
+using ValLayout = Layout<Shape<_1, _16>>;
+using SfR2SThrLayout = Layout<Shape<_16, _4>, Stride<_4, _1>>;
+using SfR2SValLayout = Layout<Shape<_1, _1>>;
+using ScaleFactorTileLayout = Layout<Shape<Shape<_32, _4>, _4>, Stride<Stride<_16, _4>, _1>>;
+
+// Fast reciprocal.
+inline __device__ float reciprocal_approximate_ftz(float a) {
+  float b;
+  asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
+  return b;
+}
+
+// Some code references TRT-LLM:
+// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/quantization.cuh
+template <typename FragmentS, typename FragmentD>
+__inline__ __device__ uint8_t cvt_warp_fp16_to_mxfp8(FragmentS& fragment_s, FragmentD& fragment_d) {
+  using FragmentSLayout = typename FragmentS::layout_type;
+  using FragmentDLayout = typename FragmentD::layout_type;
+  FragmentSLayout fragment_s_layout;
+  FragmentDLayout fragment_d_layout;
+  static_assert(is_static<FragmentSLayout>::value && size(fragment_s_layout) == 16);
+  static_assert(is_static<FragmentDLayout>::value && size(fragment_d_layout) == 16);
+
+  constexpr int eles_per_thr = 16;
+  using ValType = typename FragmentS::element_type;
+  using VecType = std::conditional_t<std::is_same_v<ValType, __nv_bfloat16>, __nv_bfloat162, __half2>;
+  VecType vec[8];
+  // Assign vals
+  vec[0].x = fragment_s(Int<0>{});
+  vec[0].y = fragment_s(Int<1>{});
+  vec[1].x = fragment_s(Int<2>{});
+  vec[1].y = fragment_s(Int<3>{});
+  vec[2].x = fragment_s(Int<4>{});
+  vec[2].y = fragment_s(Int<5>{});
+  vec[3].x = fragment_s(Int<6>{});
+  vec[3].y = fragment_s(Int<7>{});
+  vec[4].x = fragment_s(Int<8>{});
+  vec[4].y = fragment_s(Int<9>{});
+  vec[5].x = fragment_s(Int<10>{});
+  vec[5].y = fragment_s(Int<11>{});
+  vec[6].x = fragment_s(Int<12>{});
+  vec[6].y = fragment_s(Int<13>{});
+  vec[7].x = fragment_s(Int<14>{});
+  vec[7].y = fragment_s(Int<15>{});
+
+  auto local_max = __habs2(vec[0]);
+  for (int i = 1; i < eles_per_thr / 2; i++) {
+    local_max = __hmax2(__habs2(vec[i]), local_max);
+  }
+  local_max = __hmax2(__shfl_xor_sync(uint32_t(-1), local_max, 1), local_max);
+
+  // Get the final absolute maximum values.
+  float block_max(0.0f);
+  if constexpr (std::is_same_v<ValType, __nv_bfloat16>) {
+    block_max = __bfloat162float(__hmax(local_max.x, local_max.y));
+  } else {
+    block_max = __half2float(__hmax(local_max.x, local_max.y));
+  }
+  // Get the SF (max value of the vector / max value of mxfp8).
+  float sf_val = block_max * reciprocal_approximate_ftz(448.0f);
+  // 8 bits representation of the SF.
+  uint8_t fp8_sf_val;
+
+  __nv_fp8_e8m0 tmp_sf_val;
+  tmp_sf_val.__x = __nv_cvt_float_to_e8m0(sf_val, __NV_SATFINITE, cudaRoundPosInf);
+  sf_val = static_cast<float>(tmp_sf_val);
+  fp8_sf_val = tmp_sf_val.__x;
+  // Get the output scale (reciprocal of the SFValue).
+  float output_scale = block_max != 0.f ? reciprocal_approximate_ftz(sf_val) : 0.0f;
+
+  // Convert the input to float.
+  float2 fp2_vals[eles_per_thr / 2];
+
+#pragma unroll
+  for (int i = 0; i < eles_per_thr / 2; i++) {
+    if constexpr (std::is_same_v<ValType, __half>) {
+      fp2_vals[i] = __half22float2(vec[i]);
+    } else {
+      fp2_vals[i] = __bfloat1622float2(vec[i]);
+    }
+    fp2_vals[i].x *= output_scale;
+    fp2_vals[i].y *= output_scale;
+  }
+  union {
+    uint8_t bytes[16];
+    __nv_fp8x2_e4m3 elts[8];
+  } u;
+  u.elts[0] = __nv_fp8x2_e4m3(fp2_vals[0]);
+  u.elts[1] = __nv_fp8x2_e4m3(fp2_vals[1]);
+  u.elts[2] = __nv_fp8x2_e4m3(fp2_vals[2]);
+  u.elts[3] = __nv_fp8x2_e4m3(fp2_vals[3]);
+  u.elts[4] = __nv_fp8x2_e4m3(fp2_vals[4]);
+  u.elts[5] = __nv_fp8x2_e4m3(fp2_vals[5]);
+  u.elts[6] = __nv_fp8x2_e4m3(fp2_vals[6]);
+  u.elts[7] = __nv_fp8x2_e4m3(fp2_vals[7]);
+  fragment_d(Int<0>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[0]);
+  fragment_d(Int<1>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[1]);
+  fragment_d(Int<2>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[2]);
+  fragment_d(Int<3>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[3]);
+  fragment_d(Int<4>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[4]);
+  fragment_d(Int<5>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[5]);
+  fragment_d(Int<6>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[6]);
+  fragment_d(Int<7>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[7]);
+  fragment_d(Int<8>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[8]);
+  fragment_d(Int<9>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[9]);
+  fragment_d(Int<10>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[10]);
+  fragment_d(Int<11>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[11]);
+  fragment_d(Int<12>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[12]);
+  fragment_d(Int<13>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[13]);
+  fragment_d(Int<14>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[14]);
+  fragment_d(Int<15>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[15]);
+  return fp8_sf_val;
+}
+
+template <
+    typename TensorS,
+    typename TensorP,
+    typename TensorD,
+    typename TensorSharedSF,
+    typename TensorSF,
+    typename TiledCopyG2R,
+    typename TiledCopyR2G,
+    typename TiledCopyR2S>
+__inline__ __device__ void mxfp8_group_quant_tile(
+    TensorS& tensor_s,
+    TensorP& tensor_p,
+    TensorD& tensor_d,
+    TensorSharedSF& tensor_shared_sf,
+    TensorSF& tensor_sf,
+    int m,
+    TiledCopyG2R& tiled_copy_g2r,
+    TiledCopyR2G& tiled_copy_r2g,
+    TiledCopyR2S& tiled_copy_r2s) {
+  static_assert(
+      size(get<0>(typename TensorS::layout_type{})) == 128 && size(get<1>(typename TensorS::layout_type{})) == 128 &&
+      stride(get<1>(typename TensorS::layout_type{})) == 1);
+  static_assert(
+      size(get<0>(typename TensorD::layout_type{})) == 128 && size(get<1>(typename TensorD::layout_type{})) == 128 &&
+      stride(get<1>(typename TensorD::layout_type{})) == 1);
+  static_assert(
+      size(get<0>(typename TensorP::layout_type{})) == 128 && size(get<1>(typename TensorP::layout_type{})) == 128);
+  static_assert(
+      size(get<0>(typename TensorSharedSF::layout_type{})) == 128 &&
+      size(get<1>(typename TensorSharedSF::layout_type{})) == 4);
+  static_assert(
+      size(get<0>(typename TensorSF::layout_type{})) == 128 && size(get<1>(typename TensorSF::layout_type{})) == 4);
+
+  using Tiler_MN = typename TiledCopyG2R::Tiler_MN;
+  auto tiler_mn = Tiler_MN{};
+  static_assert(size<0>(tiler_mn) == 16 && size<1>(tiler_mn) == 128);
+
+  auto tiled_tensor_s = tiled_divide(tensor_s, tiler_mn);
+  auto tiled_tensor_p = tiled_divide(tensor_p, tiler_mn);
+  auto tiled_tensor_d = tiled_divide(tensor_d, tiler_mn);
+  static_assert(size<2>(tiled_tensor_s) == 1);
+  static_assert(size<2>(tiled_tensor_p) == 1);
+  static_assert(size<2>(tiled_tensor_d) == 1);
+  auto squeeze_tiled_tensor_s = take<0, 2>(tiled_tensor_s);
+  auto squeeze_tiled_tensor_p = take<0, 2>(tiled_tensor_p);
+  auto squeeze_tiled_tensor_d = take<0, 2>(tiled_tensor_d);
+
+  using SF_Tiler_MN = typename TiledCopyR2S::Tiler_MN;
+  auto sf_tiler_mn = SF_Tiler_MN{};
+  static_assert(size<0>(sf_tiler_mn) == 16 && size<1>(sf_tiler_mn) == 4);
+
+  auto tiled_tensor_sf = tiled_divide(tensor_sf, sf_tiler_mn);
+  auto tiled_tensor_shared_sf = tiled_divide(tensor_shared_sf, sf_tiler_mn);
+  auto squeeze_tiled_tensor_sf = take<0, 2>(tiled_tensor_sf);
+  auto squeeze_tiled_tensor_shared_sf = take<0, 2>(tiled_tensor_shared_sf);
+
+  constexpr int tile_loop_count = size<1>(tiled_tensor_s);
+  constexpr int rows_in_tile = 16;
+  // We don't need to clear shared memory
+  // clear(squeeze_tiled_tensor_shared_sf);
+#pragma unroll 4
+  for (int t = 0; t < tile_loop_count; t++) {
+    if (t * rows_in_tile >= m) {
+      break;
+    }
+    auto current_copy_tile_s = tensor<0>(squeeze_tiled_tensor_s(_, t));
+    auto current_copy_tile_p = tensor<0>(squeeze_tiled_tensor_p(_, t));
+    auto current_copy_tile_d = tensor<0>(squeeze_tiled_tensor_d(_, t));
+    auto current_copy_tile_sf = tensor<0>(squeeze_tiled_tensor_sf(_, t));
+    auto current_copy_tile_shared_sf = tensor<0>(squeeze_tiled_tensor_shared_sf(_, t));
+
+    // Global to Register copy
+    auto thr_copy_g2r = tiled_copy_g2r.get_thread_slice(threadIdx.x);
+    auto thr_tile_g2r_s = thr_copy_g2r.partition_S(current_copy_tile_s);
+    auto thr_tile_g2r_p = thr_copy_g2r.partition_S(current_copy_tile_p);
+    auto input_fragment = make_fragment_like(thr_tile_g2r_s);
+
+    // Register to Global copy
+    auto thr_copy_r2g = tiled_copy_r2g.get_thread_slice(threadIdx.x);
+    auto thr_tile_r2g_d = thr_copy_r2g.partition_D(current_copy_tile_d);
+    auto thr_tile_r2g_p = thr_copy_r2g.partition_D(current_copy_tile_p);
+    auto output_fragment = make_fragment_like(thr_tile_r2g_d);
+
+    // Register to Shared copy
+    auto thr_copy_r2s = tiled_copy_r2s.get_thread_slice(threadIdx.x / 2);
+    auto thr_tile_r2s_shared_sf = thr_copy_r2s.partition_D(current_copy_tile_shared_sf);
+    auto shared_sf_fragment = make_fragment_like(thr_tile_r2s_shared_sf);
+
+    // CopyG2R & convert & CopyR2G
+    copy_if(tiled_copy_g2r, thr_tile_g2r_p, thr_tile_g2r_s, input_fragment);
+    uint8_t fp8_sf_val = cvt_warp_fp16_to_mxfp8(input_fragment, output_fragment);
+    copy_if(tiled_copy_r2g, thr_tile_r2g_p, output_fragment, thr_tile_r2g_d);
+    shared_sf_fragment[0] = fp8_sf_val;
+
+    // Before first copy r2s, clear shared memory and wait previous group
+    if (t == 0 && threadIdx.x == 0) {
+      // Wait for the group to have completed reading from shared memory.
+      cuda::ptx::cp_async_bulk_wait_group_read(cuda::ptx::n32_t<0>());
+    }
+    __syncthreads();
+
+    if (threadIdx.x % 2 == 0) {
+      copy(tiled_copy_r2s, shared_sf_fragment, thr_tile_r2s_shared_sf);
+    }
+    __syncthreads();
+  }
+
+  // Wait for shared memory writes to be visible to TMA engine.
+  cuda::ptx::fence_proxy_async(cuda::ptx::space_shared);  // b)
+  __syncthreads();
+
+  if (threadIdx.x == 0) {
+    cuda::ptx::cp_async_bulk(
+        cuda::ptx::space_global,
+        cuda::ptx::space_shared,
+        squeeze_tiled_tensor_sf.data().get(),
+        squeeze_tiled_tensor_shared_sf.data().get(),
+        512);
+    // Wait for TMA transfer to have finished reading shared memory.
+    // Create a "bulk async-group" out of the previous bulk copy operation.
+    cuda::ptx::cp_async_bulk_commit_group();
+  }
+  __syncthreads();
+}
+
+template <typename T_IN, typename TiledCopyG2R, typename TiledCopyR2G, typename TiledCopyR2S>
+__global__ void mxfp8_group_quant(
+    const T_IN* input,
+    const int* tokens_per_expert,
+    const int* expert_offsets,
+    const int* blockscale_offsets,
+    cutlass::float_e4m3_t* quant_output,
+    uint8_t* scale_factor,
+    int groups,
+    int k,
+    TiledCopyG2R tiled_copy_g2r,
+    TiledCopyR2G tiled_copy_r2g,
+    TiledCopyR2S tiled_copy_r2s) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
+  __shared__ __align__(512) uint8_t shared_memory[512];
+  ScaleFactorTileLayout scale_factor_tile_layout{};
+  auto scale_factor_shared = make_tensor(
+      make_smem_ptr(shared_memory),
+      scale_factor_tile_layout);  // ((_32,_4), _4):((_16,_4), _1)
+  // Transform Groupwise Schedule into Flatten Schedule
+  uint group_total_tiles = 0;
+  uint head_cta_id = 0;
+  for (int g = 0; g < groups; g++) {
+    int m = tokens_per_expert[g];
+    int64_t expert_offset = static_cast<int64_t>(expert_offsets[g]);
+    int64_t blockscale_offset = static_cast<int64_t>(blockscale_offsets[g]);
+
+    auto input_tensor = make_tensor(
+        make_gmem_ptr(input + expert_offset * k),
+        make_layout(make_shape(m, k), LayoutRight{}));  // (M, K):(K, 1) half_t/bfloat16_t
+
+    auto quant_output_tensor = make_tensor(
+        make_gmem_ptr(quant_output + expert_offset * k),
+        make_layout(make_shape(m, k), LayoutRight{}));  // (M, K):(K, 1) cutlass::float_e4m3_t
+
+    auto scale_factor_shape = make_shape(ceil_div(m, 128) * 128, k / 32);
+    auto scale_factor_layout = tile_to_shape(scale_factor_tile_layout, scale_factor_shape, LayoutRight{});
+    // layout<0>(layout<0>(scale_factor_layout))  (_32,_4):(_16,_4) -- static
+    // layout<1>(layout<0>(scale_factor_layout))  M_align_128 / 128 -- dynamic shape dynamic stride
+    // layout<0>(layout<1>(scale_factor_layout))  _4:_1 -- static
+    // layout<1>(layout<1>(scale_factor_layout))  (K / 32) / 4 : _512 -- dynamic shape static stride
+
+    // Reshape to zipped layout for 1D indexing
+    auto zipped_scale_factor_layout = make_layout(
+        make_layout(layout<0>(layout<0>(scale_factor_layout)), layout<0>(layout<1>(scale_factor_layout))),
+        make_layout(
+            layout<1>(layout<0>(scale_factor_layout)),
+            layout<1>(layout<1>(
+                scale_factor_layout))));  // (((_32,_4),_4),(M_align_128 / 128,(K / 32) / 4)):(((_16,_4),_1),(?,_512))
+
+    auto scale_factor_tensor =
+        make_tensor(make_gmem_ptr(scale_factor + blockscale_offset * (k / 32)), zipped_scale_factor_layout);
+
+    // Used for cases where M is not divisible by 128 (most scenarios).
+    auto input_shape = shape(input_tensor);  // (M, K):(K, 1)
+    auto identity_tensor = make_identity_tensor(input_shape);
+    auto predict_tensor = cute::lazy::transform(identity_tensor, [&](auto c) { return elem_less(c, input_shape); });
+
+    // (_128, _128)
+    auto tiler = make_shape(Int<BLOCK_M>{}, Int<BLOCK_K>{});
+
+    auto tiled_input_tensor = zipped_divide(input_tensor, tiler);  // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
+    auto tiled_quant_output_tensor =
+        zipped_divide(quant_output_tensor, tiler);                     // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
+    auto tiled_predict_tensor = zipped_divide(predict_tensor, tiler);  // ((128, 128), (cdiv(M, 128), cdiv(K, 128)))
+
+    auto total_tiles = size<1>(tiled_input_tensor);  // cdiv(M, 128) * cdiv(K, 128)
+    group_total_tiles += total_tiles;
+    auto blk_offset = (blockIdx.x + (gridDim.x - head_cta_id)) % gridDim.x;
+    head_cta_id = group_total_tiles % gridDim.x;
+    while (blk_offset < total_tiles) {
+      auto current_input_tile = tensor<0>(tiled_input_tensor(_, blk_offset));
+      auto current_quant_output_tile = tensor<0>(tiled_quant_output_tensor(_, blk_offset));
+      auto current_predict_tile = tensor<0>(tiled_predict_tensor(_, blk_offset));
+      auto current_scale_factor_tile = tensor<0>(scale_factor_tensor(_, blk_offset));
+
+      mxfp8_group_quant_tile<
+          decltype(current_input_tile),
+          decltype(current_predict_tile),
+          decltype(current_quant_output_tile),
+          decltype(scale_factor_shared),
+          decltype(current_scale_factor_tile),
+          TiledCopyG2R,
+          TiledCopyR2G,
+          TiledCopyR2S>(
+          current_input_tile,
+          current_predict_tile,
+          current_quant_output_tile,
+          scale_factor_shared,
+          current_scale_factor_tile,
+          m,
+          tiled_copy_g2r,
+          tiled_copy_r2g,
+          tiled_copy_r2s);
+      blk_offset += gridDim.x;
+    }
+  }
+#endif
+}
+
+template <typename T_IN>
+void launch_es_sm100_mxfp8_blockscaled_grouped_quant(
+    const T_IN* input,
+    const int* tokens_per_expert,
+    const int* expert_offsets,
+    const int* blockscale_offsets,
+    cutlass::float_e4m3_t* quant_output,
+    uint8_t* scale_factor,
+    int num_experts,
+    int k,
+    int sm_count,
+    cudaStream_t stream) {
+  ThrLayout thr_layout{};
+  ValLayout val_layout{};
+  SfR2SThrLayout r2s_thr_layout{};
+  SfR2SValLayout r2s_val_layout{};
+
+  using CopyOpG2R = UniversalCopy<cutlass::AlignedArray<T_IN, size(val_layout)>>;
+  using CopyAtomG2R = cute::Copy_Atom<CopyOpG2R, T_IN>;
+  auto tiled_copy_g2r = cute::make_tiled_copy(CopyAtomG2R{}, thr_layout, val_layout);  // Tiler_MN: (16, 128)
+
+  using CopyOpR2G = UniversalCopy<cutlass::AlignedArray<cutlass::float_e4m3_t, size(val_layout)>>;
+  using CopyAtomR2G = cute::Copy_Atom<CopyOpR2G, cutlass::float_e4m3_t>;
+  auto tiled_copy_r2g = cute::make_tiled_copy(CopyAtomR2G{}, thr_layout, val_layout);  // Tiler_MN: (16, 128)
+
+  using CopyOpR2S = UniversalCopy<cutlass::AlignedArray<uint8_t, size(r2s_val_layout)>>;
+  using CopyAtomR2S = cute::Copy_Atom<CopyOpR2S, uint8_t>;
+  auto tiled_copy_r2s = cute::make_tiled_copy(CopyAtomR2S{}, r2s_thr_layout, r2s_val_layout);  // Tiler_MN: (16, 4)
+
+  int max_active_blocks_per_sm = -1;
+  auto error_code = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+      &max_active_blocks_per_sm,
+      mxfp8_group_quant<T_IN, decltype(tiled_copy_g2r), decltype(tiled_copy_r2g), decltype(tiled_copy_r2s)>,
+      THREAD_BLOCK_SIZE,
+      0);
+  host::RuntimeCheck(error_code == cudaSuccess, "cudaOccupancyMaxActiveBlocksPerMultiprocessor failed");
+
+  dim3 grid(sm_count * max_active_blocks_per_sm, 1, 1);
+  dim3 block(THREAD_BLOCK_SIZE, 1, 1);
+  mxfp8_group_quant<T_IN, decltype(tiled_copy_g2r), decltype(tiled_copy_r2g), decltype(tiled_copy_r2s)>
+      <<<grid, block, 0, stream>>>(
+          input,
+          tokens_per_expert,
+          expert_offsets,
+          blockscale_offsets,
+          quant_output,
+          scale_factor,
+          num_experts,
+          k,
+          tiled_copy_g2r,
+          tiled_copy_r2g,
+          tiled_copy_r2s);
+}
+
+}  // namespace expert_specialization
+
+template <typename DType>
+struct EsSm100MXFP8BlockscaledGroupQuant {
+  static void
+  run(const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView tokens_per_expert,
+      const tvm::ffi::TensorView expert_offsets,
+      const tvm::ffi::TensorView blockscale_offsets,
+      tvm::ffi::TensorView quant_output,
+      tvm::ffi::TensorView scale_factor) {
+    using namespace host;
+    auto N = SymbolicSize{"num_tokens"};
+    auto D = SymbolicSize{"hidden_size"};
+    auto G = SymbolicSize{"num_experts"};
+    auto N_SF_Alinged = SymbolicSize{"num_tokens_sf_aligned"};
+    auto D_SF = SymbolicSize{"hidden_size_sf"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, D}).with_strides({D, 1}).with_dtype<DType>().with_device(device).verify(input);
+    TensorMatcher({G}).with_dtype<int>().with_device(device).verify(tokens_per_expert);
+    TensorMatcher({G}).with_dtype<int>().with_device(device).verify(expert_offsets);
+    TensorMatcher({G}).with_dtype<int>().with_device(device).verify(blockscale_offsets);
+    RuntimeCheck(D.unwrap() % 128 == 0, "k must align to 128");
+
+    TensorMatcher({N, D}).with_strides({D, 1}).with_dtype<fp8_e4m3_t>().with_device(device).verify(quant_output);
+    TensorMatcher({N_SF_Alinged, D_SF}).with_dtype<uint8_t>().with_device(device).verify(scale_factor);
+    RuntimeCheck(D.unwrap() / 32 == D_SF.unwrap(), "Scale factor K should be hidden_size / 32");
+
+    cudaStream_t stream = LaunchKernel::resolve_device(device.unwrap());
+    expert_specialization::launch_es_sm100_mxfp8_blockscaled_grouped_quant<DType>(
+        reinterpret_cast<const DType*>(input.data_ptr()),
+        reinterpret_cast<const int*>(tokens_per_expert.data_ptr()),
+        reinterpret_cast<const int*>(expert_offsets.data_ptr()),
+        reinterpret_cast<const int*>(blockscale_offsets.data_ptr()),
+        reinterpret_cast<cutlass::float_e4m3_t*>(quant_output.data_ptr()),
+        reinterpret_cast<uint8_t*>(scale_factor.data_ptr()),
+        static_cast<int>(G.unwrap()),
+        static_cast<int>(D.unwrap()),
+        runtime::get_sm_count(device.unwrap().device_id),
+        stream);
+  }
+};

python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm.cuh

@@ -0,0 +1,217 @@
+#pragma once
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/utils.cuh>
+
+#include "cute/tensor.hpp"
+#include "es_sm100_mxfp8_blockscaled_moe_group_gemm_functor.cuh"
+#include "es_sm100_mxfp8_blockscaled_moe_group_gemm_traits.cuh"
+
+namespace expert_specialization {
+
+using namespace host;
+
+template <typename GemmTraits>
+void es_sm100_mxfp8_blockscaled_moe_group_gemm_pre_compute(
+    tvm::ffi::TensorView b,
+    tvm::ffi::TensorView sfb,
+    tvm::ffi::TensorView expert_offsets,
+    tvm::ffi::TensorView blockscale_offsets,
+    tvm::ffi::TensorView b_ptrs,
+    tvm::ffi::TensorView sfb_ptrs,
+    tvm::ffi::TensorView d,
+    tvm::ffi::TensorView d_ptrs,
+    int num_experts,
+    int m,
+    int k,
+    cudaStream_t stream) {
+  using OffsetFunctor = Sm100Mxfp8BlockScaledMoeGroupGemmOffsetFunctor<GemmTraits>;
+  using ElementB = typename OffsetFunctor::ElementB;
+  using ElementSF = typename OffsetFunctor::ElementSF;
+  using ElementD = typename OffsetFunctor::ElementD;
+
+  host::RuntimeCheck(num_experts <= 1024, "num_experts more than 1024");
+  OffsetFunctor offset_functor(
+      reinterpret_cast<int*>(expert_offsets.data_ptr()),
+      reinterpret_cast<int*>(blockscale_offsets.data_ptr()),
+      reinterpret_cast<ElementB*>(b.data_ptr()),
+      reinterpret_cast<ElementSF*>(sfb.data_ptr()),
+      reinterpret_cast<ElementD*>(d.data_ptr()),
+      reinterpret_cast<ElementB**>(b_ptrs.data_ptr()),
+      reinterpret_cast<ElementSF**>(sfb_ptrs.data_ptr()),
+      reinterpret_cast<ElementD**>(d_ptrs.data_ptr()));
+
+  sm100_mxfp8_blockscaled_moe_group_gemm_pre_compute_kernel<<<1, num_experts, 0, stream>>>(offset_functor, m, k);
+}
+
+template <typename GemmTraits>
+void es_sm100_mxfp8_blockscaled_moe_group_gemm(
+    tvm::ffi::TensorView a,
+    tvm::ffi::TensorView sfa,
+    tvm::ffi::TensorView tokens_per_expert,
+    tvm::ffi::TensorView b_ptrs,
+    tvm::ffi::TensorView sfb_ptrs,
+    tvm::ffi::TensorView d_ptrs,
+    tvm::ffi::TensorView workspace,
+    int num_experts,
+    int m,
+    int n,
+    int k,
+    int device_id,
+    int sm_count,
+    cudaStream_t stream) {
+  using Gemm = typename GemmTraits::Gemm;
+  using ElementA = typename Gemm::ElementA;
+  using ElementB = typename Gemm::ElementB;
+  using ElementSF = typename GemmTraits::ElementSF;
+  using ElementD = typename GemmTraits::ElementD;
+
+  cutlass::KernelHardwareInfo hw_info;
+  hw_info.device_id = device_id;
+  hw_info.sm_count = sm_count;
+  hw_info.cluster_shape = GemmTraits::MMAConfig::preferred_cluster;
+  hw_info.cluster_shape_fallback = GemmTraits::MMAConfig::fallback_cluster;
+
+  typename Gemm::Arguments arguments = {
+      cutlass::gemm::GemmUniversalMode::kGrouped,
+      {m, n, k, num_experts, reinterpret_cast<int*>(tokens_per_expert.data_ptr())},
+      {reinterpret_cast<const ElementA*>(a.data_ptr()),
+       reinterpret_cast<const ElementB**>(b_ptrs.data_ptr()),
+       reinterpret_cast<const ElementSF*>(sfa.data_ptr()),
+       reinterpret_cast<const ElementSF**>(sfb_ptrs.data_ptr())},
+      {{}, nullptr, nullptr, reinterpret_cast<ElementD**>(d_ptrs.data_ptr()), nullptr},
+      hw_info,
+      {}  // Scheduler
+  };
+
+  Gemm gemm;
+
+  auto can_implement_status = gemm.can_implement(arguments);
+  host::RuntimeCheck(can_implement_status == cutlass::Status::kSuccess, "Can not implement MoE Group GEMM");
+
+  auto status = gemm.initialize(arguments, reinterpret_cast<uint8_t*>(workspace.data_ptr()), stream);
+  host::RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to initialize MoE Group GEMM");
+
+  status = gemm.run(stream, nullptr);
+  host::RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to run MoE Group GEMM");
+}
+
+template <typename DType>  // CUTLASS dtype
+void es_sm100_mxfp8_blockscaled_moe_group_gemm_dispatch_dtype(
+    tvm::ffi::TensorView a,
+    tvm::ffi::TensorView b,
+    tvm::ffi::TensorView sfa,
+    tvm::ffi::TensorView sfb,
+    tvm::ffi::TensorView expert_offsets,
+    tvm::ffi::TensorView blockscale_offsets,
+    tvm::ffi::TensorView tokens_per_expert,
+    tvm::ffi::TensorView b_ptrs,
+    tvm::ffi::TensorView sfb_ptrs,
+    tvm::ffi::TensorView d,
+    tvm::ffi::TensorView d_ptrs,
+    tvm::ffi::TensorView workspace,
+    int num_experts,
+    int m,
+    int n,
+    int k,
+    int device_id,
+    int sm_count,
+    cudaStream_t stream) {
+  using GemmTraits = ExpertSpecializationSm100MXFP8BlockscaledMoeGroupGemmTraits<MMA2SMConfig, DType>;
+
+  es_sm100_mxfp8_blockscaled_moe_group_gemm_pre_compute<GemmTraits>(
+      b, sfb, expert_offsets, blockscale_offsets, b_ptrs, sfb_ptrs, d, d_ptrs, num_experts, m, k, stream);
+  es_sm100_mxfp8_blockscaled_moe_group_gemm<GemmTraits>(
+      a,
+      sfa,
+      tokens_per_expert,
+      b_ptrs,
+      sfb_ptrs,
+      d_ptrs,
+      workspace,
+      num_experts,
+      m,
+      n,
+      k,
+      device_id,
+      sm_count,
+      stream);
+}
+
+}  // namespace expert_specialization
+
+template <typename DType>
+struct EsSm100MXFP8BlockscaledMoeGroupGemm {
+  static void
+  run(tvm::ffi::TensorView a,
+      tvm::ffi::TensorView b,
+      tvm::ffi::TensorView sfa,
+      tvm::ffi::TensorView sfb,
+      tvm::ffi::TensorView expert_offsets,
+      tvm::ffi::TensorView blockscale_offsets,
+      tvm::ffi::TensorView tokens_per_expert,
+      tvm::ffi::TensorView b_ptrs,
+      tvm::ffi::TensorView sfb_ptrs,
+      tvm::ffi::TensorView d,
+      tvm::ffi::TensorView d_ptrs,
+      tvm::ffi::TensorView workspace) {
+    using namespace host;
+    auto num_tokens = SymbolicSize{"num_tokens"};
+    auto num_sf_tokens = SymbolicSize{"num_sf_tokens"};
+    auto hidden_size = SymbolicSize{"hidden_size"};
+    auto num_experts = SymbolicSize{"num_experts"};
+    auto M = SymbolicSize{"M"};
+    auto K = SymbolicSize{"K"};
+    auto M_SF = SymbolicSize{"M_SF"};
+    auto K_SF = SymbolicSize{"K_SF"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({num_experts, M, K}).with_dtype<fp8_e4m3_t>().with_device(device).verify(a);
+    TensorMatcher({num_tokens, K}).with_dtype<fp8_e4m3_t>().with_device(device).verify(b);
+    TensorMatcher({num_experts, M_SF, K_SF}).with_dtype<uint8_t>().with_device(device).verify(sfa);
+    TensorMatcher({num_sf_tokens, K_SF}).with_dtype<uint8_t>().with_device(device).verify(sfb);
+    RuntimeCheck(K.unwrap() % 128 == 0, "K should align 128");
+    RuntimeCheck(K.unwrap() / 32 == K_SF.unwrap(), "K dimension mismatch");
+
+    TensorMatcher({num_experts}).with_dtype<int>().with_device(device).verify(expert_offsets);
+    TensorMatcher({num_experts}).with_dtype<int>().with_device(device).verify(blockscale_offsets);
+    TensorMatcher({num_experts}).with_dtype<int>().with_device(device).verify(tokens_per_expert);
+    TensorMatcher({num_experts}).with_dtype<int64_t>().with_device(device).verify(b_ptrs);
+    TensorMatcher({num_experts}).with_dtype<int64_t>().with_device(device).verify(sfb_ptrs);
+    TensorMatcher({num_experts}).with_dtype<int64_t>().with_device(device).verify(d_ptrs);
+    // Check output
+    TensorMatcher({num_tokens, M}).with_strides({M, 1}).with_dtype<DType>().with_device(device).verify(d);
+
+    cudaStream_t stream = LaunchKernel::resolve_device(device.unwrap());
+    int device_id = device.unwrap().device_id;
+
+    if constexpr (std::is_same_v<DType, bf16_t> || std::is_same_v<DType, fp16_t>) {
+      using CUTLASS_DTYPE = std::conditional_t<std::is_same_v<DType, bf16_t>, cutlass::bfloat16_t, cutlass::half_t>;
+      expert_specialization::es_sm100_mxfp8_blockscaled_moe_group_gemm_dispatch_dtype<CUTLASS_DTYPE>(
+          a,
+          b,
+          sfa,
+          sfb,
+          expert_offsets,
+          blockscale_offsets,
+          tokens_per_expert,
+          b_ptrs,
+          sfb_ptrs,
+          d,
+          d_ptrs,
+          workspace,
+          static_cast<int>(num_experts.unwrap()),
+          static_cast<int>(M.unwrap()),
+          static_cast<int>(num_tokens.unwrap()),
+          static_cast<int>(K.unwrap()),
+          device_id,
+          static_cast<int>(runtime::get_sm_count(device_id)),
+          stream);
+    } else {
+      Panic("Unsupported dtype");
+    }
+  }
+};

python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_functor.cuh

@@ -0,0 +1,64 @@
+#pragma once
+#include "cute/tensor.hpp"
+#include "es_sm100_mxfp8_blockscaled_moe_group_gemm_traits.cuh"
+#include <cuda.h>
+
+namespace expert_specialization {
+
+using namespace cute;
+
+template <typename GemmTraits>
+struct Sm100Mxfp8BlockScaledMoeGroupGemmOffsetFunctor {
+  using ElementB = typename GemmTraits::Gemm::ElementB;
+  using ElementSF = typename GemmTraits::ElementSF;
+  using ElementD = typename GemmTraits::ElementD;
+  // Input
+  int* expert_offsets{nullptr};
+  int* blockscale_offsets{nullptr};
+  // Output
+  ElementB* b_base{nullptr};
+  ElementSF* sfb_base{nullptr};
+  ElementD* d_base{nullptr};
+  ElementB** b_offsets{nullptr};
+  ElementSF** sfb_offsets{nullptr};
+  ElementD** d_offsets{nullptr};
+
+  Sm100Mxfp8BlockScaledMoeGroupGemmOffsetFunctor() = default;
+  Sm100Mxfp8BlockScaledMoeGroupGemmOffsetFunctor(
+      int* _expert_offsets,
+      int* _blockscale_offsets,
+      ElementB* _b_base,
+      ElementSF* _sfb_base,
+      ElementD* _d_base,
+      ElementB** _b_offsets,
+      ElementSF** _sfb_offsets,
+      ElementD** _d_offsets)
+      : expert_offsets{_expert_offsets},
+        blockscale_offsets{_blockscale_offsets},
+        b_base(_b_base),
+        sfb_base(_sfb_base),
+        d_base(_d_base),
+        b_offsets(_b_offsets),
+        sfb_offsets(_sfb_offsets),
+        d_offsets(_d_offsets) {}
+
+  void CUTE_DEVICE operator()(int expert_id, int m, int k) {
+    int64_t expert_offset = static_cast<int64_t>(expert_offsets[expert_id]);
+    int64_t blockscale_offset = static_cast<int64_t>(blockscale_offsets[expert_id]);
+    int64_t b_stride = expert_offset * k;
+    int64_t sfb_stride = blockscale_offset * (k / 32);
+    int64_t d_stride = expert_offset * m;
+
+    b_offsets[expert_id] = b_base + b_stride;
+    sfb_offsets[expert_id] = sfb_base + sfb_stride;
+    d_offsets[expert_id] = d_base + d_stride;
+  }
+};
+
+template <typename OffsetFunctor>
+__global__ void sm100_mxfp8_blockscaled_moe_group_gemm_pre_compute_kernel(OffsetFunctor offset_functor, int m, int k) {
+  int expert_id = static_cast<int>(threadIdx.x);
+  offset_functor(expert_id, m, k);
+}
+
+}  // namespace expert_specialization

python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_traits.cuh

@@ -0,0 +1,123 @@
+#pragma once
+
+// Misc
+#include "cute/tensor.hpp"
+#include "cutlass/arch/arch.h"
+#include "cutlass/arch/mma.h"
+#include "cutlass/cutlass.h"
+#include "cutlass/detail/sm100_blockscaled_layout.hpp"
+#include "cutlass/epilogue/dispatch_policy.hpp"
+#include "cutlass/gemm/dispatch_policy.hpp"
+#include "cutlass/gemm/group_array_problem_shape.hpp"
+#include "cutlass/layout/layout.h"
+#include "cutlass/numeric_conversion.h"
+#include "cutlass/numeric_size.h"
+
+// Collective Builder
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp"
+#include "cutlass/epilogue/thread/activation.h"
+#include "cutlass/gemm/collective/collective_builder.hpp"
+
+// Integration
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
+
+namespace expert_specialization {
+
+using namespace cute;
+
+// Different configs for 1SM and 2SM MMA kernel
+struct MMA2SMConfig {
+  using MmaTileShape = Shape<_256, _128, _128>;
+  using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmMxf8f6f4Sm100;
+  using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized2Sm;
+  const static dim3 preferred_cluster;
+  const static dim3 fallback_cluster;
+};
+const dim3 MMA2SMConfig::preferred_cluster(4, 1, 1);
+const dim3 MMA2SMConfig::fallback_cluster(2, 1, 1);
+
+template <typename _MMAConfig, typename OutputDtype>
+struct ExpertSpecializationSm100MXFP8BlockscaledMoeGroupGemmTraits {
+  using MMAConfig = _MMAConfig;
+  using ElementInput = cutlass::float_e4m3_t;
+  using ElementOutput = OutputDtype;
+  using ProblemShape = cutlass::gemm::MoEProblemShape<Shape<int, int, int>>;  // <M,N,K> per group
+
+  // A matrix configuration
+  using ElementA = cutlass::mx_float8_t<ElementInput>;
+  using LayoutA = cutlass::layout::RowMajor;
+  constexpr static int AlignmentA = 16;
+
+  // B matrix configuration
+  using ElementB = cutlass::mx_float8_t<ElementInput>;
+  using LayoutB = cutlass::layout::ColumnMajor;
+  constexpr static int AlignmentB = 16;
+
+  // C/D matrix configuration
+  using ElementC = void;
+  using ElementD = ElementOutput;
+  using LayoutC = cutlass::layout::ColumnMajor;
+  using LayoutD = cutlass::layout::ColumnMajor;
+  constexpr static int AlignmentC = 128 / cutlass::sizeof_bits<ElementD>::value;
+  constexpr static int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
+  using ElementAccumulator = float;
+
+  static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest;
+  using CustomEVTIdentity =  // acc
+      cutlass::epilogue::fusion::Sm90EVT<
+          cutlass::epilogue::fusion::
+              Sm90Compute<cutlass::epilogue::thread::Identity, ElementD, ElementAccumulator, RoundStyle>,
+          cutlass::epilogue::fusion::Sm90AccFetch>;
+
+  // Core kernel configurations
+  using ArchTag = cutlass::arch::Sm100;
+  using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
+  using StageCountType = cutlass::gemm::collective::StageCountAuto;
+
+  // Runtime Cluster Shape
+  using ClusterShape = Shape<int32_t, int32_t, _1>;
+
+  // Define Epilogue
+  using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
+      ArchTag,
+      OperatorClass,
+      typename MMAConfig::MmaTileShape,
+      ClusterShape,
+      cutlass::epilogue::collective::EpilogueTileAuto,
+      ElementAccumulator,
+      ElementAccumulator,
+      ElementC,
+      LayoutC*,
+      AlignmentC,
+      ElementD,
+      LayoutD*,
+      AlignmentD,
+      typename MMAConfig::EpilogueSchedule,
+      CustomEVTIdentity>::CollectiveOp;
+
+  // Define Mainloop
+  using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
+      ArchTag,
+      OperatorClass,
+      ElementA,
+      LayoutA,
+      AlignmentA,
+      ElementB,
+      LayoutB*,
+      AlignmentB,
+      ElementAccumulator,
+      typename MMAConfig::MmaTileShape,
+      ClusterShape,
+      cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
+          sizeof(typename CollectiveEpilogue::SharedStorage))>,
+      typename MMAConfig::KernelSchedule>::CollectiveOp;
+
+  // Define GemmKernel
+  using GemmKernel = cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop, CollectiveEpilogue>;
+  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+  using ElementSF = typename GemmKernel::ElementSF;
+};
+
+}  // namespace expert_specialization

python/sglang/jit_kernel/csrc/moe/moe_align_kernel.cu

@@ -0,0 +1,582 @@
+/* Copyright 2025 SGLang Team. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <algorithm>
+
+#ifndef WARP_SIZE
+#define WARP_SIZE 32
+#endif
+
+#define CEILDIV(x, y) (((x) + (y) - 1) / (y))
+
+#define VEC_SIZE 4
+using Vec = int4;
+
+inline uint32_t next_pow2(uint32_t x) noexcept {
+  --x;
+  x |= x >> 1;
+  x |= x >> 2;
+  x |= x >> 4;
+  x |= x >> 8;
+  x |= x >> 16;
+  return x + 1;
+}
+
+namespace moe {
+
+__device__ __forceinline__ int warp_exclusive_scan(int v, unsigned mask = 0xffffffffu) {
+  int original = v;
+#pragma unroll
+  for (int offset = 1; offset < WARP_SIZE; offset <<= 1) {
+    int n = __shfl_up_sync(mask, v, offset);
+    if ((threadIdx.x & (WARP_SIZE - 1)) >= offset) v += n;
+  }
+  return v - original;
+}
+
+template <typename scalar_t>
+__global__ void count_and_sort_expert_tokens_kernel(
+    const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ sorted_token_ids,
+    int32_t* __restrict__ cumsum_buffer,
+    size_t numel) {
+  const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const size_t stride = blockDim.x * gridDim.x;
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int32_t expert_id = topk_ids[i] + 1;
+    int32_t rank_post_pad = atomicAdd(&cumsum_buffer[expert_id], 1);
+    sorted_token_ids[rank_post_pad] = i;
+  }
+}
+
+template <typename scalar_t>
+__global__ void moe_align_block_size_kernel(
+    const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ sorted_token_ids,
+    int32_t* __restrict__ expert_ids,
+    int32_t* __restrict__ total_tokens_post_pad,
+    int32_t num_experts,
+    int32_t block_size,
+    size_t numel,
+    int32_t* __restrict__ cumsum,
+    bool pad_sorted_token_ids,
+    const int32_t scan_size,
+    int32_t max_num_tokens_padded) {
+  // Use a separate thread block to populate sorted_token_ids
+  if (blockIdx.x == 1) {
+    if (pad_sorted_token_ids) {
+      Vec fill_vec;
+      fill_vec.x = fill_vec.y = fill_vec.z = fill_vec.w = numel;
+      int32_t total_vecs = (max_num_tokens_padded + VEC_SIZE - 1) / VEC_SIZE;
+      Vec* out_ptr = reinterpret_cast<Vec*>(sorted_token_ids);
+      for (int32_t i = threadIdx.x; i < total_vecs; i += blockDim.x) {
+        out_ptr[i] = fill_vec;
+      }
+    }
+    return;
+  }
+
+  extern __shared__ int32_t smem[];
+  int32_t* shared_counts = smem;                  // [num_experts]
+  int32_t* prefix = shared_counts + num_experts;  // [num_experts + 1]
+  int32_t* scan_buf = prefix + num_experts + 1;   // [scan_size]
+  __shared__ int32_t s_total_tokens_post_pad;
+
+  const size_t tid = threadIdx.x;
+  const size_t stride = blockDim.x;
+
+  if (tid < num_experts) {
+    shared_counts[tid] = 0;
+  }
+
+  __syncthreads();
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int expert_id = topk_ids[i] + 1;
+    atomicAdd(&shared_counts[expert_id], 1);
+  }
+
+  __syncthreads();
+
+  int32_t padded_count = 0;
+  if (tid < num_experts) {
+    int32_t count = shared_counts[tid];
+    padded_count = (count + block_size - 1) / block_size * block_size;
+    scan_buf[tid] = padded_count;
+  }
+
+#ifndef __CUDA_ARCH__  // HIP
+
+  if (tid >= num_experts && tid < scan_size) {
+    scan_buf[tid] = 0;
+  }
+
+  __syncthreads();
+
+  // Blelloch scan
+  int offset = 1;
+#pragma unroll
+  for (int d = scan_size >> 1; d > 0; d >>= 1) {
+    if (tid < d) {
+      int ai = offset * (2 * tid + 1) - 1;
+      int bi = offset * (2 * tid + 2) - 1;
+      scan_buf[bi] += scan_buf[ai];
+    }
+    offset <<= 1;
+    __syncthreads();
+  }
+
+  // down-sweep
+  if (tid == 0) {
+    prefix[num_experts] = scan_buf[scan_size - 1];
+    scan_buf[scan_size - 1] = 0;
+  }
+  __syncthreads();
+
+#pragma unroll
+  for (int d = 1; d < scan_size; d <<= 1) {
+    offset >>= 1;
+    if (tid < d) {
+      int ai = offset * (2 * tid + 1) - 1;
+      int bi = offset * (2 * tid + 2) - 1;
+      if (bi < scan_size) {
+        int temp = scan_buf[ai];
+        scan_buf[ai] = scan_buf[bi];
+        scan_buf[bi] += temp;
+      }
+    }
+    __syncthreads();
+  }
+
+  if (tid < num_experts) {
+    prefix[tid] = scan_buf[tid];
+  }
+
+  if (tid == 0) {
+    s_total_tokens_post_pad = prefix[num_experts];
+    *total_tokens_post_pad = s_total_tokens_post_pad;
+  }
+  __syncthreads();
+
+#else  // CUDA
+
+  // Intra warp prefix sum
+  int32_t* warp_sums = scan_buf + scan_size;  // [<= 32]
+  const int warp_id = tid / WARP_SIZE;
+  const int lane_id = tid & (WARP_SIZE - 1);
+  const int num_warps_for_scan = (scan_size + WARP_SIZE - 1) / WARP_SIZE;
+  const int warp_sum = warp_exclusive_scan(padded_count) + padded_count;
+  if (lane_id == WARP_SIZE - 1) warp_sums[warp_id] = warp_sum;
+  __syncthreads();
+
+  // warp0 accumulate all the block's prefix sum
+  if (tid < WARP_SIZE) {
+    int val = (tid < num_warps_for_scan) ? warp_sums[tid] : 0;
+    int incl = warp_exclusive_scan(val) + val;
+    warp_sums[tid] = incl;
+  }
+  __syncthreads();
+
+  // Every thread obtains the whole block's sum
+  if (tid == 0) {
+    prefix[num_experts] = warp_sums[num_warps_for_scan - 1];
+    s_total_tokens_post_pad = prefix[num_experts];
+    *total_tokens_post_pad = s_total_tokens_post_pad;
+  }
+  __syncthreads();
+
+  // Fill 0 to scan_buf extended area (tid >= num_expert)
+  if (tid >= num_experts && tid < scan_size) scan_buf[tid] = 0;
+  __syncthreads();
+
+  // Perform 2 level exclusive-prefix-sum to scan_buf
+  int v = (tid < scan_size) ? scan_buf[tid] : 0;
+  int pre = warp_exclusive_scan(v);
+  if (lane_id == WARP_SIZE - 1) warp_sums[warp_id] = pre + v;
+  __syncthreads();
+
+  if (warp_id == 0) {
+    int val = (lane_id < num_warps_for_scan) ? warp_sums[lane_id] : 0;
+    warp_sums[lane_id] = warp_exclusive_scan(val);
+  }
+  __syncthreads();
+
+  int offset = warp_sums[warp_id];
+  if (tid < scan_size) scan_buf[tid] = pre + offset;
+  __syncthreads();
+
+  // Write prefix[0..num_experts - 1] and cumsum
+  if (tid < num_experts) prefix[tid] = scan_buf[tid];
+#endif
+
+  if (tid <= num_experts) {
+    cumsum[tid] = prefix[tid];
+  }
+  // fill expert_ids
+  const int32_t num_blocks = s_total_tokens_post_pad / block_size;
+  for (int32_t i = tid; i < num_blocks; i += stride) {
+    int32_t block_start = i * block_size;
+    int left = 0, right = num_experts;
+    while (left < right) {
+      int mid = (left + right) >> 1;
+      if (prefix[mid] <= block_start) {
+        left = mid + 1;
+      } else {
+        right = mid;
+      }
+    }
+    expert_ids[i] = left - 2;
+  }
+}
+
+template <typename scalar_t, int32_t fill_threads>
+__global__ void moe_align_block_size_small_batch_expert_kernel(
+    const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ sorted_token_ids,
+    int32_t* __restrict__ expert_ids,
+    int32_t* __restrict__ total_tokens_post_pad,
+    int32_t num_experts,
+    int32_t block_size,
+    size_t numel,
+    bool pad_sorted_token_ids,
+    int32_t max_num_tokens_padded) {
+  // Adapted from
+  // https://github.com/vllm-project/vllm/pull/29642/files#diff-5647b1413f4ae9aacba904eca8f8a8aee9079321eadff4c10101a2c6962dcc53R226
+  // Use an additional group of threads to fill sorted_token_ids.
+  // Since the kernel will use sorted_token_ids afterward,
+  // we fill sorted_token_ids within the same threadblock to make
+  // synchronization easier.
+  if (threadIdx.x < fill_threads) {
+    // Initialize sorted_token_ids with numel
+    if (pad_sorted_token_ids) {
+      for (int32_t it = threadIdx.x; it < max_num_tokens_padded; it += fill_threads) {
+        sorted_token_ids[it] = numel;
+      }
+    }
+    // Three __syncthreads() corresponding to the other threads
+    __syncthreads();
+    __syncthreads();
+    __syncthreads();
+    return;
+  }
+
+  const size_t tid = threadIdx.x - fill_threads;
+  const size_t stride = blockDim.x - fill_threads;
+
+  extern __shared__ int32_t shared_mem[];
+  int32_t* cumsum = shared_mem;
+  int32_t* tokens_cnts = (int32_t*)(shared_mem + num_experts + 1);
+
+  for (int i = 0; i < num_experts; ++i) {
+    tokens_cnts[(tid + 1) * num_experts + i] = 0;
+  }
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int32_t expert_id = topk_ids[i] + 1;
+    ++tokens_cnts[(tid + 1) * num_experts + expert_id];
+  }
+
+  __syncthreads();
+
+  if (tid < num_experts) {
+    tokens_cnts[tid] = 0;
+    for (int i = 1; i <= stride; ++i) {
+      tokens_cnts[i * num_experts + tid] += tokens_cnts[(i - 1) * num_experts + tid];
+    }
+  }
+
+  __syncthreads();
+
+  if (tid == 0) {
+    cumsum[0] = 0;
+    for (int i = 1; i <= num_experts; ++i) {
+      cumsum[i] = cumsum[i - 1] + CEILDIV(tokens_cnts[stride * num_experts + i - 1], block_size) * block_size;
+    }
+    *total_tokens_post_pad = static_cast<int32_t>(cumsum[num_experts]);
+  }
+
+  __syncthreads();
+
+  if (tid < num_experts) {
+    for (int i = cumsum[tid]; i < cumsum[tid + 1]; i += block_size) {
+      expert_ids[i / block_size] = tid - 1;
+    }
+  }
+
+  for (size_t i = tid; i < numel; i += stride) {
+    int32_t expert_id = topk_ids[i] + 1;
+    int32_t rank_post_pad = tokens_cnts[tid * num_experts + expert_id] + cumsum[expert_id];
+    sorted_token_ids[rank_post_pad] = i;
+    ++tokens_cnts[tid * num_experts + expert_id];
+  }
+}
+
+// v2 kernel: supports >1024 experts via EXPERTS_PER_THREAD templating
+// and a two-level warp scan (no cub dependency). Uses the same +1 offset
+// convention as the original kernel (topk_ids shifted by +1 so -1 maps to 0).
+// Launched with <<<2, 1024>>>: block 1 fills sorted_token_ids in parallel
+// with block 0 doing the alignment compute.
+//
+// With 1024 threads and EXPERTS_PER_THREAD=4, covers at most 4096 expert
+// indices. Since num_experts includes the +1 offset bucket, this supports
+// up to 4095 real experts.
+template <typename scalar_t, int EXPERTS_PER_THREAD>
+__global__ void moe_align_block_size_kernel_v2(
+    const scalar_t* __restrict__ topk_ids,
+    int32_t* __restrict__ sorted_token_ids,
+    int32_t* __restrict__ expert_ids,
+    int32_t* __restrict__ total_tokens_post_pad,
+    int32_t num_experts,
+    int32_t padded_num_experts,
+    int32_t block_size,
+    size_t numel,
+    int32_t* __restrict__ cumsum,
+    bool pad_sorted_token_ids,
+    int32_t max_num_tokens_padded) {
+  // Use a separate thread block to populate sorted_token_ids
+  if (blockIdx.x == 1) {
+    if (pad_sorted_token_ids) {
+      Vec fill_vec;
+      fill_vec.x = fill_vec.y = fill_vec.z = fill_vec.w = numel;
+      int32_t total_vecs = (max_num_tokens_padded + VEC_SIZE - 1) / VEC_SIZE;
+      Vec* out_ptr = reinterpret_cast<Vec*>(sorted_token_ids);
+      for (int32_t i = threadIdx.x; i < total_vecs; i += blockDim.x) {
+        out_ptr[i] = fill_vec;
+      }
+    }
+    return;
+  }
+
+  extern __shared__ int32_t smem[];
+  // Layout: shared_counts[padded_num_experts] | warp_sums[WARP_SIZE]
+  int32_t* shared_counts = smem;
+  int32_t* warp_sums = smem + padded_num_experts;
+
+  const size_t tid = threadIdx.x;
+  const int warp_id = tid / WARP_SIZE;
+  const int lane_id = tid & (WARP_SIZE - 1);
+
+  // Phase 1: Zero shared counts and count tokens per expert
+  const int my_start = tid * EXPERTS_PER_THREAD;
+  for (size_t i = tid; i < padded_num_experts; i += blockDim.x) {
+    shared_counts[i] = 0;
+  }
+
+  __syncthreads();
+
+  for (size_t i = tid; i < numel; i += blockDim.x) {
+    int expert_id = topk_ids[i] + 1;  // +1 offset convention
+    if (expert_id < num_experts) {
+      atomicAdd(&shared_counts[expert_id], 1);
+    }
+  }
+
+  __syncthreads();
+
+  // Phase 2: Compute padded counts and two-level warp exclusive prefix sum
+  int32_t local_padded[EXPERTS_PER_THREAD];
+  int32_t thread_sum = 0;
+  for (int i = 0; i < EXPERTS_PER_THREAD; ++i) {
+    int eid = my_start + i;
+    if (eid < num_experts) {
+      local_padded[i] = CEILDIV(shared_counts[eid], block_size) * block_size;
+    } else {
+      local_padded[i] = 0;
+    }
+    thread_sum += local_padded[i];
+  }
+
+  // Level 1: intra-warp exclusive scan on thread_sum
+  int32_t warp_prefix = warp_exclusive_scan(thread_sum);
+  int32_t warp_total = warp_prefix + thread_sum;
+  if (lane_id == WARP_SIZE - 1) warp_sums[warp_id] = warp_total;
+  __syncthreads();
+
+  // Level 2: warp 0 scans the per-warp totals
+  const int num_warps = (blockDim.x + WARP_SIZE - 1) / WARP_SIZE;
+  if (tid < WARP_SIZE) {
+    int val = (tid < num_warps) ? warp_sums[tid] : 0;
+    warp_sums[tid] = warp_exclusive_scan(val);
+  }
+  __syncthreads();
+
+  // Combine: thread_prefix = warp_sums[warp_id] + warp_prefix
+  int32_t thread_prefix = warp_sums[warp_id] + warp_prefix;
+
+  // Local sequential prefix sum within each thread's expert group
+  int32_t running = 0;
+  for (int i = 0; i < EXPERTS_PER_THREAD; ++i) {
+    int eid = my_start + i;
+    if (eid <= num_experts) {
+      cumsum[eid] = thread_prefix + running;
+    }
+    running += local_padded[i];
+  }
+
+  // Last thread writes total
+  if (tid == blockDim.x - 1) {
+    cumsum[num_experts] = thread_prefix + thread_sum;
+    *total_tokens_post_pad = thread_prefix + thread_sum;
+  }
+
+  __syncthreads();
+
+  // Phase 3: Fill expert_ids (eid - 1 to match sgl-kernel convention)
+  for (int i = 0; i < EXPERTS_PER_THREAD; ++i) {
+    int eid = my_start + i;
+    if (eid < num_experts) {
+      for (int j = cumsum[eid]; j < cumsum[eid + 1]; j += block_size) {
+        expert_ids[j / block_size] = eid - 1;
+      }
+    }
+  }
+}
+
+}  // namespace moe
+
+namespace {
+
+template <typename scalar_t>
+struct MoeAlignBlockSizeKernel {
+  static void
+  run(tvm::ffi::TensorView topk_ids,
+      int64_t num_experts,
+      int64_t block_size,
+      tvm::ffi::TensorView sorted_token_ids,
+      tvm::ffi::TensorView expert_ids,
+      tvm::ffi::TensorView num_tokens_post_pad,
+      tvm::ffi::TensorView cumsum_buffer,
+      bool pad_sorted_token_ids) {
+    using namespace host;
+
+    auto device = topk_ids.device();
+    const cudaStream_t stream = LaunchKernel::resolve_device(device);
+
+    int threads = 1024;
+    threads = ((threads + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
+
+    int64_t max_num_tokens_padded = sorted_token_ids.size(0);
+
+    // num_experts from Python is actual_num_experts + 1 (for EP offset convention).
+    // The v2 kernel (>1024 experts) uses 1024 threads with EXPERTS_PER_THREAD up
+    // to 8, covering at most 8192 expert indices. This supports up to 8191 real
+    // experts, sufficient for LoRA virtual experts (num_moe_experts * max_loras).
+    RuntimeCheck(num_experts <= 8192, "moe_align_block_size: num_experts must be <= 8192, got ", num_experts);
+
+    const scalar_t* topk_ids_ptr = static_cast<const scalar_t*>(topk_ids.data_ptr());
+    int32_t* sorted_token_ids_ptr = static_cast<int32_t*>(sorted_token_ids.data_ptr());
+    int32_t* expert_ids_ptr = static_cast<int32_t*>(expert_ids.data_ptr());
+    int32_t* num_tokens_post_pad_ptr = static_cast<int32_t*>(num_tokens_post_pad.data_ptr());
+    int32_t* cumsum_buffer_ptr = static_cast<int32_t*>(cumsum_buffer.data_ptr());
+    size_t numel = topk_ids.numel();
+
+    bool small_batch_expert_mode = (numel < 1024) && (num_experts <= 64);
+
+    if (small_batch_expert_mode) {
+      const int32_t num_thread = std::max((int32_t)num_experts, (int32_t)WARP_SIZE);
+      constexpr int32_t fill_threads = 256;
+      const int32_t shared_mem_size = ((num_thread + 1) * num_experts + (num_experts + 1)) * sizeof(int32_t);
+
+      auto kernel = moe::moe_align_block_size_small_batch_expert_kernel<scalar_t, fill_threads>;
+      LaunchKernel(dim3(1), dim3(fill_threads + num_thread), stream, shared_mem_size)(
+          kernel,
+          topk_ids_ptr,
+          sorted_token_ids_ptr,
+          expert_ids_ptr,
+          num_tokens_post_pad_ptr,
+          (int32_t)num_experts,
+          (int32_t)block_size,
+          numel,
+          pad_sorted_token_ids,
+          (int32_t)max_num_tokens_padded);
+    } else if (num_experts <= 1024) {
+      const size_t scan_size = next_pow2(num_experts);
+      const size_t shared_mem_size = (num_experts + (num_experts + 1) + scan_size + WARP_SIZE) * sizeof(int32_t);
+
+      auto align_kernel = moe::moe_align_block_size_kernel<scalar_t>;
+      LaunchKernel(dim3(2), dim3(threads), stream, shared_mem_size)(
+          align_kernel,
+          topk_ids_ptr,
+          sorted_token_ids_ptr,
+          expert_ids_ptr,
+          num_tokens_post_pad_ptr,
+          (int32_t)num_experts,
+          (int32_t)block_size,
+          numel,
+          cumsum_buffer_ptr,
+          pad_sorted_token_ids,
+          (int32_t)scan_size,
+          (int32_t)max_num_tokens_padded);
+
+      const int block_threads = std::min(256, threads);
+      const int num_blocks = (numel + block_threads - 1) / block_threads;
+      const int max_blocks = 65535;
+      const int actual_blocks = std::min(num_blocks, max_blocks);
+
+      auto sort_kernel = moe::count_and_sort_expert_tokens_kernel<scalar_t>;
+      LaunchKernel(dim3(actual_blocks), dim3(block_threads), stream)(
+          sort_kernel, topk_ids_ptr, sorted_token_ids_ptr, cumsum_buffer_ptr, numel);
+    } else {
+      // v2 path for >1024 experts: two-level warp scan with EXPERTS_PER_THREAD
+      int64_t padded_num_experts = ((num_experts + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
+      size_t shared_mem_size = (padded_num_experts + WARP_SIZE) * sizeof(int32_t);
+
+      auto launch_v2 = [&](auto ept_tag) {
+        constexpr int EPT = decltype(ept_tag)::value;
+        auto v2_kernel = moe::moe_align_block_size_kernel_v2<scalar_t, EPT>;
+        LaunchKernel(dim3(2), dim3(threads), stream, shared_mem_size)(
+            v2_kernel,
+            topk_ids_ptr,
+            sorted_token_ids_ptr,
+            expert_ids_ptr,
+            num_tokens_post_pad_ptr,
+            (int32_t)num_experts,
+            (int32_t)padded_num_experts,
+            (int32_t)block_size,
+            numel,
+            cumsum_buffer_ptr,
+            pad_sorted_token_ids,
+            (int32_t)max_num_tokens_padded);
+      };
+
+      if (padded_num_experts <= 2048) {
+        launch_v2(std::integral_constant<int, 2>{});
+      } else if (padded_num_experts <= 4096) {
+        launch_v2(std::integral_constant<int, 4>{});
+      } else {
+        launch_v2(std::integral_constant<int, 8>{});
+      }
+
+      const int block_threads = std::min(256, threads);
+      const int num_blocks = (numel + block_threads - 1) / block_threads;
+      const int max_blocks = 65535;
+      const int actual_blocks = std::min(num_blocks, max_blocks);
+
+      auto sort_kernel = moe::count_and_sort_expert_tokens_kernel<scalar_t>;
+      LaunchKernel(dim3(actual_blocks), dim3(block_threads), stream)(
+          sort_kernel, topk_ids_ptr, sorted_token_ids_ptr, cumsum_buffer_ptr, numel);
+    }
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/moe/moe_finalize_fuse_shared.cu

@@ -0,0 +1,418 @@
+// Copyright (c) 2026 LightSeek Foundation
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
+/*
+ * Fused MoE finalize + shared-output add (bf16 output, SM>=90 for PDL).
+ *
+ * Forked from flashinfer's ``finalizeKernel`` and ``finalizeKernelVecLoad``
+ * (trtllm_fused_moe_dev_kernel.cu:639 and :803), stripped of the MoE
+ * backend's KernelParams / UsePdl templating, and extended with an
+ * optional shared_output residual add on the epilogue side.
+ *
+ * For each token t, computes:
+ *     out[t] = Σ_k expert_weights[t, k] * gemm2_out[permuted_idx(t, k)]
+ *            + shared_output[t]                      // if non-null
+ *
+ * Eliminates the native PyTorch ``routed + shared_output`` add (and the
+ * separate ``*= routed_scaling_factor`` kernel when applicable) from
+ * ``DeepseekV3MoE.forward``, and gives the downstream allreduce+rmsnorm
+ * a clean PDL handoff.
+ *
+ * Expert-weight dtype is templated on ``TypeExpW`` so we support both the
+ * bf16 and fp32 topk-weight paths (DSv3/K2.5 trtllm backends use fp32
+ * because their ``_routing_logits_dtype = torch.float32``; other backends
+ * use bf16).
+ *
+ * Expert-weight scale convention: in our target backends
+ * (flashinfer trtllm nvfp4 + unquantized), ``apply_routed_scaling_factor_on_output``
+ * is True, so the routed scaling factor is already folded into
+ * ``expert_weights`` at topk time. This kernel does not apply any
+ * additional scale.
+ */
+
+#include <cutlass/array.h>
+#include <cutlass/numeric_conversion.h>
+#include <cutlass/numeric_types.h>
+
+#include "tvm_ffi_utils.h"
+#include <cuda_runtime.h>
+
+namespace sglang {
+
+using BF16 = cutlass::bfloat16_t;
+
+constexpr int FINALIZE_THREADS_PER_BLOCK = 256;
+constexpr int MAX_TOPK = 64;
+
+// ---------------------------------------------------------------------------
+// General kernel — one CTA per (hidden_chunk, token). Picks up small-to-mid
+// workloads where the block count fits in a few waves.
+// ---------------------------------------------------------------------------
+template <typename TypeExpW>
+__global__ void moeFinalizeKernel(
+    int numTokens,
+    int hiddenDim,
+    int hiddenDimPadded,
+    int topK,
+    BF16 const* __restrict__ inPtr,
+    int const* __restrict__ expandedIdxToPermutedIdx,
+    TypeExpW const* __restrict__ expertWeightsPtr,
+    BF16 const* __restrict__ sharedBiasPtr,
+    BF16* __restrict__ outPtr) {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  cudaGridDependencySynchronize();
+#endif
+
+  for (int64_t tokenIdx = blockIdx.y; tokenIdx < numTokens; tokenIdx += gridDim.y) {
+    for (int64_t hiddenIdx = threadIdx.x + blockDim.x * blockIdx.x; hiddenIdx < hiddenDim;
+         hiddenIdx += blockDim.x * gridDim.x) {
+      float acc = 0.0f;
+      for (int k = 0; k < topK; k++) {
+        int64_t const expandedIdx = tokenIdx * topK + k;
+        int64_t const permutedIdx = expandedIdxToPermutedIdx[expandedIdx];
+        if (permutedIdx == -1) {
+          continue;
+        }
+        float const scale = static_cast<float>(expertWeightsPtr[expandedIdx]);
+        float const val = static_cast<float>(inPtr[permutedIdx * hiddenDimPadded + hiddenIdx]);
+        acc += scale * val;
+      }
+      if (sharedBiasPtr != nullptr) {
+        acc += static_cast<float>(sharedBiasPtr[tokenIdx * hiddenDim + hiddenIdx]);
+      }
+      outPtr[tokenIdx * hiddenDim + hiddenIdx] = static_cast<BF16>(acc);
+    }
+  }
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  cudaTriggerProgrammaticLaunchCompletion();
+#endif
+}
+
+// ---------------------------------------------------------------------------
+// Vectorized-load kernel — one CTA per token, 128-bit loads, topK unrolled.
+// Better at prefill shapes where the general kernel's block count saturates
+// many waves and the indirect gather from gemm2_out dominates.
+// ---------------------------------------------------------------------------
+
+__device__ inline float4 vectorizedLoadPtx(float4 const* ptr) {
+  float4 ret;
+  asm volatile("ld.global.v4.f32 {%0, %1, %2, %3}, [%4];"
+               : "=f"(ret.x), "=f"(ret.y), "=f"(ret.z), "=f"(ret.w)
+               : "l"(ptr));
+  return ret;
+}
+
+template <int TopKUnrollFactor>
+struct IdxPackedTraits;
+template <>
+struct IdxPackedTraits<1> {
+  using Packed = int;
+};
+template <>
+struct IdxPackedTraits<2> {
+  using Packed = int2;
+};
+template <>
+struct IdxPackedTraits<4> {
+  using Packed = int4;
+};
+
+template <typename TypeExpW, int TopKUnrollFactor>
+__global__ void moeFinalizeKernelVecLoad(
+    int numTokens,
+    int hiddenDim,
+    int hiddenDimPadded,
+    int topK,
+    BF16 const* __restrict__ inPtr,
+    int const* __restrict__ expandedIdxToPermutedIdx,
+    TypeExpW const* __restrict__ expertWeightsPtr,
+    BF16 const* __restrict__ sharedBiasPtr,
+    BF16* __restrict__ outPtr) {
+  static_assert(
+      TopKUnrollFactor == 1 || TopKUnrollFactor == 2 || TopKUnrollFactor == 4, "TopKUnrollFactor must be 1, 2, or 4");
+  using IdxPackedType = typename IdxPackedTraits<TopKUnrollFactor>::Packed;
+  using IdxArrayType = cutlass::Array<int, TopKUnrollFactor>;
+  using ScaleArrayType = cutlass::Array<TypeExpW, TopKUnrollFactor>;
+
+  // 128 bits per thread → 8 bf16 elements.
+  constexpr int FINALIZE_ELEM_PER_THREAD = 8;
+  using InputElem = cutlass::Array<BF16, FINALIZE_ELEM_PER_THREAD>;
+  using OutputElem = cutlass::Array<BF16, FINALIZE_ELEM_PER_THREAD>;
+  using ComputeElem = cutlass::Array<float, FINALIZE_ELEM_PER_THREAD>;
+
+  int64_t const tokenIdx = blockIdx.x;
+  int64_t const startOffset = threadIdx.x;
+  int64_t const stride = FINALIZE_THREADS_PER_BLOCK;
+  int64_t const numElemsInPaddedCol = hiddenDimPadded / FINALIZE_ELEM_PER_THREAD;
+  int64_t const numElemsInCol = hiddenDim / FINALIZE_ELEM_PER_THREAD;
+
+  // Stage the per-token (topK/unroll) indices + scales into smem.
+  __shared__ ScaleArrayType scaleArrSmem[MAX_TOPK / TopKUnrollFactor];
+  __shared__ IdxArrayType permutedIdxArrSmem[MAX_TOPK / TopKUnrollFactor];
+
+  for (int kChunkIdx = threadIdx.x; kChunkIdx < topK / TopKUnrollFactor; kChunkIdx += blockDim.x) {
+    int64_t const expandedIdx = tokenIdx * topK + kChunkIdx * TopKUnrollFactor;
+    auto const permutedIdxPacked =
+        reinterpret_cast<IdxPackedType const*>(expandedIdxToPermutedIdx)[expandedIdx / TopKUnrollFactor];
+    permutedIdxArrSmem[kChunkIdx] = *reinterpret_cast<IdxArrayType const*>(&permutedIdxPacked);
+#pragma unroll
+    for (int ki = 0; ki < TopKUnrollFactor; ++ki) {
+      scaleArrSmem[kChunkIdx][ki] = expertWeightsPtr[expandedIdx + ki];
+    }
+  }
+
+  BF16* outputPtr = outPtr + tokenIdx * hiddenDim;
+  auto* outElemPtr = reinterpret_cast<OutputElem*>(outputPtr);
+  auto const* inElemPtr = reinterpret_cast<InputElem const*>(inPtr);
+  auto const* sharedElemPtr =
+      sharedBiasPtr != nullptr ? reinterpret_cast<InputElem const*>(sharedBiasPtr + tokenIdx * hiddenDim) : nullptr;
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  cudaGridDependencySynchronize();
+#endif
+  __syncthreads();
+
+  for (int elemIndex = startOffset; elemIndex < numElemsInCol; elemIndex += stride) {
+    ComputeElem threadOutput;
+    threadOutput.fill(0.0f);
+
+    for (int kChunkIdx = 0; kChunkIdx < topK / TopKUnrollFactor; kChunkIdx++) {
+      IdxArrayType permutedIdxArr = permutedIdxArrSmem[kChunkIdx];
+      InputElem inputElemArr[TopKUnrollFactor];
+#pragma unroll
+      for (int ki = 0; ki < TopKUnrollFactor; ++ki) {
+        int const permutedIdx = permutedIdxArr[ki];
+        if (permutedIdx == -1) {
+          continue;
+        }
+        auto const* inputPermutedPtr = inElemPtr + permutedIdx * numElemsInPaddedCol;
+        float4 input = vectorizedLoadPtx(reinterpret_cast<float4 const*>(&inputPermutedPtr[elemIndex]));
+        inputElemArr[ki] = *reinterpret_cast<InputElem const*>(&input);
+      }
+      ScaleArrayType scaleArr = scaleArrSmem[kChunkIdx];
+#pragma unroll
+      for (int ki = 0; ki < TopKUnrollFactor; ++ki) {
+        int const permutedIdx = permutedIdxArr[ki];
+        if (permutedIdx == -1) {
+          continue;
+        }
+        float const scale = static_cast<float>(scaleArr[ki]);
+        cutlass::NumericArrayConverter<float, BF16, FINALIZE_ELEM_PER_THREAD> toFloat;
+        ComputeElem expertResult = toFloat(inputElemArr[ki]);
+#pragma unroll
+        for (int e = 0; e < FINALIZE_ELEM_PER_THREAD; ++e) {
+          threadOutput[e] += scale * expertResult[e];
+        }
+      }
+    }
+
+    if (sharedElemPtr != nullptr) {
+      float4 shared = vectorizedLoadPtx(reinterpret_cast<float4 const*>(&sharedElemPtr[elemIndex]));
+      InputElem sharedElem = *reinterpret_cast<InputElem const*>(&shared);
+      cutlass::NumericArrayConverter<float, BF16, FINALIZE_ELEM_PER_THREAD> toFloat;
+      ComputeElem sharedFloat = toFloat(sharedElem);
+#pragma unroll
+      for (int e = 0; e < FINALIZE_ELEM_PER_THREAD; ++e) {
+        threadOutput[e] += sharedFloat[e];
+      }
+    }
+
+    cutlass::NumericArrayConverter<BF16, float, FINALIZE_ELEM_PER_THREAD> toBF16;
+    outElemPtr[elemIndex] = toBF16(threadOutput);
+  }
+
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
+  cudaTriggerProgrammaticLaunchCompletion();
+#endif
+}
+
+// ---------------------------------------------------------------------------
+// Typed dispatch
+// ---------------------------------------------------------------------------
+template <typename TypeExpW>
+void dispatchFinalize(
+    int numTokens,
+    int hiddenDim,
+    int hiddenDimPadded,
+    int topK,
+    BF16 const* inPtr,
+    int const* expandedIdxPtr,
+    void const* weightsPtrVoid,
+    BF16 const* sharedPtr,
+    BF16* outPtr,
+    bool useVecLoad,
+    cudaStream_t stream,
+    cudaLaunchAttribute const* attrs,
+    int numAttrs) {
+  auto const* weightsPtr = static_cast<TypeExpW const*>(weightsPtrVoid);
+  constexpr int kNumThreads = 256;
+
+  if (!useVecLoad) {
+    int const numBlocksX = (hiddenDim + kNumThreads - 1) / kNumThreads;
+    int const numBlocksY = std::min(8192, numTokens);
+    cudaLaunchConfig_t config;
+    config.gridDim = dim3(numBlocksX, numBlocksY);
+    config.blockDim = dim3(kNumThreads);
+    config.dynamicSmemBytes = 0;
+    config.stream = stream;
+    config.numAttrs = numAttrs;
+    config.attrs = const_cast<cudaLaunchAttribute*>(attrs);
+
+    cudaLaunchKernelEx(
+        &config,
+        moeFinalizeKernel<TypeExpW>,
+        numTokens,
+        hiddenDim,
+        hiddenDimPadded,
+        topK,
+        inPtr,
+        expandedIdxPtr,
+        weightsPtr,
+        sharedPtr,
+        outPtr);
+    return;
+  }
+
+  auto launch = [&](auto unroll_tag) {
+    constexpr int UNROLL = decltype(unroll_tag)::value;
+    cudaLaunchConfig_t config;
+    config.gridDim = dim3(numTokens);
+    config.blockDim = dim3(FINALIZE_THREADS_PER_BLOCK);
+    config.dynamicSmemBytes = 0;
+    config.stream = stream;
+    config.numAttrs = numAttrs;
+    config.attrs = const_cast<cudaLaunchAttribute*>(attrs);
+    cudaLaunchKernelEx(
+        &config,
+        moeFinalizeKernelVecLoad<TypeExpW, UNROLL>,
+        numTokens,
+        hiddenDim,
+        hiddenDimPadded,
+        topK,
+        inPtr,
+        expandedIdxPtr,
+        weightsPtr,
+        sharedPtr,
+        outPtr);
+  };
+  // Match flashinfer's LAUNCH_TOPK_EXPW dispatch order.
+  if (topK % 4 == 0) {
+    launch(std::integral_constant<int, 4>{});
+  } else if (topK % 2 == 0) {
+    launch(std::integral_constant<int, 2>{});
+  } else {
+    launch(std::integral_constant<int, 1>{});
+  }
+}
+
+}  // namespace sglang
+
+// ---------------------------------------------------------------------------
+// Host launcher
+// ---------------------------------------------------------------------------
+void moe_finalize_fuse_shared(
+    TensorView out,
+    TensorView gemm2_out,
+    TensorView expanded_idx_to_permuted_idx,
+    TensorView expert_weights,
+    TensorView shared_output,
+    int64_t top_k,
+    bool enable_pdl) {
+  TVM_FFI_ICHECK_EQ(out.ndim(), 2) << "out must be 2-D [numTokens, hiddenDim]";
+  TVM_FFI_ICHECK_EQ(gemm2_out.ndim(), 2) << "gemm2_out must be 2-D [totalNumPaddedTokens, hiddenDimPadded]";
+  TVM_FFI_ICHECK_EQ(expanded_idx_to_permuted_idx.ndim(), 1);
+  TVM_FFI_ICHECK_EQ(expert_weights.ndim(), 2) << "expert_weights must be 2-D [numTokens, topK]";
+
+  int const numTokens = int(out.size(0));
+  int const hiddenDim = int(out.size(1));
+  int const hiddenDimPadded = int(gemm2_out.size(1));
+  TVM_FFI_ICHECK_LE(top_k, sglang::MAX_TOPK);
+  TVM_FFI_ICHECK_EQ(expanded_idx_to_permuted_idx.size(0), numTokens * top_k);
+  TVM_FFI_ICHECK_EQ(expert_weights.size(0), numTokens);
+  TVM_FFI_ICHECK_EQ(expert_weights.size(1), top_k);
+
+  bool const hasShared = shared_output.numel() > 0;
+  if (hasShared) {
+    TVM_FFI_ICHECK_EQ(shared_output.ndim(), 2);
+    TVM_FFI_ICHECK_EQ(shared_output.size(0), numTokens);
+    TVM_FFI_ICHECK_EQ(shared_output.size(1), hiddenDim);
+  }
+
+  auto const* inPtr = static_cast<sglang::BF16 const*>(gemm2_out.data_ptr());
+  auto const* expandedIdxPtr = static_cast<int const*>(expanded_idx_to_permuted_idx.data_ptr());
+  auto const* sharedPtr = hasShared ? static_cast<sglang::BF16 const*>(shared_output.data_ptr()) : nullptr;
+  auto* outPtr = static_cast<sglang::BF16*>(out.data_ptr());
+
+  cudaSetDevice(out.device().device_id);
+  cudaStream_t const stream = get_stream(out.device());
+
+  // Dispatch heuristic (matches flashinfer): few waves → general kernel,
+  // many waves → vectorized. The 1184 threshold comes from 148 SMs × 8
+  // blocks/SM on Blackwell.
+  constexpr int kNumThreads = 256;
+  int const numBlocksX = (hiddenDim + kNumThreads - 1) / kNumThreads;
+  int const numBlocksY = std::min(8192, numTokens);
+  bool const useVecLoad = (numBlocksX * numBlocksY) >= 1184 && (hiddenDim % 8 == 0) && (hiddenDimPadded % 8 == 0);
+
+  cudaLaunchAttribute attrs[1];
+  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attrs[0].val.programmaticStreamSerializationAllowed = enable_pdl;
+
+  auto ew_dtype = expert_weights.dtype();
+  if (ew_dtype == DLDataType{kDLFloat, 32, 1}) {
+    sglang::dispatchFinalize<float>(
+        numTokens,
+        hiddenDim,
+        hiddenDimPadded,
+        int(top_k),
+        inPtr,
+        expandedIdxPtr,
+        expert_weights.data_ptr(),
+        sharedPtr,
+        outPtr,
+        useVecLoad,
+        stream,
+        attrs,
+        1);
+  } else if (ew_dtype == DLDataType{kDLBfloat, 16, 1}) {
+    sglang::dispatchFinalize<sglang::BF16>(
+        numTokens,
+        hiddenDim,
+        hiddenDimPadded,
+        int(top_k),
+        inPtr,
+        expandedIdxPtr,
+        expert_weights.data_ptr(),
+        sharedPtr,
+        outPtr,
+        useVecLoad,
+        stream,
+        attrs,
+        1);
+  } else {
+    TVM_FFI_ICHECK(false) << "expert_weights dtype must be float32 or bfloat16";
+  }
+
+  cudaError_t const err = cudaGetLastError();
+  TVM_FFI_ICHECK(err == cudaSuccess) << "moe_finalize_fuse_shared launch failed: " << cudaGetErrorString(err);
+}
+
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(moe_finalize_fuse_shared, moe_finalize_fuse_shared);

python/sglang/jit_kernel/csrc/moe/moe_fused_gate.cuh

@@ -0,0 +1,367 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cfloat>
+#include <cstdint>
+
+namespace {
+
+constexpr uint32_t kWarpSize = 32;
+constexpr uint32_t kWarpsPerCTA = 6;
+constexpr uint32_t kSmallTokenThreshold = 512;
+constexpr uint32_t kMaxExperts = 512;
+constexpr uint32_t kMaxTopK = 16;
+
+enum class ScoringFunc : uint32_t {
+  kSigmoid = 0,
+  kSqrtSoftplus = 1,
+};
+
+struct MoEFusedGateParams {
+  const float* __restrict__ input;
+  const float* __restrict__ bias;
+  float* __restrict__ output;
+  int32_t* __restrict__ indices;
+  uint32_t num_rows;
+  uint32_t num_experts;
+  uint32_t topk;
+  uint32_t num_fused_shared_experts;
+  bool renormalize;
+  float routed_scaling_factor;
+  bool apply_routed_scaling_factor_on_output;
+};
+
+template <ScoringFunc kScoringFunc>
+__device__ __forceinline__ float compute_score(float x) {
+  if constexpr (kScoringFunc == ScoringFunc::kSigmoid) {
+    // sigmoid(x) = 1 / (1 + exp(-x))
+    return 1.0f / (1.0f + expf(-x));
+  } else {
+    // sqrt(softplus(x)) = sqrt(log(1 + exp(x)))
+    float softplus = log1pf(expf(x));
+    return sqrtf(softplus);
+  }
+}
+
+template <uint32_t kWarpsPerToken, ScoringFunc kScoringFunc>
+__global__ void moe_fused_gate_kernel_small_token(const MoEFusedGateParams __grid_constant__ params) {
+  const auto& [input, bias, output, indices, num_rows, num_experts, topk, num_fused_shared_experts, renormalize, routed_scaling_factor, apply_routed_scaling_factor_on_output] =
+      params;
+
+  uint32_t row_idx = blockIdx.x;
+  if (row_idx >= num_rows) return;
+
+  // number of routed experts to select (excluding fused shared experts)
+  const uint32_t topk_routed = topk - num_fused_shared_experts;
+
+  uint32_t tid = threadIdx.x;
+  uint32_t warp_id = tid / kWarpSize;
+  uint32_t lane_id = tid % kWarpSize;
+  // Actual warps launched (<= kWarpsPerToken). num_experts that need fewer than
+  // kWarpsPerToken warps leave the upper warp_maxs/warp_experts slots unwritten,
+  // so the cross-warp reduction below must only read the launched warps.
+  const uint32_t num_warps = blockDim.x / kWarpSize;
+
+  extern __shared__ float shared_mem[];
+  float* shared_scores = shared_mem;
+  float* shared_original_scores = shared_mem + num_experts;
+
+  // For warp-level reduction
+  __shared__ float warp_maxs[kWarpsPerToken];
+  __shared__ int warp_experts[kWarpsPerToken];
+  __shared__ int selected_experts[kMaxTopK];
+
+  for (uint32_t e = tid; e < num_experts; e += blockDim.x) {
+    float input_val = input[row_idx * num_experts + e];
+    float bias_val = bias[e];
+    float score_val = compute_score<kScoringFunc>(input_val);
+    float biased_val = score_val + bias_val;
+    shared_scores[e] = biased_val;
+    shared_original_scores[e] = score_val;
+  }
+
+  __syncthreads();
+
+  // only select topk_routed experts (excluding shared experts)
+  for (uint32_t k = 0; k < topk_routed; k++) {
+    float my_val = -FLT_MAX;
+    int my_expert = -1;
+    for (uint32_t e = tid; e < num_experts; e += blockDim.x) {
+      if (shared_scores[e] > my_val) {
+        my_val = shared_scores[e];
+        my_expert = e;
+      }
+    }
+
+    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);
+      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;
+      }
+    }
+
+    if (lane_id == 0 && warp_id < kWarpsPerToken) {
+      warp_maxs[warp_id] = warp_max_val;
+      warp_experts[warp_id] = warp_max_expert;
+    }
+
+    __syncthreads();
+
+    if (warp_id == 0) {
+      float final_max = (lane_id < num_warps) ? warp_maxs[lane_id] : -FLT_MAX;
+      int final_expert = (lane_id < num_warps) ? 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) {
+        selected_experts[k] = final_expert;
+      }
+    }
+
+    __syncthreads();
+
+    int selected = selected_experts[k];
+    if (selected >= 0 && tid == 0) {
+      shared_scores[selected] = -FLT_MAX;
+    }
+
+    __syncthreads();
+  }
+
+  static_assert(kMaxTopK <= device::kWarpThreads);
+  if (tid >= device::kWarpThreads) return;
+
+  // only use the first warp to perform write to global operation
+  float routed_weight = 0.0f;
+  int32_t selected_expert = 0;
+  if (tid < topk_routed) {
+    int expert_id = selected_experts[tid];
+    float score = shared_original_scores[expert_id];
+    if (expert_id >= 0 && expert_id < static_cast<int>(num_experts)) {
+      routed_weight = score;
+      selected_expert = expert_id;
+    }
+  }
+  const auto routed_sum = device::warp::reduce_sum<kMaxTopK>(routed_weight);
+  if (tid < topk) {
+    const bool is_shared = tid >= topk_routed;
+    const auto output_offset = row_idx * topk + tid;
+    const auto weight = is_shared ? (routed_sum / routed_scaling_factor) : routed_weight;
+    const auto expert_id = is_shared ? (num_experts + tid - topk_routed) : selected_expert;
+    const auto scale = apply_routed_scaling_factor_on_output ? routed_scaling_factor : 1.0f;
+    const auto norm = renormalize && routed_sum > 0.0f ? routed_sum : 1.0f;
+    output[output_offset] = weight / norm * scale;
+    indices[output_offset] = expert_id;
+  }
+}
+
+template <ScoringFunc kScoringFunc>
+__global__ void moe_fused_gate_kernel(const MoEFusedGateParams __grid_constant__ params) {
+  const auto& [input, bias, output, indices, num_rows, num_experts, topk, num_fused_shared_experts, renormalize, routed_scaling_factor, apply_routed_scaling_factor_on_output] =
+      params;
+
+  uint32_t row_idx = blockIdx.x * kWarpsPerCTA + threadIdx.y;
+  if (row_idx >= num_rows) return;
+
+  // number of routed experts to select (excluding fused shared experts)
+  const uint32_t topk_routed = topk - num_fused_shared_experts;
+
+  uint32_t lane_id = threadIdx.x;
+  uint32_t warp_id = threadIdx.y;
+
+  extern __shared__ float shared_mem[];
+  float* shared_scores = shared_mem + warp_id * num_experts * 2;
+  float* shared_original_scores = shared_scores + num_experts;
+  __shared__ int selected_experts[kWarpsPerCTA][kMaxTopK];
+  int* warp_selected_experts = selected_experts[warp_id];
+
+  for (uint32_t e = lane_id; e < num_experts; e += kWarpSize) {
+    float input_val = input[row_idx * num_experts + e];
+    float bias_val = bias[e];
+    float score_val = compute_score<kScoringFunc>(input_val);
+    float biased_val = score_val + bias_val;
+    shared_scores[e] = biased_val;
+    shared_original_scores[e] = score_val;
+  }
+
+  __syncwarp();
+
+  // only select topk_routed experts
+  for (uint32_t k = 0; k < topk_routed; k++) {
+    float max_val = -FLT_MAX;
+    int max_expert = -1;
+
+    for (uint32_t expert = lane_id; expert < num_experts; expert += kWarpSize) {
+      if (shared_scores[expert] > max_val) {
+        max_val = shared_scores[expert];
+        max_expert = expert;
+      }
+    }
+
+    for (int offset = kWarpSize / 2; 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;
+      }
+    }
+
+    if (lane_id == 0) {
+      warp_selected_experts[k] = max_expert;
+      if (max_expert != -1) {
+        shared_scores[max_expert] = -FLT_MAX;
+      }
+    }
+
+    __syncwarp();
+  }
+
+  static_assert(kMaxTopK <= device::kWarpThreads);
+
+  float routed_weight = 0.0f;
+  int32_t selected_expert = 0;
+  if (lane_id < topk_routed) {
+    int expert_id = warp_selected_experts[lane_id];
+    if (expert_id >= 0 && expert_id < static_cast<int>(num_experts)) {
+      routed_weight = shared_original_scores[expert_id];
+      selected_expert = expert_id;
+    }
+  }
+  const auto routed_sum = device::warp::reduce_sum<kMaxTopK>(routed_weight);
+  if (lane_id < topk) {
+    const bool is_shared = lane_id >= topk_routed;
+    const auto output_idx = row_idx * topk + lane_id;
+    const auto weight = is_shared ? (routed_sum / routed_scaling_factor) : routed_weight;
+    const auto expert_id = is_shared ? (num_experts + lane_id - topk_routed) : selected_expert;
+    const auto scale = apply_routed_scaling_factor_on_output ? routed_scaling_factor : 1.0f;
+    const auto norm = renormalize && routed_sum > 0.0f ? routed_sum : 1.0f;
+    output[output_idx] = weight / norm * scale;
+    indices[output_idx] = expert_id;
+  }
+}
+
+template <ScoringFunc kScoringFunc>
+void dispatch_small_token_kernel(
+    uint32_t num_rows,
+    uint32_t threads_per_block,
+    uint32_t warps_per_token,
+    DLDevice device,
+    size_t smem_per_row,
+    const MoEFusedGateParams& params) {
+  using namespace host;
+  if (warps_per_token <= 8) {
+    LaunchKernel(num_rows, threads_per_block, device, smem_per_row)(
+        moe_fused_gate_kernel_small_token<8, kScoringFunc>, params);
+  } else if (warps_per_token <= 12) {
+    LaunchKernel(num_rows, threads_per_block, device, smem_per_row)(
+        moe_fused_gate_kernel_small_token<12, kScoringFunc>, params);
+  } else {
+    LaunchKernel(num_rows, threads_per_block, device, smem_per_row)(
+        moe_fused_gate_kernel_small_token<16, kScoringFunc>, params);
+  }
+}
+
+struct MoEFusedGateKernel {
+  static void
+  run(const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView bias,
+      const tvm::ffi::TensorView output,
+      const tvm::ffi::TensorView indices,
+      uint32_t topk,
+      uint32_t scoring_func,  // 0 = sigmoid, 1 = sqrtsoftplus
+      uint32_t num_fused_shared_experts,
+      bool renormalize,
+      float routed_scaling_factor,
+      bool apply_routed_scaling_factor_on_output) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_rows"};
+    auto E = SymbolicSize{"num_experts"};
+    auto K = SymbolicSize{"topk"};
+    auto device = SymbolicDevice{};
+    K.set_value(topk);
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, E}).with_dtype<float>().with_device(device).verify(input);
+    TensorMatcher({E}).with_dtype<float>().with_device(device).verify(bias);
+    TensorMatcher({N, K}).with_dtype<float>().with_device(device).verify(output);
+    TensorMatcher({N, K}).with_dtype<int32_t>().with_device(device).verify(indices);
+
+    const auto num_rows = static_cast<uint32_t>(N.unwrap());
+    const auto num_experts = static_cast<uint32_t>(E.unwrap());
+
+    RuntimeCheck(num_experts <= kMaxExperts, "num_experts exceeds maximum supported value");
+    RuntimeCheck(scoring_func <= 1, "scoring_func must be 0 (sigmoid) or 1 (sqrtsoftplus)");
+    RuntimeCheck(topk > num_fused_shared_experts, "topk must be greater than num_fused_shared_experts");
+
+    const auto params = MoEFusedGateParams{
+        .input = static_cast<const float*>(input.data_ptr()),
+        .bias = static_cast<const float*>(bias.data_ptr()),
+        .output = static_cast<float*>(output.data_ptr()),
+        .indices = static_cast<int32_t*>(indices.data_ptr()),
+        .num_rows = num_rows,
+        .num_experts = num_experts,
+        .topk = topk,
+        .num_fused_shared_experts = num_fused_shared_experts,
+        .renormalize = renormalize,
+        .routed_scaling_factor = routed_scaling_factor,
+        .apply_routed_scaling_factor_on_output = apply_routed_scaling_factor_on_output,
+    };
+
+    const size_t smem_per_row = 2 * num_experts * sizeof(float);
+
+    bool use_small_token_kernel = num_rows <= kSmallTokenThreshold;
+
+    if (use_small_token_kernel) {
+      // 1 token per block
+      uint32_t warps_per_token = div_ceil(num_experts, kWarpSize);
+      warps_per_token = std::min(warps_per_token, 16u);
+      uint32_t threads_per_block = warps_per_token * kWarpSize;
+
+      if (scoring_func == 0) {
+        dispatch_small_token_kernel<ScoringFunc::kSigmoid>(
+            num_rows, threads_per_block, warps_per_token, device.unwrap(), smem_per_row, params);
+      } else {
+        dispatch_small_token_kernel<ScoringFunc::kSqrtSoftplus>(
+            num_rows, threads_per_block, warps_per_token, device.unwrap(), smem_per_row, params);
+      }
+    } else {
+      // multiple tokens per block
+      uint32_t num_blocks = div_ceil(num_rows, kWarpsPerCTA);
+      dim3 block_dim(kWarpSize, kWarpsPerCTA);
+      size_t large_smem = smem_per_row * kWarpsPerCTA;
+
+      if (scoring_func == 0) {
+        LaunchKernel(num_blocks, block_dim, device.unwrap(), large_smem)(
+            moe_fused_gate_kernel<ScoringFunc::kSigmoid>, params);
+      } else {
+        LaunchKernel(num_blocks, block_dim, device.unwrap(), large_smem)(
+            moe_fused_gate_kernel<ScoringFunc::kSqrtSoftplus>, params);
+      }
+    }
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/moe/moe_topk_sigmoid.cuh

@@ -0,0 +1,544 @@
+
+// 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
+#pragma once
+
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <cub/cub.cuh>
+#include <cub/util_type.cuh>
+#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>
+using MaxReduceOp = cuda::maximum<>;
+using MinReduceOp = cuda::minimum<>;
+#else
+using MaxReduceOp = cub::Max;
+using MinReduceOp = cub::Min;
+#endif
+
+#include <cfloat>
+#include <cstdint>
+#include <type_traits>
+
+using tvm::ffi::TensorView;
+
+#ifndef MOE_TOPK_SIGMOID_WARP_SIZE
+#define MOE_TOPK_SIGMOID_WARP_SIZE 32
+#endif
+
+namespace {
+
+static constexpr int WARP_SIZE = MOE_TOPK_SIGMOID_WARP_SIZE;
+
+#define MAX(a, b) ((a) > (b) ? (a) : (b))
+#define MIN(a, b) ((a) < (b) ? (a) : (b))
+
+// ---------------------------------------------------------------------------
+// Aligned array — avoids CUTLASS dependency; identical semantics.
+// ---------------------------------------------------------------------------
+template <typename T, int N, int Alignment = static_cast<int>(sizeof(T) * N)>
+class alignas(Alignment) AlignedArray {
+  T data[N];
+};
+
+// ---------------------------------------------------------------------------
+// Type conversion helper
+// ---------------------------------------------------------------------------
+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 {
+    return static_cast<float>(x);
+  }
+}
+
+// ---------------------------------------------------------------------------
+// moeSigmoid — fallback sigmoid kernel (used for non-power-of-2 experts)
+// ---------------------------------------------------------------------------
+template <typename T, int TPB>
+__launch_bounds__(TPB) __global__
+    void moeSigmoid(const T* input, const bool* finished, float* output, const int num_cols) {
+  const int thread_row_offset = blockIdx.x * num_cols;
+
+  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]);
+    val = 1.0f / (1.0f + expf(-val));
+    output[idx] = val;
+  }
+}
+
+// ---------------------------------------------------------------------------
+// moeTopK — fallback top-k kernel (used for non-power-of-2 experts)
+// ---------------------------------------------------------------------------
+template <int TPB>
+__launch_bounds__(TPB) __global__ void moeTopK(
+    const float* inputs_after_sigmoid,
+    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,
+    const float* correction_bias,
+    double routed_scaling_factor,
+    int num_fused_shared_experts) {
+  using cub_kvp = cub::KeyValuePair<int, float>;
+  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 int topk = k + num_fused_shared_experts;
+
+  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_sigmoid[idx];
+      if (correction_bias != nullptr) {
+        inp_kvp.value += correction_bias[expert];
+      }
+
+      for (int prior_k = 0; prior_k < k_idx; ++prior_k) {
+        const int prior_winning_expert = indices[topk * block_row + prior_k];
+        if (prior_winning_expert == expert) {
+          inp_kvp = thread_kvp;
+        }
+      }
+
+      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 = topk * block_row + k_idx;
+      float val;
+      if (correction_bias != nullptr) {
+        val = inputs_after_sigmoid[thread_read_offset + expert];
+      } else {
+        val = result_kvp.value;
+      }
+      output[idx] = val;
+      indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
+      assert(indices[idx] >= 0);
+      row_sum_for_renormalize += val;
+    }
+    __syncthreads();
+  }
+
+  if (num_fused_shared_experts > 0 && threadIdx.x == 0) {
+    const int last_idx = topk * block_row + k;
+    if (renormalize) {
+      output[last_idx] = 1.0f;
+    } else {
+      output[last_idx] = row_sum_for_renormalize / routed_scaling_factor;
+    }
+    indices[last_idx] = num_experts;
+  }
+  if (renormalize && threadIdx.x == 0) {
+    float row_sum_for_renormalize_inv = routed_scaling_factor / (row_sum_for_renormalize + 1e-20f);
+    for (int k_idx = 0; k_idx < k; ++k_idx) {
+      const int idx = topk * block_row + k_idx;
+      output[idx] = output[idx] * row_sum_for_renormalize_inv;
+    }
+  }
+}
+
+// ---------------------------------------------------------------------------
+// topkGatingSigmoid — optimised kernel for power-of-2 expert counts
+// ---------------------------------------------------------------------------
+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 topkGatingSigmoid(
+    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* correction_bias,
+    double routed_scaling_factor,
+    int num_fused_shared_experts) {
+  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, "");
+  static_assert(WARP_SIZE % THREADS_PER_ROW == 0, "");
+  static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW), "");
+  static_assert(THREADS_PER_ROW <= 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, "");
+
+  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;
+  const int topk = k + num_fused_shared_experts;
+
+  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) {
+    float val = convert_to_float<T>(row_chunk_temp[ii]);
+    val = 1.0f / (1.0f + expf(-val));
+    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;
+  }
+
+  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 = __shfl_xor_sync(0xffffffff, max_val, mask, THREADS_PER_ROW);
+      int other_expert = __shfl_xor_sync(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 = topk * thread_row + k_idx;
+      float out_val;
+      if (correction_bias != nullptr) {
+        out_val = convert_to_float<T>(thread_row_ptr[expert]);
+        out_val = 1.0f / (1.0f + expf(-out_val));
+      } else {
+        out_val = max_val;
+      }
+      output[idx] = out_val;
+      indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
+      row_sum_for_renormalize += out_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 (num_fused_shared_experts > 0 && thread_group_idx == 0) {
+    const int last_idx = topk * thread_row + k;
+    if (renormalize) {
+      output[last_idx] = 1.0f;
+    } else {
+      output[last_idx] = row_sum_for_renormalize / routed_scaling_factor;
+    }
+    indices[last_idx] = NUM_EXPERTS;
+  }
+  if (renormalize && thread_group_idx == 0) {
+    float row_sum_for_renormalize_inv = routed_scaling_factor / (row_sum_for_renormalize + 1e-20f);
+#pragma unroll
+    for (int k_idx = 0; k_idx < k; ++k_idx) {
+      const int idx = topk * thread_row + k_idx;
+      output[idx] = output[idx] * row_sum_for_renormalize_inv;
+    }
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Compile-time constants helper
+// ---------------------------------------------------------------------------
+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 = 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
+
+// ---------------------------------------------------------------------------
+// Per-expert-count launcher helper
+// ---------------------------------------------------------------------------
+template <typename T, int EXPERTS, int WARPS_PER_TB>
+void topkGatingSigmoidLauncherHelper(
+    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* correction_bias,
+    double routed_scaling_factor,
+    int num_fused_shared_experts,
+    cudaStream_t stream) {
+  static constexpr std::size_t MAX_BYTES_PER_LDG = 16;
+  static constexpr int BYTES_PER_LDG = 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);
+  topkGatingSigmoid<T, VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG><<<num_blocks, block_dim, 0, stream>>>(
+      input,
+      finished,
+      output,
+      num_rows,
+      indices,
+      k,
+      start_expert,
+      end_expert,
+      renormalize,
+      correction_bias,
+      routed_scaling_factor,
+      num_fused_shared_experts);
+}
+
+// ---------------------------------------------------------------------------
+// Dispatch macro — used inside topkGatingSigmoidKernelLauncher
+// ---------------------------------------------------------------------------
+#define LAUNCH_SIGMOID(TYPE, NUM_EXPERTS, WARPS_PER_TB)             \
+  topkGatingSigmoidLauncherHelper<TYPE, NUM_EXPERTS, WARPS_PER_TB>( \
+      gating_output,                                                \
+      nullptr,                                                      \
+      topk_weights,                                                 \
+      topk_indices,                                                 \
+      num_tokens,                                                   \
+      topk - num_fused_shared_experts,                              \
+      0,                                                            \
+      num_experts,                                                  \
+      renormalize,                                                  \
+      correction_bias,                                              \
+      routed_scaling_factor,                                        \
+      num_fused_shared_experts,                                     \
+      stream)
+
+// ---------------------------------------------------------------------------
+// Main launcher: dispatches on num_experts
+// ---------------------------------------------------------------------------
+template <typename T>
+void topkGatingSigmoidKernelLauncher(
+    const T* gating_output,
+    float* topk_weights,
+    int* topk_indices,
+    float* sigmoid_workspace,
+    const int num_tokens,
+    const int num_experts,
+    const int topk,
+    const bool renormalize,
+    const float* correction_bias,
+    double routed_scaling_factor,
+    int num_fused_shared_experts,
+    cudaStream_t stream) {
+  static constexpr int WARPS_PER_TB = 4;
+  switch (num_experts) {
+    case 1:
+      LAUNCH_SIGMOID(T, 1, WARPS_PER_TB);
+      break;
+    case 2:
+      LAUNCH_SIGMOID(T, 2, WARPS_PER_TB);
+      break;
+    case 4:
+      LAUNCH_SIGMOID(T, 4, WARPS_PER_TB);
+      break;
+    case 8:
+      LAUNCH_SIGMOID(T, 8, WARPS_PER_TB);
+      break;
+    case 16:
+      LAUNCH_SIGMOID(T, 16, WARPS_PER_TB);
+      break;
+    case 32:
+      LAUNCH_SIGMOID(T, 32, WARPS_PER_TB);
+      break;
+    case 64:
+      LAUNCH_SIGMOID(T, 64, WARPS_PER_TB);
+      break;
+    case 128:
+      LAUNCH_SIGMOID(T, 128, WARPS_PER_TB);
+      break;
+    case 256:
+      LAUNCH_SIGMOID(T, 256, WARPS_PER_TB);
+      break;
+    default: {
+      // Fallback: non-power-of-2 or >256 experts
+      using namespace host;
+      RuntimeCheck(
+          sigmoid_workspace != nullptr, "sigmoid_workspace must be provided for num_experts that are not a power of 2");
+      static constexpr int TPB = 256;
+      moeSigmoid<T, TPB><<<num_tokens, TPB, 0, stream>>>(gating_output, nullptr, sigmoid_workspace, num_experts);
+      moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
+          sigmoid_workspace,
+          nullptr,
+          topk_weights,
+          topk_indices,
+          num_experts,
+          topk - num_fused_shared_experts,
+          0,
+          num_experts,
+          renormalize,
+          correction_bias,
+          routed_scaling_factor,
+          num_fused_shared_experts);
+    }
+  }
+}
+
+#undef LAUNCH_SIGMOID
+
+}  // namespace
+
+// ---------------------------------------------------------------------------
+// Host launcher (tvm-ffi interface)
+// ---------------------------------------------------------------------------
+template <typename T>
+void topk_sigmoid(
+    TensorView gating_output,
+    TensorView topk_weights,
+    TensorView topk_ids,
+    TensorView workspace,
+    bool renormalize,
+    tvm::ffi::Optional<TensorView> correction_bias,
+    double routed_scaling_factor,
+    int num_fused_shared_experts) {
+  using namespace host;
+
+  // --- Input validation ---
+  RuntimeCheck(gating_output.dim() == 2, "gating_output must be 2-D");
+  RuntimeCheck(topk_weights.dim() == 2, "topk_weights must be 2-D");
+  RuntimeCheck(topk_ids.dim() == 2, "topk_ids must be 2-D");
+
+  const int64_t num_tokens = gating_output.shape()[0];
+  const int64_t num_experts = gating_output.shape()[1];
+  const int64_t topk = topk_weights.shape()[1];
+
+  RuntimeCheck(
+      topk_weights.shape()[0] == num_tokens && topk_ids.shape()[0] == num_tokens,
+      "topk_weights and topk_ids must have num_tokens rows");
+  RuntimeCheck(topk_ids.shape()[1] == topk, "topk_ids second dim must match topk_weights");
+  RuntimeCheck(topk <= num_experts, "topk must be <= num_experts");
+  RuntimeCheck(num_fused_shared_experts <= 1, "num_fused_shared_experts must be <= 1");
+
+  // correction_bias validation
+  if (correction_bias.has_value()) {
+    const auto& bias = correction_bias.value();
+    RuntimeCheck(bias.dim() == 1, "correction_bias must be 1-D");
+    RuntimeCheck(bias.shape()[0] == num_experts, "correction_bias size must equal num_experts");
+    RuntimeCheck(
+        bias.dtype().code == DLDataTypeCode::kDLFloat && bias.dtype().bits == 32, "correction_bias must be float32");
+  }
+
+  const T* gating_ptr = static_cast<const T*>(gating_output.data_ptr());
+  float* weights_ptr = static_cast<float*>(topk_weights.data_ptr());
+  int* indices_ptr = static_cast<int*>(topk_ids.data_ptr());
+  float* workspace_ptr = static_cast<float*>(workspace.data_ptr());
+  const float* bias_ptr =
+      correction_bias.has_value() ? static_cast<const float*>(correction_bias.value().data_ptr()) : nullptr;
+
+  cudaStream_t stream = LaunchKernel::resolve_device(gating_output.device());
+
+  topkGatingSigmoidKernelLauncher<T>(
+      gating_ptr,
+      weights_ptr,
+      indices_ptr,
+      workspace_ptr,
+      static_cast<int>(num_tokens),
+      static_cast<int>(num_experts),
+      static_cast<int>(topk),
+      renormalize,
+      bias_ptr,
+      routed_scaling_factor,
+      num_fused_shared_experts,
+      stream);
+}

python/sglang/jit_kernel/csrc/moe/nvfp4_blockwise_moe.cuh

@@ -0,0 +1,882 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/utils.cuh>
+
+#include <cutlass/arch/arch.h>
+#include <cutlass/cutlass.h>
+
+#include "cute/tensor.hpp"
+#include "cutlass/epilogue/collective/collective_builder.hpp"
+#include "cutlass/epilogue/collective/default_epilogue.hpp"
+#include "cutlass/epilogue/thread/linear_combination.h"
+#include "cutlass/gemm/collective/collective_builder.hpp"
+#include "cutlass/gemm/device/gemm_universal_adapter.h"
+#include "cutlass/gemm/dispatch_policy.hpp"
+#include "cutlass/gemm/group_array_problem_shape.hpp"
+#include "cutlass/gemm/kernel/gemm_universal.hpp"
+#include "cutlass/tensor_ref.h"
+#include "cutlass/util/command_line.h"
+#include "cutlass/util/distribution.h"
+#include "cutlass/util/host_tensor.h"
+#include "cutlass/util/packed_stride.hpp"
+#include "cutlass/util/reference/device/gemm.h"
+#include "cutlass/util/reference/device/tensor_compare.h"
+#include "cutlass/util/reference/host/gett.hpp"
+#include "cutlass/util/reference/host/tensor_compare.h"
+#include "cutlass/util/reference/host/tensor_fill.h"
+#include "cutlass/util/reference/host/tensor_norm.h"
+#include "cutlass/util/tensor_view_io.h"
+#include <algorithm>
+#include <cassert>
+#include <cstdint>
+#include <limits>
+#include <unordered_map>
+
+using namespace host;
+using namespace cute;
+
+struct WorkspaceKey {
+  int device_id;
+  uintptr_t stream;
+  auto operator==(const WorkspaceKey&) const -> bool = default;
+};
+
+struct WorkspaceKeyHash {
+  auto operator()(const WorkspaceKey& key) const -> size_t {
+    size_t h1 = std::hash<int>{}(key.device_id);
+    size_t h2 = std::hash<uintptr_t>{}(key.stream);
+    return h1 ^ (h2 + 0x9e3779b97f4a7c15ULL + (h1 << 6) + (h1 >> 2));
+  }
+};
+
+struct WorkspaceState {
+  void* ptr = nullptr;
+  size_t bytes = 0;
+};
+
+inline auto get_cached_workspace(size_t required_bytes, int device_id, cudaStream_t stream) -> void* {
+  if (required_bytes == 0) {
+    return nullptr;
+  }
+
+  thread_local std::unordered_map<WorkspaceKey, WorkspaceState, WorkspaceKeyHash> cache;
+  WorkspaceKey key{device_id, reinterpret_cast<uintptr_t>(stream)};
+  auto& ws = cache[key];
+
+  if (ws.ptr != nullptr && ws.bytes >= required_bytes) {
+    return ws.ptr;
+  }
+
+  RuntimeDeviceCheck(cudaSetDevice(device_id));
+  if (ws.ptr != nullptr) {
+    RuntimeDeviceCheck(cudaFreeAsync(ws.ptr, stream));
+    ws.ptr = nullptr;
+    ws.bytes = 0;
+  }
+  RuntimeDeviceCheck(cudaMallocAsync(&ws.ptr, required_bytes, stream));
+  ws.bytes = required_bytes;
+  return ws.ptr;
+}
+
+inline int getSMVersion(int device_id) {
+  int sm_major = 0;
+  int sm_minor = 0;
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_major, cudaDevAttrComputeCapabilityMajor, device_id));
+  RuntimeDeviceCheck(cudaDeviceGetAttribute(&sm_minor, cudaDevAttrComputeCapabilityMinor, device_id));
+  return sm_major * 10 + sm_minor;
+}
+
+template <
+    typename ElementAB,
+    typename ElementC,
+    typename ElementSF,
+    typename ElementAccumulator,
+    typename LayoutSFA,
+    typename LayoutSFB,
+    typename ScaleConfig>
+__global__ void __get_group_gemm_starts(
+    ElementAB** a_offsets,
+    ElementAB** b_offsets,
+    ElementC** out_offsets,
+    ElementSF** a_scales_offsets,
+    ElementSF** b_scales_offsets,
+    ElementAccumulator** alpha_offsets,
+    LayoutSFA* layout_sfa_base_as_int,
+    LayoutSFB* layout_sfb_base_as_int,
+    ElementAB* a_base_as_int,
+    ElementAB* b_base_as_int,
+    ElementC* out_base_as_int,
+    ElementSF* a_scales_base_as_int,
+    ElementSF* b_scales_base_as_int,
+    ElementAccumulator* alphas_base_as_int,
+    const int32_t* expert_offsets,
+    const int32_t* sf_offsets,
+    const int32_t* problem_sizes_as_shapes,
+    const int K,
+    const int N) {
+  int64_t expert_id = threadIdx.x;
+  if (expert_id >= gridDim.x * blockDim.x) {
+    return;
+  }
+  // Originally int32_t but upcasting to int64_t to avoid overflow
+  // during offset calculations
+  int64_t expert_offset = static_cast<int64_t>(expert_offsets[expert_id]);
+  int64_t sf_offset = static_cast<int64_t>(sf_offsets[expert_id]);
+  // size for block in block scale.
+  int64_t group_size = 16;
+  int64_t m = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3]);
+  int64_t n = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 1]);
+  int64_t k = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 2]);
+  assert((m >= 0 && n == N && k == K && k % 2 == 0) && "unexpected problem sizes");
+
+  int64_t half_k = static_cast<int64_t>(k / 2);
+  int64_t group_k = static_cast<int64_t>(k / group_size);
+  // Shape of A as uint8/byte = [M, K // 2]
+  // Shape of B as uint8/byte = [E, N, K // 2]
+  a_offsets[expert_id] = a_base_as_int + expert_offset * half_k;
+
+  b_offsets[expert_id] = b_base_as_int + expert_id * n * half_k;
+  // Shape of C = [M, N]
+  out_offsets[expert_id] = out_base_as_int + expert_offset * n;
+  // Shape of a_scale = [sum(sf_sizes), K // group_size]
+  a_scales_offsets[expert_id] = a_scales_base_as_int + sf_offset * group_k;
+
+  assert((reinterpret_cast<uintptr_t>(a_scales_offsets[expert_id]) % 128) == 0 && "TMA requires 128-byte alignment");
+
+  // Shape of B scale = [E, N, K // group_size]
+  b_scales_offsets[expert_id] = b_scales_base_as_int + expert_id * n * group_k;
+  assert((reinterpret_cast<uintptr_t>(b_scales_offsets[expert_id]) % 128) == 0 && "TMA requires 128-byte alignment");
+  // Shape of alpha = [E]
+  alpha_offsets[expert_id] = alphas_base_as_int + expert_id;
+
+  LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id;
+  LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id;
+
+  *layout_sfa_ptr = ScaleConfig::tile_atom_to_shape_SFA(
+      cute::make_shape(static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
+  *layout_sfb_ptr = ScaleConfig::tile_atom_to_shape_SFB(
+      cute::make_shape(static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
+}
+
+#define __CALL_GET_STARTS_KERNEL_BLOCKSCALE(                                                            \
+    ELEMENT_AB_TYPE, SF_TYPE, TYPE_CHECK, C_TYPE, LayoutSFA, LayoutSFB, ScaleConfig)                    \
+  else if (TYPE_CHECK) {                                                                                \
+    __get_group_gemm_starts<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, float, LayoutSFA, LayoutSFB, ScaleConfig> \
+        <<<1, num_experts, 0, stream>>>(                                                                \
+            static_cast<ELEMENT_AB_TYPE**>(a_starts.data_ptr()),                                        \
+            static_cast<ELEMENT_AB_TYPE**>(b_starts.data_ptr()),                                        \
+            static_cast<C_TYPE**>(out_starts.data_ptr()),                                               \
+            static_cast<SF_TYPE**>(a_scales_starts.data_ptr()),                                         \
+            static_cast<SF_TYPE**>(b_scales_starts.data_ptr()),                                         \
+            static_cast<float**>(alpha_starts.data_ptr()),                                              \
+            reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),                                        \
+            reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()),                                        \
+            static_cast<ELEMENT_AB_TYPE*>(a_tensors.data_ptr()),                                        \
+            static_cast<ELEMENT_AB_TYPE*>(b_tensors.data_ptr()),                                        \
+            static_cast<C_TYPE*>(out_tensors.data_ptr()),                                               \
+            static_cast<SF_TYPE*>(a_scales.data_ptr()),                                                 \
+            static_cast<SF_TYPE*>(b_scales.data_ptr()),                                                 \
+            static_cast<float*>(alphas.data_ptr()),                                                     \
+            static_cast<int32_t*>(expert_offsets.data_ptr()),                                           \
+            static_cast<int32_t*>(sf_offsets.data_ptr()),                                               \
+            static_cast<int32_t*>(problem_sizes.data_ptr()),                                            \
+            K,                                                                                          \
+            N);                                                                                         \
+  }
+
+template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
+void run_get_group_gemm_starts(
+    const tvm::ffi::TensorView a_starts,
+    const tvm::ffi::TensorView b_starts,
+    const tvm::ffi::TensorView out_starts,
+    const tvm::ffi::TensorView a_scales_starts,
+    const tvm::ffi::TensorView b_scales_starts,
+    const tvm::ffi::TensorView alpha_starts,
+    const tvm::ffi::TensorView layout_sfa,
+    const tvm::ffi::TensorView layout_sfb,
+    /*these are used for their base addresses*/
+    tvm::ffi::TensorView const& a_tensors,
+    tvm::ffi::TensorView const& b_tensors,
+    tvm::ffi::TensorView const& out_tensors,
+    tvm::ffi::TensorView const& a_scales,
+    tvm::ffi::TensorView const& b_scales,
+    tvm::ffi::TensorView const& alphas,
+    tvm::ffi::TensorView const& expert_offsets,
+    tvm::ffi::TensorView const& sf_offsets,
+    tvm::ffi::TensorView const& problem_sizes,
+    int M,
+    int N,
+    int K) {
+  int num_experts = static_cast<int>(expert_offsets.size(0));
+  auto stream = LaunchKernel::resolve_device(a_tensors.device());
+
+  RuntimeCheck(out_tensors.size(1) == N, "Output tensor shape doesn't match expected shape");
+  RuntimeCheck(
+      K / 2 == b_tensors.size(2),
+      "b_tensors(dim = 2) and a_tensors(dim = 1) trailing"
+      " dimension must match");
+  if (false) {
+  }
+  //(ELEMENT_AB_TYPE, BS_TYPE, TENSOR_C_TYPE, C_TYPE, LayoutSFA, LayoutSFB,
+  // ScaleConfig)
+  __CALL_GET_STARTS_KERNEL_BLOCKSCALE(
+      cutlass::float_e2m1_t,
+      cutlass::float_ue4m3_t,
+      host::is_type<bf16_t>(out_tensors.dtype()),
+      cutlass::bfloat16_t,
+      LayoutSFA,
+      LayoutSFB,
+      ScaleConfig)
+  __CALL_GET_STARTS_KERNEL_BLOCKSCALE(
+      cutlass::float_e2m1_t,
+      cutlass::float_ue4m3_t,
+      host::is_type<fp16_t>(out_tensors.dtype()),
+      cutlass::half_t,
+      LayoutSFA,
+      LayoutSFB,
+      ScaleConfig)
+  else {
+    Panic("Invalid output type (must be float16 or bfloat16)");
+  }
+}
+
+void run_fp4_blockwise_scaled_group_mm_sm120(
+    tvm::ffi::TensorView output,
+    const tvm::ffi::TensorView a,
+    const tvm::ffi::TensorView b,
+    const tvm::ffi::TensorView a_blockscale,
+    const tvm::ffi::TensorView b_blockscales,
+    const tvm::ffi::TensorView alphas,
+    const tvm::ffi::TensorView ab_strides,
+    const tvm::ffi::TensorView c_strides,
+    const tvm::ffi::TensorView problem_sizes,
+    const tvm::ffi::TensorView expert_offsets,
+    const tvm::ffi::TensorView sf_offsets,
+    const tvm::ffi::TensorView a_ptrs,
+    const tvm::ffi::TensorView b_ptrs,
+    const tvm::ffi::TensorView out_ptrs,
+    const tvm::ffi::TensorView a_scales_ptrs,
+    const tvm::ffi::TensorView b_scales_ptrs,
+    const tvm::ffi::TensorView alpha_ptrs,
+    const tvm::ffi::TensorView layout_sfa,
+    const tvm::ffi::TensorView layout_sfb,
+    int M,
+    int N,
+    int K) {
+  using ProblemShape = cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
+  using ElementType = cutlass::float_e2m1_t;
+  using ElementSFType = cutlass::float_ue4m3_t;
+  using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+  using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+
+  using ElementC = cutlass::bfloat16_t;
+  using ElementD = cutlass::bfloat16_t;
+  using ElementAccumulator = float;
+  // Layout definitions
+  using LayoutA = cutlass::layout::RowMajor;
+  using LayoutB = cutlass::layout::ColumnMajor;
+  using LayoutC = cutlass::layout::RowMajor;
+  using LayoutD = cutlass::layout::RowMajor;
+
+  // Alignment constraints
+  static constexpr int AlignmentA = 32;
+  static constexpr int AlignmentB = 32;
+  static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
+  static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
+
+  // Architecture definitions
+  using ArchTag = cutlass::arch::Sm120;
+  using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;
+  using StageCountType = cutlass::gemm::collective::StageCountAuto;
+  using ThreadBlockShape = Shape<_128, _128, _128>;
+  // on the tile size
+
+  using ClusterShape = Shape<_1, _1, _1>;
+
+  using FusionOperation =
+      cutlass::epilogue::fusion::LinearCombination<ElementD, ElementAccumulator, ElementC, ElementAccumulator>;
+
+  using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
+      ArchTag,
+      OperatorClass,
+      ThreadBlockShape,
+      ClusterShape,
+      cutlass::epilogue::collective::EpilogueTileAuto,
+      ElementAccumulator,
+      ElementAccumulator,
+      ElementC,
+      LayoutC*,
+      AlignmentC,
+      ElementD,
+      LayoutC*,
+      AlignmentD,
+      cutlass::epilogue::collective::EpilogueScheduleAuto,
+      FusionOperation>::CollectiveOp;
+
+  using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
+      ArchTag,
+      OperatorClass,
+      ElementA,
+      LayoutA*,
+      AlignmentA,
+      ElementB,
+      LayoutB*,
+      AlignmentB,
+      ElementAccumulator,
+      ThreadBlockShape,
+      ClusterShape,
+      cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
+          sizeof(typename CollectiveEpilogue::SharedStorage))>,
+      cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpong>::CollectiveOp;
+
+  using GemmKernel = cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop, CollectiveEpilogue>;
+
+  using Gemm1SM = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+  using Gemm = Gemm1SM;
+  using StrideA = typename Gemm::GemmKernel::InternalStrideA;
+  using StrideB = typename Gemm::GemmKernel::InternalStrideB;
+  using StrideC = typename Gemm::GemmKernel::InternalStrideC;
+  using StrideD = typename Gemm::GemmKernel::InternalStrideD;
+
+  using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA;
+  using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB;
+  using ScaleConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
+
+  using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
+  int num_experts = static_cast<int>(expert_offsets.size(0));
+
+  run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
+      a_ptrs,
+      b_ptrs,
+      out_ptrs,
+      a_scales_ptrs,
+      b_scales_ptrs,
+      alpha_ptrs,
+      layout_sfa,
+      layout_sfb,
+      a,
+      b,
+      output,
+      a_blockscale,
+      b_blockscales,
+      alphas,
+      expert_offsets,
+      sf_offsets,
+      problem_sizes,
+      M,
+      N,
+      K);
+
+  // Create an instance of the GEMM
+  Gemm gemm_op;
+
+  // Initialize problem_sizes_as_shapes correctly
+  UnderlyingProblemShape* problem_sizes_as_shapes = static_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
+
+  // Set the Scheduler info
+  cutlass::KernelHardwareInfo hw_info;
+
+  using RasterOrderOptions = cutlass::gemm::kernel::detail::RasterOrderOptions;
+  typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
+  scheduler.raster_order = RasterOrderOptions::AlongM;
+  hw_info.device_id = a.device().device_id;
+  static std::unordered_map<int, int> cached_sm_counts;
+  if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
+    cached_sm_counts[hw_info.device_id] =
+        cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
+  }
+  hw_info.sm_count = std::min(cached_sm_counts[hw_info.device_id], std::numeric_limits<int>::max());
+
+  // Mainloop Arguments
+  typename GemmKernel::MainloopArguments mainloop_args{
+      static_cast<const ElementType**>(a_ptrs.data_ptr()),
+      static_cast<StrideA*>(ab_strides.data_ptr()),
+      static_cast<const ElementType**>(b_ptrs.data_ptr()),
+      static_cast<StrideB*>(ab_strides.data_ptr()),
+      static_cast<const ElementSFType**>(a_scales_ptrs.data_ptr()),
+      reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
+      static_cast<const ElementSFType**>(b_scales_ptrs.data_ptr()),
+      reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())};
+
+  // Epilogue Arguments
+  typename GemmKernel::EpilogueArguments epilogue_args{
+      {},  // epilogue.thread
+      nullptr,
+      static_cast<StrideC*>(c_strides.data_ptr()),
+      static_cast<ElementD**>(out_ptrs.data_ptr()),
+      static_cast<StrideC*>(c_strides.data_ptr())};
+  auto& fusion_args = epilogue_args.thread;
+  fusion_args.alpha_ptr_array = reinterpret_cast<float**>(alpha_ptrs.data_ptr());
+  fusion_args.dAlpha = {_0{}, _0{}, 1};
+  fusion_args.beta = 0.0f;
+
+  // Gemm Arguments
+  typename GemmKernel::Arguments args{
+      cutlass::gemm::GemmUniversalMode::kGrouped,
+      {num_experts, problem_sizes_as_shapes, nullptr},
+      mainloop_args,
+      epilogue_args,
+      hw_info,
+      scheduler};
+
+  size_t workspace_size = Gemm::get_workspace_size(args);
+  const cudaStream_t stream = LaunchKernel::resolve_device(a.device());
+  void* workspace = get_cached_workspace(workspace_size, hw_info.device_id, stream);
+
+  auto can_implement_status = gemm_op.can_implement(args);
+  RuntimeCheck(
+      can_implement_status == cutlass::Status::kSuccess,
+      "Failed to implement GEMM: ",
+      cutlassGetStatusString(can_implement_status));
+
+  // Run the GEMM
+  auto status = gemm_op.initialize(args, workspace);
+  RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to initialize GEMM: ", cutlassGetStatusString(status));
+
+  status = gemm_op.run(args, workspace, stream);
+  RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to run GEMM: ", cutlassGetStatusString(status));
+}
+
+template <typename OutType>
+void run_fp4_blockwise_scaled_group_mm_sm100(
+    tvm::ffi::TensorView output,
+    const tvm::ffi::TensorView a,
+    const tvm::ffi::TensorView b,
+    const tvm::ffi::TensorView a_blockscale,
+    const tvm::ffi::TensorView b_blockscales,
+    const tvm::ffi::TensorView alphas,
+    const tvm::ffi::TensorView ab_strides,
+    const tvm::ffi::TensorView c_strides,
+    const tvm::ffi::TensorView problem_sizes,
+    const tvm::ffi::TensorView expert_offsets,
+    const tvm::ffi::TensorView sf_offsets,
+    const tvm::ffi::TensorView a_ptrs,
+    const tvm::ffi::TensorView b_ptrs,
+    const tvm::ffi::TensorView out_ptrs,
+    const tvm::ffi::TensorView a_scales_ptrs,
+    const tvm::ffi::TensorView b_scales_ptrs,
+    const tvm::ffi::TensorView alpha_ptrs,
+    const tvm::ffi::TensorView layout_sfa,
+    const tvm::ffi::TensorView layout_sfb,
+    int M,
+    int N,
+    int K) {
+  using ProblemShape = cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
+  using ElementType = cutlass::float_e2m1_t;
+  using ElementSFType = cutlass::float_ue4m3_t;
+  using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+  using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
+
+  using ElementC = OutType;
+  using ElementD = ElementC;
+  using ElementAccumulator = float;
+  // Layout definitions
+  using LayoutA = cutlass::layout::RowMajor;
+  using LayoutB = cutlass::layout::ColumnMajor;
+  using LayoutC = cutlass::layout::RowMajor;
+  using LayoutD = LayoutC;
+
+  // Alignment constraints
+  static constexpr int AlignmentA = 32;
+  static constexpr int AlignmentB = 32;
+  static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
+  static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
+
+  // Architecture definitions
+  using ArchTag = cutlass::arch::Sm100;
+  using EpilogueOperatorClass = cutlass::arch::OpClassTensorOp;             // Epilogue Operator class tag
+  using MainloopOperatorClass = cutlass::arch::OpClassBlockScaledTensorOp;  // Mainloop Operator class tag
+  using StageCountType = cutlass::gemm::collective::StageCountAuto;         // Stage count maximized based
+                                                                            // on the tile size
+
+  using ClusterShape = Shape<_1, _1, _1>;
+  struct MMA1SMConfig {
+    using MmaTileShape = Shape<_128, _128, _128>;
+    using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmNvf4Sm100;  // Kernel to launch
+    using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm;           // Epilogue to launch
+  };
+
+  using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
+      ArchTag,
+      EpilogueOperatorClass,
+      typename MMA1SMConfig::MmaTileShape,
+      ClusterShape,
+      Shape<_128, _64>,
+      ElementAccumulator,
+      ElementAccumulator,
+      ElementC,
+      LayoutC*,
+      AlignmentC,
+      ElementD,
+      LayoutC*,
+      AlignmentD,
+      typename MMA1SMConfig::EpilogueSchedule>::CollectiveOp;
+
+  using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
+      ArchTag,
+      MainloopOperatorClass,
+      ElementA,
+      LayoutA*,
+      AlignmentA,
+      ElementB,
+      LayoutB*,
+      AlignmentB,
+      ElementAccumulator,
+      typename MMA1SMConfig::MmaTileShape,
+      ClusterShape,
+      cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
+          sizeof(typename CollectiveEpilogue::SharedStorage))>,
+      typename MMA1SMConfig::KernelSchedule>::CollectiveOp;
+
+  using GemmKernel = cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop, CollectiveEpilogue>;
+
+  using Gemm1SM = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+  using Gemm = Gemm1SM;
+  using StrideA = typename Gemm::GemmKernel::InternalStrideA;
+  using StrideB = typename Gemm::GemmKernel::InternalStrideB;
+  using StrideC = typename Gemm::GemmKernel::InternalStrideC;
+  using StrideD = typename Gemm::GemmKernel::InternalStrideD;
+
+  using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA;
+  using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB;
+  using ScaleConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
+
+  using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
+  int num_experts = static_cast<int>(expert_offsets.size(0));
+
+  run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
+      a_ptrs,
+      b_ptrs,
+      out_ptrs,
+      a_scales_ptrs,
+      b_scales_ptrs,
+      alpha_ptrs,
+      layout_sfa,
+      layout_sfb,
+      a,
+      b,
+      output,
+      a_blockscale,
+      b_blockscales,
+      alphas,
+      expert_offsets,
+      sf_offsets,
+      problem_sizes,
+      M,
+      N,
+      K);
+
+  // Create an instance of the GEMM
+  Gemm gemm_op;
+
+  // Initialize problem_sizes_as_shapes correctly
+  UnderlyingProblemShape* problem_sizes_as_shapes = static_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
+
+  // Set the Scheduler info
+  cutlass::KernelHardwareInfo hw_info;
+  using RasterOrderOptions = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm100GroupParams<
+      typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions;
+  typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
+  scheduler.raster_order = RasterOrderOptions::AlongM;
+  hw_info.device_id = a.device().device_id;
+  static std::unordered_map<int, int> cached_sm_counts;
+  if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
+    cached_sm_counts[hw_info.device_id] =
+        cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
+  }
+  hw_info.sm_count = std::min(cached_sm_counts[hw_info.device_id], std::numeric_limits<int>::max());
+
+  // Mainloop Arguments
+  typename GemmKernel::MainloopArguments mainloop_args{
+      static_cast<const ElementType**>(a_ptrs.data_ptr()),
+      static_cast<StrideA*>(ab_strides.data_ptr()),
+      static_cast<const ElementType**>(b_ptrs.data_ptr()),
+      static_cast<StrideB*>(ab_strides.data_ptr()),
+      static_cast<const ElementSFType**>(a_scales_ptrs.data_ptr()),
+      reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
+      static_cast<const ElementSFType**>(b_scales_ptrs.data_ptr()),
+      reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())};
+
+  // Epilogue Arguments
+  typename GemmKernel::EpilogueArguments epilogue_args{
+      {},  // epilogue.thread
+      nullptr,
+      static_cast<StrideC*>(c_strides.data_ptr()),
+      static_cast<ElementD**>(out_ptrs.data_ptr()),
+      static_cast<StrideC*>(c_strides.data_ptr())};
+  auto& fusion_args = epilogue_args.thread;
+  fusion_args.alpha_ptr_array = reinterpret_cast<float**>(alpha_ptrs.data_ptr());
+  fusion_args.dAlpha = {_0{}, _0{}, 1};
+
+  // Gemm Arguments
+  typename GemmKernel::Arguments args{
+      cutlass::gemm::GemmUniversalMode::kGrouped,
+      {num_experts, problem_sizes_as_shapes, nullptr},
+      mainloop_args,
+      epilogue_args,
+      hw_info,
+      scheduler};
+
+  size_t workspace_size = Gemm::get_workspace_size(args);
+  const cudaStream_t stream = LaunchKernel::resolve_device(a.device());
+  void* workspace = get_cached_workspace(workspace_size, hw_info.device_id, stream);
+
+  auto can_implement_status = gemm_op.can_implement(args);
+  RuntimeCheck(
+      can_implement_status == cutlass::Status::kSuccess,
+      "Failed to implement GEMM: ",
+      cutlassGetStatusString(can_implement_status));
+
+  // Run the GEMM
+  auto status = gemm_op.initialize(args, workspace);
+  RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to initialize GEMM: ", cutlassGetStatusString(status));
+
+  status = gemm_op.run(args, workspace, stream);
+  RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to run GEMM: ", cutlassGetStatusString(status));
+}
+
+void cutlass_fp4_group_mm_sm100a_sm120a(
+    tvm::ffi::TensorView output,
+    const tvm::ffi::TensorView a,
+    const tvm::ffi::TensorView b,
+    const tvm::ffi::TensorView a_blockscale,
+    const tvm::ffi::TensorView b_blockscales,
+    const tvm::ffi::TensorView alphas,
+    const tvm::ffi::TensorView ab_strides,
+    const tvm::ffi::TensorView c_strides,
+    const tvm::ffi::TensorView problem_sizes,
+    const tvm::ffi::TensorView expert_offsets,
+    const tvm::ffi::TensorView sf_offsets,
+    const tvm::ffi::TensorView a_ptrs,
+    const tvm::ffi::TensorView b_ptrs,
+    const tvm::ffi::TensorView out_ptrs,
+    const tvm::ffi::TensorView a_scales_ptrs,
+    const tvm::ffi::TensorView b_scales_ptrs,
+    const tvm::ffi::TensorView alpha_ptrs,
+    const tvm::ffi::TensorView layout_sfa,
+    const tvm::ffi::TensorView layout_sfb) {
+  auto check_cuda_contig = [](const tvm::ffi::TensorView t, const char* name) {
+    RuntimeCheck(t.device().device_type == kDLCUDA, name, " must be a CUDA tensor");
+    RuntimeCheck(t.is_contiguous(), name, " must be contiguous");
+  };
+
+  check_cuda_contig(output, "output");
+  check_cuda_contig(a, "a");
+  check_cuda_contig(b, "b");
+  check_cuda_contig(a_blockscale, "a_blockscale");
+  check_cuda_contig(b_blockscales, "b_blockscales");
+  check_cuda_contig(alphas, "alphas");
+  check_cuda_contig(ab_strides, "ab_strides");
+  check_cuda_contig(c_strides, "c_strides");
+  check_cuda_contig(problem_sizes, "problem_sizes");
+  check_cuda_contig(expert_offsets, "expert_offsets");
+  check_cuda_contig(sf_offsets, "sf_offsets");
+  check_cuda_contig(a_ptrs, "a_ptrs");
+  check_cuda_contig(b_ptrs, "b_ptrs");
+  check_cuda_contig(out_ptrs, "out_ptrs");
+  check_cuda_contig(a_scales_ptrs, "a_scales_ptrs");
+  check_cuda_contig(b_scales_ptrs, "b_scales_ptrs");
+  check_cuda_contig(alpha_ptrs, "alpha_ptrs");
+  check_cuda_contig(layout_sfa, "layout_sfa");
+  check_cuda_contig(layout_sfb, "layout_sfb");
+
+  RuntimeCheck(
+      output.device() == a.device() && a.device() == b.device() && a.device() == a_blockscale.device() &&
+          a.device() == b_blockscales.device() && a.device() == alphas.device() && a.device() == ab_strides.device() &&
+          a.device() == c_strides.device() && a.device() == problem_sizes.device() &&
+          a.device() == expert_offsets.device() && a.device() == sf_offsets.device() && a.device() == a_ptrs.device() &&
+          a.device() == b_ptrs.device() && a.device() == out_ptrs.device() && a.device() == a_scales_ptrs.device() &&
+          a.device() == b_scales_ptrs.device() && a.device() == alpha_ptrs.device() &&
+          a.device() == layout_sfa.device() && a.device() == layout_sfb.device(),
+      "all tensors must be on the same device");
+
+  RuntimeCheck(host::is_type<uint8_t>(a.dtype()), "a must be uint8");
+  RuntimeCheck(host::is_type<uint8_t>(b.dtype()), "b must be uint8");
+  RuntimeCheck(host::is_type<fp8_e4m3_t>(a_blockscale.dtype()), "a_blockscale must be float8_e4m3fn");
+  RuntimeCheck(host::is_type<fp8_e4m3_t>(b_blockscales.dtype()), "b_blockscales must be float8_e4m3fn");
+  RuntimeCheck(host::is_type<float>(alphas.dtype()), "alphas must be float32");
+  RuntimeCheck(host::is_type<int64_t>(ab_strides.dtype()), "ab_strides must be int64");
+  RuntimeCheck(host::is_type<int64_t>(c_strides.dtype()), "c_strides must be int64");
+  RuntimeCheck(host::is_type<int32_t>(problem_sizes.dtype()), "problem_sizes must be int32");
+  RuntimeCheck(host::is_type<int32_t>(expert_offsets.dtype()), "expert_offsets must be int32");
+  RuntimeCheck(host::is_type<int32_t>(sf_offsets.dtype()), "sf_offsets must be int32");
+  RuntimeCheck(host::is_type<int64_t>(a_ptrs.dtype()), "a_ptrs must be int64");
+  RuntimeCheck(host::is_type<int64_t>(b_ptrs.dtype()), "b_ptrs must be int64");
+  RuntimeCheck(host::is_type<int64_t>(out_ptrs.dtype()), "out_ptrs must be int64");
+  RuntimeCheck(host::is_type<int64_t>(a_scales_ptrs.dtype()), "a_scales_ptrs must be int64");
+  RuntimeCheck(host::is_type<int64_t>(b_scales_ptrs.dtype()), "b_scales_ptrs must be int64");
+  RuntimeCheck(host::is_type<int64_t>(alpha_ptrs.dtype()), "alpha_ptrs must be int64");
+  RuntimeCheck(host::is_type<int64_t>(layout_sfa.dtype()), "layout_sfa must be int64");
+  RuntimeCheck(host::is_type<int64_t>(layout_sfb.dtype()), "layout_sfb must be int64");
+  RuntimeCheck(
+      host::is_type<bf16_t>(output.dtype()) || host::is_type<fp16_t>(output.dtype()),
+      "output must be bfloat16 or float16");
+
+  RuntimeCheck(a.dim() == 2, "a must be 2D");
+  RuntimeCheck(b.dim() == 3, "b must be 3D");
+  RuntimeCheck(a_blockscale.dim() == 2, "a_blockscale must be 2D");
+  RuntimeCheck(b_blockscales.dim() == 3, "b_blockscales must be 3D");
+  RuntimeCheck(alphas.dim() == 1, "alphas must be 1D");
+  RuntimeCheck(ab_strides.dim() == 1, "ab_strides must be 1D");
+  RuntimeCheck(c_strides.dim() == 1, "c_strides must be 1D");
+  RuntimeCheck(problem_sizes.dim() == 2, "problem_sizes must be 2D");
+  RuntimeCheck(expert_offsets.dim() == 1, "expert_offsets must be 1D");
+  RuntimeCheck(sf_offsets.dim() == 1, "sf_offsets must be 1D");
+  RuntimeCheck(a_ptrs.dim() == 1, "a_ptrs must be 1D");
+  RuntimeCheck(b_ptrs.dim() == 1, "b_ptrs must be 1D");
+  RuntimeCheck(out_ptrs.dim() == 1, "out_ptrs must be 1D");
+  RuntimeCheck(a_scales_ptrs.dim() == 1, "a_scales_ptrs must be 1D");
+  RuntimeCheck(b_scales_ptrs.dim() == 1, "b_scales_ptrs must be 1D");
+  RuntimeCheck(alpha_ptrs.dim() == 1, "alpha_ptrs must be 1D");
+  RuntimeCheck(layout_sfa.dim() == 2, "layout_sfa must be 2D");
+  RuntimeCheck(layout_sfb.dim() == 2, "layout_sfb must be 2D");
+  RuntimeCheck(problem_sizes.size(1) == 3, "problem_sizes must have shape (num_experts, 3)");
+
+  const int num_experts = static_cast<int>(expert_offsets.size(0));
+  RuntimeCheck(problem_sizes.size(0) == num_experts, "problem_sizes size mismatch with expert_offsets");
+  RuntimeCheck(sf_offsets.size(0) == num_experts, "sf_offsets size mismatch with expert_offsets");
+  RuntimeCheck(alphas.size(0) == num_experts, "alphas size mismatch with expert_offsets");
+  RuntimeCheck(ab_strides.size(0) == num_experts, "ab_strides size mismatch with expert_offsets");
+  RuntimeCheck(c_strides.size(0) == num_experts, "c_strides size mismatch with expert_offsets");
+  RuntimeCheck(a_ptrs.size(0) == num_experts, "a_ptrs size mismatch with expert_offsets");
+  RuntimeCheck(b_ptrs.size(0) == num_experts, "b_ptrs size mismatch with expert_offsets");
+  RuntimeCheck(out_ptrs.size(0) == num_experts, "out_ptrs size mismatch with expert_offsets");
+  RuntimeCheck(a_scales_ptrs.size(0) == num_experts, "a_scales_ptrs size mismatch with expert_offsets");
+  RuntimeCheck(b_scales_ptrs.size(0) == num_experts, "b_scales_ptrs size mismatch with expert_offsets");
+  RuntimeCheck(alpha_ptrs.size(0) == num_experts, "alpha_ptrs size mismatch with expert_offsets");
+  RuntimeCheck(layout_sfa.size(0) == num_experts && layout_sfa.size(1) == 5, "layout_sfa must be [num_experts, 5]");
+  RuntimeCheck(layout_sfb.size(0) == num_experts && layout_sfb.size(1) == 5, "layout_sfb must be [num_experts, 5]");
+
+  int M = static_cast<int>(a.size(0));
+  int N = static_cast<int>(b.size(1));
+  int K = static_cast<int>(2 * b.size(2));
+  RuntimeCheck(output.dim() == 2, "output must be 2D");
+  RuntimeCheck(output.size(0) == M && output.size(1) == N, "output shape mismatch");
+
+  auto sm_version = getSMVersion(a.device().device_id);
+  if (sm_version == 100 || sm_version == 103) {
+    if (host::is_type<bf16_t>(output.dtype())) {
+      run_fp4_blockwise_scaled_group_mm_sm100<cutlass::bfloat16_t>(
+          output,
+          a,
+          b,
+          a_blockscale,
+          b_blockscales,
+          alphas,
+          ab_strides,
+          c_strides,
+          problem_sizes,
+          expert_offsets,
+          sf_offsets,
+          a_ptrs,
+          b_ptrs,
+          out_ptrs,
+          a_scales_ptrs,
+          b_scales_ptrs,
+          alpha_ptrs,
+          layout_sfa,
+          layout_sfb,
+          M,
+          N,
+          K);
+    } else {
+      run_fp4_blockwise_scaled_group_mm_sm100<cutlass::half_t>(
+          output,
+          a,
+          b,
+          a_blockscale,
+          b_blockscales,
+          alphas,
+          ab_strides,
+          c_strides,
+          problem_sizes,
+          expert_offsets,
+          sf_offsets,
+          a_ptrs,
+          b_ptrs,
+          out_ptrs,
+          a_scales_ptrs,
+          b_scales_ptrs,
+          alpha_ptrs,
+          layout_sfa,
+          layout_sfb,
+          M,
+          N,
+          K);
+    }
+  } else if (sm_version >= 120) {
+    if (host::is_type<bf16_t>(output.dtype())) {
+      run_fp4_blockwise_scaled_group_mm_sm120(
+          output,
+          a,
+          b,
+          a_blockscale,
+          b_blockscales,
+          alphas,
+          ab_strides,
+          c_strides,
+          problem_sizes,
+          expert_offsets,
+          sf_offsets,
+          a_ptrs,
+          b_ptrs,
+          out_ptrs,
+          a_scales_ptrs,
+          b_scales_ptrs,
+          alpha_ptrs,
+          layout_sfa,
+          layout_sfb,
+          M,
+          N,
+          K);
+    } else {
+      Panic("SM120 path currently supports only bfloat16 output");
+    }
+  } else {
+    RuntimeCheck(false, "Unsupported SM version: ", sm_version);
+  }
+}
+
+void cutlass_fp4_group_mm(
+    tvm::ffi::TensorView output,
+    const tvm::ffi::TensorView a,
+    const tvm::ffi::TensorView b,
+    const tvm::ffi::TensorView a_blockscale,
+    const tvm::ffi::TensorView b_blockscales,
+    const tvm::ffi::TensorView alphas,
+    const tvm::ffi::TensorView ab_strides,
+    const tvm::ffi::TensorView c_strides,
+    const tvm::ffi::TensorView problem_sizes,
+    const tvm::ffi::TensorView expert_offsets,
+    const tvm::ffi::TensorView sf_offsets,
+    const tvm::ffi::TensorView a_ptrs,
+    const tvm::ffi::TensorView b_ptrs,
+    const tvm::ffi::TensorView out_ptrs,
+    const tvm::ffi::TensorView a_scales_ptrs,
+    const tvm::ffi::TensorView b_scales_ptrs,
+    const tvm::ffi::TensorView alpha_ptrs,
+    const tvm::ffi::TensorView layout_sfa,
+    const tvm::ffi::TensorView layout_sfb) {
+  cutlass_fp4_group_mm_sm100a_sm120a(
+      output,
+      a,
+      b,
+      a_blockscale,
+      b_blockscales,
+      alphas,
+      ab_strides,
+      c_strides,
+      problem_sizes,
+      expert_offsets,
+      sf_offsets,
+      a_ptrs,
+      b_ptrs,
+      out_ptrs,
+      a_scales_ptrs,
+      b_scales_ptrs,
+      alpha_ptrs,
+      layout_sfa,
+      layout_sfb);
+}

python/sglang/jit_kernel/csrc/moe/tvm_ffi_utils.h

@@ -0,0 +1,105 @@
+/*
+ * Copyright (c) 2023 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/dtype.h>
+#include <tvm/ffi/error.h>
+#include <tvm/ffi/extra/c_env_api.h>
+#include <tvm/ffi/function.h>
+
+#include "dlpack/dlpack.h"
+
+using tvm::ffi::Tensor;
+using tvm::ffi::TensorView;
+namespace ffi = tvm::ffi;
+
+inline constexpr int64_t encode_dlpack_dtype(DLDataType dtype) {
+  return (dtype.code << 16) | (dtype.bits << 8) | dtype.lanes;
+}
+
+constexpr DLDataType dl_uint8 = DLDataType{kDLUInt, 8, 1};
+constexpr DLDataType dl_uint16 = DLDataType{kDLUInt, 16, 1};
+constexpr DLDataType dl_uint32 = DLDataType{kDLUInt, 32, 1};
+constexpr DLDataType dl_uint64 = DLDataType{kDLUInt, 64, 1};
+constexpr DLDataType dl_int8 = DLDataType{kDLInt, 8, 1};
+constexpr DLDataType dl_int16 = DLDataType{kDLInt, 16, 1};
+constexpr DLDataType dl_int32 = DLDataType{kDLInt, 32, 1};
+constexpr DLDataType dl_int64 = DLDataType{kDLInt, 64, 1};
+constexpr DLDataType dl_float16 = DLDataType{kDLFloat, 16, 1};
+constexpr DLDataType dl_float32 = DLDataType{kDLFloat, 32, 1};
+constexpr DLDataType dl_float64 = DLDataType{kDLFloat, 64, 1};
+constexpr DLDataType dl_float8_e4m3fn = DLDataType{kDLFloat8_e4m3fn, 8, 1};
+constexpr DLDataType dl_float8_e5m2 = DLDataType{kDLFloat8_e5m2, 8, 1};
+constexpr DLDataType dl_float4_e2m1fn = DLDataType{kDLFloat4_e2m1fn, 4, 1};
+constexpr DLDataType dl_float4_e2m1fn_x2 = DLDataType{kDLFloat4_e2m1fn, 4, 2};
+constexpr DLDataType dl_bfloat16 = DLDataType{kDLBfloat, 16, 1};
+constexpr DLDataType dl_bool = DLDataType{kDLBool, 8, 1};
+
+constexpr int64_t float16_code = encode_dlpack_dtype(dl_float16);
+constexpr int64_t bfloat16_code = encode_dlpack_dtype(dl_bfloat16);
+constexpr int64_t float32_code = encode_dlpack_dtype(dl_float32);
+constexpr int64_t uint8_code = encode_dlpack_dtype(dl_uint8);
+constexpr int64_t int32_code = encode_dlpack_dtype(dl_int32);
+constexpr int64_t int64_code = encode_dlpack_dtype(dl_int64);
+constexpr int64_t float8_e4m3fn_code = encode_dlpack_dtype(dl_float8_e4m3fn);
+constexpr int64_t float8_e5m2_code = encode_dlpack_dtype(dl_float8_e5m2);
+constexpr int64_t float4_e2m1fn_code = encode_dlpack_dtype(dl_float4_e2m1fn);
+
+constexpr DLDevice cpu = DLDevice{kDLCPU, 0};
+
+#define CHECK_CUDA(x) TVM_FFI_ICHECK_EQ(x.device().device_type, kDLCUDA) << #x " must be a CUDA tensor";
+#define CHECK_CPU(x) TVM_FFI_ICHECK_EQ(x.device().device_type, kDLCPU) << #x " must be a host tensor";
+#define CHECK_CONTIGUOUS(x) TVM_FFI_ICHECK(x.IsContiguous()) << #x " must be contiguous";
+#define CHECK_LAST_DIM_CONTIGUOUS(x) \
+  TVM_FFI_ICHECK_EQ(x.stride(-1), 1) \
+  #x "must be contiguous at last dimension";
+#define CHECK_INPUT(x) \
+  CHECK_CUDA(x);       \
+  CHECK_CONTIGUOUS(x)
+#define CHECK_INPUT_TYPE(x, st) TVM_FFI_ICHECK_EQ(x.dtype(), st) << "Inconsistency of Tensor type: " #x;
+#define CHECK_INPUT_AND_TYPE(x, st) \
+  CHECK_CUDA(x);                    \
+  CHECK_CONTIGUOUS(x);              \
+  CHECK_INPUT_TYPE(x, st)
+#define CHECK_LAST_DIM_CONTIGUOUS_INPUT(x) \
+  CHECK_CUDA(x);                           \
+  CHECK_LAST_DIM_CONTIGUOUS(x)
+#define CHECK_DIM(d, x) TVM_FFI_ICHECK_EQ(x.ndim(), d) << #x " must be a " #d "D tensor";
+#define CHECK_DEVICE(a, b)                                           \
+  TVM_FFI_ICHECK_EQ(a.device().device_type, b.device().device_type); \
+  TVM_FFI_ICHECK_EQ(a.device().device_id, b.device().device_id);
+
+inline cudaStream_t get_current_stream() {
+  int device;
+  cudaGetDevice(&device);
+  return static_cast<cudaStream_t>(TVMFFIEnvGetStream(kDLCUDA, device));
+}
+
+inline cudaStream_t get_stream(DLDevice device) {
+  return static_cast<cudaStream_t>(TVMFFIEnvGetStream(device.device_type, device.device_id));
+}
+
+inline int64_t get_element_size(ffi::Tensor x) {
+  return (x.dtype().bits * x.dtype().lanes) / 8;
+}
+
+inline int64_t get_element_size(ffi::TensorView x) {
+  return (x.dtype().bits * x.dtype().lanes) / 8;
+}
+
+inline ffi::Tensor alloc_tensor(tvm::ffi::Shape shape, DLDataType dtype, DLDevice device) {
+  return ffi::Tensor::FromEnvAlloc(TVMFFIEnvTensorAlloc, shape, dtype, device);
+}