chore: import upstream snapshot with attribution

2026-07-13 12:38:16 +08:00
commit 94057c3d3e
7152 changed files with 2120455 additions and 0 deletions
@@ -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);
+  }
+};
@@ -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");
+    }
+  }
+};
@@ -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
@@ -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