[CK_TILE] fused-moe first version (#1634)

* moe pipeline

* update code

* compile OK

* update

* update cpu reference

* update pipeline_gemm0

* compiler ok

* update pipeline

* rename to ex pipeline

* block-asm

* update

* update

* update first gemm ok

* compute correct

* update file structure

* update README

* update

* update

* update code

* update API

* return unsupport case

* add comment

* update readme

* update

* uncomment

* update

* fix build err

---------
Co-authored-by: valarLip <340077269@qq.com>

example/ck_tile/06_permute/alternative_impl/matrix_core_swizzle.cpp

@@ -40,7 +40,7 @@ float matrix_core_swizzle(matrix_core_swizzle_traits t,
             else if(t.permute.compare("0,1,3,4,2,5") == 0)
             {
                 constexpr matrix_core_permute_style pstyle =
-                    matrix_core_permute_style::permute_b_nr_kr_kw_nw_kv;
+                    matrix_core_permute_style::b_nr_kr_kw_nw_kv;
                 using Kernel =
                     matrix_core_swizzle_kernel<BLOCK_SIZE, NPerBlock, KPerBlock, pstyle, Inst>;
 
@@ -83,7 +83,7 @@ float matrix_core_swizzle(matrix_core_swizzle_traits t,
             else if(t.permute.compare("0,1,3,4,2,5") == 0)
             {
                 constexpr matrix_core_permute_style pstyle =
-                    matrix_core_permute_style::permute_b_nr_kr_kw_nw_kv;
+                    matrix_core_permute_style::b_nr_kr_kw_nw_kv;
                 using Kernel =
                     matrix_core_swizzle_kernel<BLOCK_SIZE, NPerBlock, KPerBlock, pstyle, Inst>;
 

example/ck_tile/06_permute/alternative_impl/matrix_core_swizzle_kernel.hpp

@@ -42,8 +42,8 @@ enum class matrix_core_permute_style
 {
     permute_b_n0_k0_n1_k1_n2_k2 = 0, // 0,1,4,2,5,3,6
     permute_b_n0_n1_k0_k1_n2_k2 = 1, // 0,1,2,4,5,3,6
-    permute_b_nr_kr_kw_nw_kv    = 2, // 0,1,3,4,2,5
-    permute_b_nr_kr_waveflatten = permute_b_nr_kr_kw_nw_kv,
+    b_nr_kr_kw_nw_kv            = 2, // 0,1,3,4,2,5
+    b_nr_kr_waveflatten         = b_nr_kr_kw_nw_kv,
 };
 
 // assume this is B matrix, originally we have batch*n*k
@@ -203,7 +203,7 @@ struct matrix_core_swizzle_kernel
             else
             {
                 // clang-format off
-                // permute_b_nr_kr_kw_nw_kv or permute_b_nr_kr_waveflatten
+                // b_nr_kr_kw_nw_kv or b_nr_kr_waveflatten
                 constexpr index_t Kv = Alignment;
                 constexpr index_t Nw = WarpGemm::WarpGemmAttribute::Impl::kAMLane;
                 constexpr index_t Kw = WarpGemm::WarpGemmAttribute::Impl::kABKLane;
@@ -332,7 +332,7 @@ struct matrix_core_swizzle_kernel
                         make_tuple(sequence<0>{}, sequence<1>{}));
                     return tmp_1;
 #else
-                    // permute_b_nr_kr_waveflatten = permute_b_nr_kr_kw_nw_kv,
+                    // b_nr_kr_waveflatten = b_nr_kr_kw_nw_kv,
                     constexpr index_t kv = Alignment;
                     constexpr index_t nw = WarpGemm::WarpGemmAttribute::Impl::kAMLane;
                     constexpr index_t kw = WarpGemm::WarpGemmAttribute::Impl::kABKLane;
@@ -376,13 +376,13 @@ struct matrix_core_swizzle_kernel
                 else
                 {
 #if MERGE_2D_013425
-                    // permute_b_nr_kr_waveflatten = permute_b_nr_kr_kw_nw_kv
+                    // b_nr_kr_waveflatten = b_nr_kr_kw_nw_kv
                     return make_tile_window(dst_view,
                                             make_tuple(number<NPerBlock>{}, number<KPerBlock>{}),
                                             {i_n * NPerBlock, i_k * KPerBlock},
                                             get_dst_dist());
 #else
-                    // permute_b_nr_kr_waveflatten = permute_b_nr_kr_kw_nw_kv
+                    // b_nr_kr_waveflatten = b_nr_kr_kw_nw_kv
                     constexpr index_t kv = Alignment;
                     constexpr index_t nw = WarpGemm::WarpGemmAttribute::Impl::kAMLane;
                     constexpr index_t kw = WarpGemm::WarpGemmAttribute::Impl::kABKLane;

example/ck_tile/06_permute/permute.cpp

@@ -264,7 +264,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
     {
         if(arg_parser.get_str("perm") == std::string("0,1,3,4,2,5"))
         {
-            // permute_b_nr_kr_kw_nw_kv = 2,   // 0,1,3,4,2,5
+            // b_nr_kr_kw_nw_kv = 2,   // 0,1,3,4,2,5
             matrix_core_swizzle_traits t;
             t.data_type = data_type;
             t.permute   = arg_parser.get_str("perm");

example/ck_tile/13_moe_sorting/moe_sorting_api.hpp

@@ -5,7 +5,7 @@
 #include <string>
 #include "ck_tile/core.hpp"
 #include "ck_tile/host.hpp"
-#include "ck_tile/ops/moe_sorting.hpp"
+#include "ck_tile/ops/fused_moe.hpp"
 
 struct moe_sorting_trait
 {

example/ck_tile/15_fused_moe/CMakeLists.txt

+set(TILE_EXAPMLE_FUSED_MOE "tile_example_fused_moe")
+# not using add_example_executable() to add this target, since we don't want this to have
+# to be included in "make all/install/check"
+message("adding ${TILE_EXAPMLE_FUSED_MOE}")
+file(GLOB INSTANCE_SRCS instances/*.cpp)
+add_executable(${TILE_EXAPMLE_FUSED_MOE} EXCLUDE_FROM_ALL main.cpp)
+target_include_directories(${TILE_EXAPMLE_FUSED_MOE} PRIVATE ${CMAKE_CURRENT_LIST_DIR})
+target_sources(${TILE_EXAPMLE_FUSED_MOE} PRIVATE ${INSTANCE_SRCS})
+
+set(TILE_EXAPMLE_FUSED_MOE_COMPILE_OPTIONS)
+
+# NOTE: we turn off undefined-func-template to let source compile without explicit declare function specializations
+list(APPEND TILE_EXAPMLE_FUSED_MOE_COMPILE_OPTIONS -Wno-undefined-func-template -Wno-float-equal)
+list(APPEND TILE_EXAPMLE_FUSED_MOE_COMPILE_OPTIONS -DCK_TILE_BUFFER_LOAD_AGPR=1) # TODO: enable load to a
+list(APPEND TILE_EXAPMLE_FUSED_MOE_COMPILE_OPTIONS -DCK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=4) # rta
+# list(APPEND TILE_EXAPMLE_FUSED_MOE_COMPILE_OPTIONS  -mllvm -greedy-reverse-local-assignment=1)
+# list(APPEND TILE_EXAPMLE_FUSED_MOE_COMPILE_OPTIONS -v --save-temps -Wno-gnu-line-marker)
+
+target_compile_options(${TILE_EXAPMLE_FUSED_MOE} PRIVATE ${TILE_EXAPMLE_FUSED_MOE_COMPILE_OPTIONS})

example/ck_tile/15_fused_moe/README.md

+# fused-moe
+Implementing the fused-moe block operator using ck-tile. This is a scatter/gather-group-gemm based solution, similiar to that of [vllm moe](https://github.com/vllm-project/vllm/blob/main/benchmarks/kernels/benchmark_moe.py), but we introduce more kernel fusion to boost performance
+![](misc/moe-0.png)
+
+The benifit of this fused-moe:
+* 1.5~2x perf boost compared with current vllm solution
+* zero workspace to reduce memory footprint
+* much less kernel instance, easy to maintain
+
+# Implementation and feature support
+## moe-sorting
+this is a common pre-process step before the actual moe-gemm. The purpose is to transform the moe loop over from token-by-token to expert-by-expert, make sure very workgroup is working for a single expert (B matrix). Besides, we extend this op to do the zeroing of the output buffer(to be used for reduce buffer with atomic)
+
+## moe-gemm
+`moe-gemm` is a group-gemm based back-to-back gemm, where the row-id of input token comes from another buffer. Naive understanding of fused-moe is from token-by-token view as below picture:
+![](misc/moe-1.png)
+After `moe-sorting`, we can view this algorithm as expert-by-expert, as below:
+![](misc/moe-2.png)
+
+## optimization
+summary of the key design of this fused-moe operator:
+* fuse 2 group-gemm + activation + `topk-weight` multiply into single kernel, using atomic for 2nd gemm accumualation
+* fuse buffer-zeroing in `moe-sorgin`, user no longer need call extra torch.zero() for the out buffer
+* fused scatter-gather for row index(same as vllm)
+* pre-shuffle B matric(weight) to maximize memory throughput. input(activation) keep original layout `[batch, hidden]`.
+* extrem optimized pipeline using block-inline-asm(we call it `micro-kernel` or `uk`), while not breaking the *composable* design of ck
+
+## 
+```
+// [indexing implementation-1]
+// using M_a as constexpr block_size to partition all tokens into different slices
+// each slice map to one expert, and one expert can have multiple slices
+// e.g. num_experts = 6, topk=3, M_a = 4, input_tokens = 5
+// before sort, topk_ids is : [[0, 3, 5], [2, 3, 5], [1, 3, 5], [1, 2, 3], [1, 3, 5]]
+//                            tok-0      tok-1      tok-2      tok-3      tok-4
+//           topk_weight is : [[a, b, c], [d, e, f], [g, h, i], [j, k, l], [m, n, o]] (some float number)
+//
+// token_id_per_expert is : [[0], [2, 3, 4], [1, 3], [0, 1, 2, 3, 4], [], [0, 1, 2, 5]]
+//  (only for reference)    exp-0  exp-1     exp-2   exp-3          exp-4  exp-5
+// weight_id_per_expert is: [[a], [g, j, m], [d, k], [b, e, h, l, n], [], [c, f, i, o]]
+//
+// max_num_tokens_padded : topk * input_tokens + num_experts * (M_a - 1)
+// * this could be larger than actual, since actual tokens are on GPU
+//
+// sorted_token_ids_ptr   : [0, 6, 6, 6, 2, 3, 4, 6, 1, 3, 6, 6, 0, 1, 2, 3, 4, 6, 6, 6, 6, 6, 6, 6, 0, 1, 2, 5]
+//                          |-  exp-0  -|-  exp-1  -|-  exp-2  -|-      exp-3          -|-  exp-4 -|-  exp-5  -|
+// sorted_weight_ptr      : [a, *, *, *, g, j, m, *, d, k, *, *, b, e, h, l, n, *, *, *, *, *, *, *, c, f, i, o]
+//
+// * length is max_num_tokens_padded, actual size is num_tokens_post_padded_ptr
+//
+// sorted_expert_ids_ptr  : [0, 1, 2, 3, 3, 4, 5]
+// * length is (max_num_tokens_padded + block_size - 1) / block_size
+//
+// num_tokens_post_padded_ptr : [28]
+// num_sorted_tiles_ptr : [7]
+//
+// * different from vLLM
+//   1) token_id stored in sorted_token_ids_ptr is actual token_id, not token_id*top_K expanded id
+//   2）need sorted_weight_ptr
+//   3) use num_sorted_tiles_ptr, already divided by M_a
+//
+// * below used for indexing
+//  1) sorted_token_ids_ptr [max_num_tokens_padded]
+//  2) sorted_weight_ptr
+//  3) sorted_expert_ids_ptr
+//  4）num_tokens_post_padded_ptr/num_sorted_tiles_ptr (select one)
+//
+//   max_num_tokens_padded: opk_ids.numel() + num_experts * (block_size - 1)
+```
\ No newline at end of file

example/ck_tile/15_fused_moe/fused_moe.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "fused_moesorting.hpp"
+#include "fused_moegemm.hpp"
+
+struct fused_moe_args
+{
+    const void* a_ptr;              // [m, k], input token
+    const void* a_scale_ptr;        // [m, 1], token scale
+    const void* g_ptr;              // [e, n, k]/[e, 2*n, k], pre-shuffle([e, nr, kr, w])
+    const void* d_ptr;              // [e, n, k], pre-shuffle([e, nr, kr, w])
+    const void* g_scale_ptr;        // [e, 1, n], gate(up) scale
+    const void* d_scale_ptr;        // [e, 1, k], down scale
+    const void* y_smooth_scale_ptr; // [e, 1, n], smooth-quant-scale for 2nd gemm input
+    void* o_ptr;                    // [m, k], output token (no need to do zeroing)
+
+    const void* topk_ids_ptr;    // [tokens, topk]
+    const void* topk_weight_ptr; // [tokens, topk]
+    void* sorted_token_ids_ptr;  // [max_num_tokens_padded]
+    void* sorted_weight_ptr;     // [max_num_tokens_padded]
+    void* sorted_expert_ids_ptr; // [(max_num_tokens_padded + block_size - 1) / block_size]
+    void* num_sorted_tiles_ptr;  // [1]
+
+    ck_tile::index_t block_m;           // block_m, used to devide the input
+    ck_tile::index_t hidden_size;       // k
+    ck_tile::index_t intermediate_size; // n / TP, for Gate. if Gate+Up, Down need divide by 2
+    ck_tile::index_t num_tokens;        // input number of tokens for current iteration
+    ck_tile::index_t num_experts;       // number of groups
+    ck_tile::index_t topk;              // need this?
+
+    ck_tile::index_t stride_token; // for input/output, stride for each row, should >= hidden_size
+};
+
+// This is the public API, will be generated by script
+struct fused_moe_traits
+{
+    std::string prec_i;  // input precision
+    std::string prec_w;  // weight precision
+    std::string prec_o;  // output precision
+    std::string prec_st; // token scale data type
+    std::string prec_sw; // weight scale data type
+    std::string prec_sq; // smooth quant scale
+    std::string prec_kw; // topk-weight data type
+    int block_m;
+    int gate_only;
+    int fused_quant; // 0:no-sweep, 1:smooth-dynamic-quant, 2:dynamic-quant
+};
+
+float fused_moe(fused_moe_traits, fused_moe_args, const ck_tile::stream_config&);

example/ck_tile/15_fused_moe/fused_moegemm.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "ck_tile/core.hpp"
+#include "ck_tile/host/kernel_launch.hpp"
+#include "ck_tile/ops/fused_moe.hpp"
+#include <string>
+
+// this is only a convenient structure for creating an example
+// this is not part of the host API
+template <typename I, typename W, typename O, typename ST, typename SW, typename SQ, typename KW>
+struct FusedMoeGemmTypeConfig;
+
+template <typename ST, typename SW, typename SQ, typename KW>
+struct FusedMoeGemmTypeConfig<ck_tile::bf16_t, ck_tile::bf16_t, ck_tile::bf16_t, ST, SW, SQ, KW>
+{
+    using ADataType            = ck_tile::bf16_t;
+    using GDataType            = ck_tile::bf16_t;
+    using DDataType            = ck_tile::bf16_t;
+    using AccDataType          = float;
+    using ODataType            = ck_tile::bf16_t;
+    using AScaleDataType       = ck_tile::remove_cvref_t<ST>;
+    using GScaleDataType       = ck_tile::remove_cvref_t<SW>;
+    using DScaleDataType       = ck_tile::remove_cvref_t<SW>;
+    using YSmoothScaleDataType = ck_tile::remove_cvref_t<SQ>;
+    using TopkWeightDataType   = ck_tile::remove_cvref_t<KW>;
+    using IndexDataType        = ck_tile::index_t;
+};
+
+template <typename ST, typename SW, typename SQ, typename KW>
+struct FusedMoeGemmTypeConfig<ck_tile::fp16_t, ck_tile::fp16_t, ck_tile::fp16_t, ST, SW, SQ, KW>
+{
+    using ADataType            = ck_tile::fp16_t;
+    using GDataType            = ck_tile::fp16_t;
+    using DDataType            = ck_tile::fp16_t;
+    using AccDataType          = float;
+    using ODataType            = ck_tile::fp16_t;
+    using AScaleDataType       = ck_tile::remove_cvref_t<ST>;
+    using GScaleDataType       = ck_tile::remove_cvref_t<SW>;
+    using DScaleDataType       = ck_tile::remove_cvref_t<SW>;
+    using YSmoothScaleDataType = ck_tile::remove_cvref_t<SQ>;
+    using TopkWeightDataType   = ck_tile::remove_cvref_t<KW>;
+    using IndexDataType        = ck_tile::index_t;
+};
+
+template <typename ST, typename SW, typename SQ, typename KW>
+struct FusedMoeGemmTypeConfig<ck_tile::int8_t, ck_tile::int8_t, ck_tile::bf16_t, ST, SW, SQ, KW>
+{
+    using ADataType            = ck_tile::int8_t;
+    using GDataType            = ck_tile::int8_t;
+    using DDataType            = ck_tile::int8_t;
+    using AccDataType          = int32_t;
+    using ODataType            = ck_tile::bf16_t;
+    using AScaleDataType       = ck_tile::remove_cvref_t<ST>;
+    using GScaleDataType       = ck_tile::remove_cvref_t<SW>;
+    using DScaleDataType       = ck_tile::remove_cvref_t<SW>;
+    using YSmoothScaleDataType = ck_tile::remove_cvref_t<SQ>;
+    using TopkWeightDataType   = ck_tile::remove_cvref_t<KW>;
+    using IndexDataType        = ck_tile::index_t;
+};
+
+// runtime args
+struct fused_moegemm_args : public ck_tile::FusedMoeGemmHostArgs
+{
+};
+
+// This is the public API, will be generated by script
+struct fused_moegemm_traits
+{
+    std::string prec_i;  // input precision
+    std::string prec_w;  // weight precision
+    std::string prec_o;  // output precision
+    std::string prec_st; // token scale data type
+    std::string prec_sw; // weight scale data type
+    std::string prec_sq; // smooth quant scale
+    std::string prec_kw; // topk-weight data type
+    int block_m;
+    int gate_only;
+    int fused_quant; // 0:no-sweep, 1:smooth-dynamic-quant, 2:dynamic-quant
+};
+
+float fused_moegemm(fused_moegemm_traits, fused_moegemm_args, const ck_tile::stream_config&);

example/ck_tile/15_fused_moe/fused_moesorting.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+#include <string>
+#include "ck_tile/core.hpp"
+#include "ck_tile/host.hpp"
+#include "ck_tile/ops/fused_moe.hpp"
+
+struct fused_moesorting_trait
+{
+    std::string index_type;
+    std::string weight_type; // currently always float
+};
+
+struct fused_moesorting_args : public ck_tile::MoeSortingHostArgs
+{
+};
+
+float fused_moesorting(fused_moesorting_trait t, fused_moesorting_args a, ck_tile::stream_config s);

example/ck_tile/15_fused_moe/instances/fused_moe_api.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "fused_moe.hpp"
+
+float fused_moe(fused_moe_traits t, fused_moe_args a, const ck_tile::stream_config& s)
+{
+    auto s_sub = ck_tile::stream_config{s.stream_id_, false, s.log_level_, 0, 1};
+
+    auto o_data_bytes = [&]() {
+        if(t.prec_o == "fp32")
+            return 4;
+        else if(t.prec_o == "fp16" || t.prec_o == "bf16")
+            return 2;
+        else if(t.prec_o == "int8" || t.prec_o == "fp8")
+            return 1;
+        return 1;
+    }();
+
+    auto t0 = fused_moesorting_trait{"int32", "fp32"};
+    auto a0 = fused_moesorting_args{
+        a.topk_ids_ptr,                              // const void* p_topk_ids;
+        a.topk_weight_ptr,                           // const void* p_weights;
+        a.sorted_token_ids_ptr,                      // void* p_sorted_token_ids;
+        a.sorted_weight_ptr,                         // void* p_sorted_weights;
+        a.sorted_expert_ids_ptr,                     // void* p_sorted_expert_ids;
+        a.num_sorted_tiles_ptr,                      // void* p_total_tokens_post_pad;
+        a.o_ptr,                                     // void* p_moe_buf;
+        a.num_tokens,                                // index_t tokens;
+        a.block_m,                                   // index_t unit_size;
+        a.num_experts,                               // index_t num_experts;
+        a.topk,                                      // index_t topk;
+        a.num_tokens * a.stride_token * o_data_bytes // index_t moe_buf_bytes;
+    };
+
+    auto t1 = fused_moegemm_traits{t.prec_i,
+                                   t.prec_w,
+                                   t.prec_o,
+                                   t.prec_st,
+                                   t.prec_sw,
+                                   t.prec_sq,
+                                   t.prec_kw,
+                                   t.block_m,
+                                   t.gate_only,
+                                   t.fused_quant};
+    auto a1 = fused_moegemm_args{
+        a.a_ptr,                 // const void* a_ptr;
+        a.a_scale_ptr,           // const void* a_scale_ptr;
+        a.g_ptr,                 // const void* g_ptr;
+        a.d_ptr,                 // const void* d_ptr;
+        a.g_scale_ptr,           // const void* g_scale_ptr;
+        a.d_scale_ptr,           // const void* d_scale_ptr;
+        a.y_smooth_scale_ptr,    // const void* y_smooth_scale_ptr;
+        a.o_ptr,                 // void* o_ptr;
+        a.sorted_token_ids_ptr,  // const void* sorted_token_ids_ptr;
+        a.sorted_weight_ptr,     // const void* sorted_weight_ptr;
+        a.sorted_expert_ids_ptr, // const void* sorted_expert_ids_ptr;
+        a.num_sorted_tiles_ptr,  // const void* num_sorted_tiles_ptr;
+        a.hidden_size,           // index_t hidden_size;
+        a.intermediate_size,     // index_t intermediate_size;
+        a.num_tokens,            // index_t num_tokens;
+        a.num_experts,           // index_t num_experts;
+        a.topk,                  // index_t topk;
+        a.stride_token           // index_t stride_token;
+    };
+
+    float r0 = -1;
+    float r1 = -1;
+
+    float r = ck_tile::launch_kernel(
+        s,
+        [=, &r0](const ck_tile::stream_config&) { r0 = fused_moesorting(t0, a0, s_sub); },
+        [=, &r1](const ck_tile::stream_config&) { r1 = fused_moegemm(t1, a1, s_sub); });
+
+    // keep unsupported case return negative
+    if(r0 < 0 || r1 < 0)
+        return -1;
+
+    return r;
+}

example/ck_tile/15_fused_moe/instances/fused_moegemm_api.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <ck_tile/core.hpp>
+#include "fused_moegemm.hpp"
+#include "fused_moegemm_api_traits.hpp"
+
+// Note: this internal API only declare, not define here, otherwise will block `make -j`
+template <typename Traits_>
+float fused_moegemm_(const ck_tile::stream_config& s, fused_moegemm_args a);
+
+template <ck_tile::index_t... Is>
+using S = ck_tile::sequence<Is...>;
+
+float fused_moegemm(fused_moegemm_traits t, fused_moegemm_args a, const ck_tile::stream_config& s)
+{
+    // clang-format off
+    float r = -1;
+    if(t.prec_i == "bf16" && t.prec_w == "bf16" && t.prec_o == "bf16" && t.prec_st == "fp32" &&
+       t.prec_sw == "fp32" && t.prec_sq == "fp32" && t.prec_kw == "fp32" && t.block_m == 32 && t.gate_only == 1)
+    {
+        using t_ = fmoe_<ck_tile::bf16_t, ck_tile::bf16_t, ck_tile::bf16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, 1, 0>;
+        r = fused_moegemm_<t_>(s, a);
+    }
+    else if(t.prec_i == "fp16" && t.prec_w == "fp16" && t.prec_o == "fp16" && t.prec_st == "fp32" &&
+       t.prec_sw == "fp32" && t.prec_sq == "fp32" && t.prec_kw == "fp32" && t.block_m == 32 && t.gate_only == 1)
+    {
+        using t_ = fmoe_<ck_tile::fp16_t, ck_tile::fp16_t, ck_tile::fp16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, 1, 0>;
+        r = fused_moegemm_<t_>(s, a);
+    }
+    // clang-format on
+    return r;
+}

example/ck_tile/15_fused_moe/instances/fused_moegemm_api_internal.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "fused_moegemm_api_traits.hpp"
+#include "ck_tile/ops/fused_moe.hpp"
+#include <iostream>
+
+template <ck_tile::index_t... Is>
+using S = ck_tile::sequence<Is...>;
+
+// do not the define of this tepmlate function inside the _api.cpp, otherwise will block make -j
+template <typename Ts_>
+float fused_moegemm_(const ck_tile::stream_config& s, fused_moegemm_args a)
+{
+    using f_traits = ck_tile::FusedMoeGemmTraits<Ts_::GateOnly, Ts_::FusedQuant == 1, 1 /*atomic*/>;
+    using f_shape  = ck_tile::FusedMoeGemmShape<typename Ts_::BlockTile_0,
+                                               typename Ts_::WarpPerBlock_0,
+                                               typename Ts_::WarpTile_0,
+                                               typename Ts_::BlockTile_1,
+                                               typename Ts_::WarpPerBlock_0,
+                                               typename Ts_::WarpTile_0>;
+    using f_problem =
+        ck_tile::FusedMoeGemmPipelineProblem<typename Ts_::ADataType,
+                                             typename Ts_::GDataType,
+                                             typename Ts_::DDataType,
+                                             typename Ts_::AccDataType,
+                                             typename Ts_::ODataType,
+                                             typename Ts_::AScaleDataType,
+                                             typename Ts_::GScaleDataType,
+                                             typename Ts_::DScaleDataType,
+                                             typename Ts_::YSmoothScaleDataType,
+                                             typename Ts_::TopkWeightDataType,
+                                             typename Ts_::IndexDataType,
+                                             ck_tile::element_wise::FastGeluAsm, // TODO: hardcoded
+                                             f_shape,
+                                             f_traits>;
+
+    // using f_pipeline    = ck_tile::FusedMoeGemmPipeline_FlatmmEx<f_problem>;
+    using f_pipeline    = ck_tile::FusedMoeGemmPipeline_FlatmmUk<f_problem>;
+    using f_partitioner = ck_tile::FusedMoeGemmTilePartitioner_Linear<f_shape>;
+    using f_kernel      = ck_tile::FusedMoeGemmKernel<f_partitioner, f_pipeline, void>;
+
+    const dim3 grids                       = f_kernel::GridSize(a);
+    constexpr dim3 blocks                  = f_kernel::BlockSize();
+    constexpr ck_tile::index_t kBlockPerCu = 1;
+
+    static int printed = 0;
+
+    auto kargs = f_kernel::MakeKargs(a);
+    if(s.log_level_ > 0 && printed == 0)
+    {
+        std::cout << ", " << f_kernel::GetName() << std::flush;
+        printed = 1;
+    }
+
+    return ck_tile::launch_kernel(
+        s, ck_tile::make_kernel<blocks.x, kBlockPerCu>(f_kernel{}, grids, blocks, 0, kargs));
+}

example/ck_tile/15_fused_moe/instances/fused_moegemm_api_traits.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <ck_tile/core.hpp>
+
+// this is used to pattern-match internl kernel implementation, not to instantiate kernel
+template <typename I,
+          typename W,
+          typename O,
+          typename ST,
+          typename SW,
+          typename SQ,
+          typename KW,
+          typename BlockTIle_, // seq<b_token, b_interm, b_hidden, b_down>
+          typename WarpPerBlock_,
+          typename WarpTile_, // seq<*,*,*>, used to select mfma
+          ck_tile::index_t GateOnly_   = 0,
+          ck_tile::index_t FusedQuant_ = 0>
+struct fmoe_ // traits, ugly name, only used for internal
+{
+    using TypeConfig = FusedMoeGemmTypeConfig<I, W, O, ST, SW, SQ, KW>;
+
+    using ADataType            = ck_tile::remove_cvref_t<typename TypeConfig::ADataType>;
+    using GDataType            = ck_tile::remove_cvref_t<typename TypeConfig::GDataType>;
+    using DDataType            = ck_tile::remove_cvref_t<typename TypeConfig::DDataType>;
+    using AccDataType          = ck_tile::remove_cvref_t<typename TypeConfig::AccDataType>;
+    using ODataType            = ck_tile::remove_cvref_t<typename TypeConfig::ODataType>;
+    using AScaleDataType       = ck_tile::remove_cvref_t<typename TypeConfig::AScaleDataType>;
+    using GScaleDataType       = ck_tile::remove_cvref_t<typename TypeConfig::GScaleDataType>;
+    using DScaleDataType       = ck_tile::remove_cvref_t<typename TypeConfig::DScaleDataType>;
+    using YSmoothScaleDataType = ck_tile::remove_cvref_t<typename TypeConfig::YSmoothScaleDataType>;
+    using TopkWeightDataType   = ck_tile::remove_cvref_t<typename TypeConfig::TopkWeightDataType>;
+    using IndexDataType        = ck_tile::remove_cvref_t<typename TypeConfig::IndexDataType>;
+
+    static constexpr ck_tile::index_t BT_ = BlockTIle_::at(ck_tile::number<0>{}); // block token
+    static constexpr ck_tile::index_t BI_ =
+        BlockTIle_::at(ck_tile::number<1>{}); // block intermediate
+    static constexpr ck_tile::index_t BH_ = BlockTIle_::at(ck_tile::number<2>{}); // block hidden
+    static constexpr ck_tile::index_t BD_ = BlockTIle_::at(ck_tile::number<3>{}); // block down
+
+    using BlockTile_0    = ck_tile::sequence<BT_, BI_, BH_>;
+    using WarpPerBlock_0 = ck_tile::remove_cvref_t<WarpPerBlock_>;
+    using WarpTile_0     = ck_tile::remove_cvref_t<WarpTile_>;
+
+    using BlockTile_1    = ck_tile::sequence<BT_, BD_, BI_ / (GateOnly_ ? 1 : 2)>;
+    using WarpPerBlock_1 = ck_tile::remove_cvref_t<WarpPerBlock_>;
+    using WarpTile_1     = ck_tile::remove_cvref_t<WarpTile_>;
+
+    static constexpr ck_tile::index_t GateOnly   = GateOnly_;
+    static constexpr ck_tile::index_t FusedQuant = FusedQuant_;
+};

example/ck_tile/15_fused_moe/instances/fused_moegemm_bf16_m32.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <ck_tile/core.hpp>
+#include "fused_moegemm.hpp"
+#include "fused_moegemm_api_traits.hpp"
+#include "fused_moegemm_api_internal.hpp"
+
+// clang-format off
+template float fused_moegemm_<
+    fmoe_<ck_tile::bf16_t, ck_tile::bf16_t, ck_tile::bf16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, 1, 0>
+>(const ck_tile::stream_config& s, fused_moegemm_args a);
+
+// clang-format on

example/ck_tile/15_fused_moe/instances/fused_moegemm_fp16_m32.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include <ck_tile/core.hpp>
+#include "fused_moegemm.hpp"
+#include "fused_moegemm_api_traits.hpp"
+#include "fused_moegemm_api_internal.hpp"
+
+// clang-format off
+template float fused_moegemm_<
+    fmoe_<ck_tile::fp16_t, ck_tile::fp16_t, ck_tile::fp16_t, float, float, float, float, S<32, 512, 128, 128>, S<1, 4, 1>, S<16, 16, 32>, 1, 0>
+>(const ck_tile::stream_config& s, fused_moegemm_args a);
+
+// clang-format on

example/ck_tile/15_fused_moe/instances/fused_moesorting_api.cpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#include "fused_moesorting.hpp"
+
+#define MOE_SORTING_DISPATCH(unroll_num_)                                                   \
+    constexpr ck_tile::index_t unroll_num = unroll_num_;                                    \
+    using ms_problem     = ck_tile::MoeSortingProblem<index_t, ms_weight_type, unroll_num>; \
+    using kernel         = ck_tile::MoeSortingKernel<ms_problem>;                           \
+    auto kargs           = kernel::MakeKargs(a);                                            \
+    const dim3 grids     = kernel::GridSize(a);                                             \
+    const dim3 blocks    = kernel::BlockSize(a);                                            \
+    const auto lds_bytes = kernel::GetSmemSize(a);                                          \
+    float ave_time       = ck_tile::launch_kernel(                                          \
+        s, ck_tile::make_kernel(kernel{}, grids, blocks, lds_bytes, kargs));          \
+    return ave_time;
+
+float fused_moesorting(fused_moesorting_trait t, fused_moesorting_args a, ck_tile::stream_config s)
+{
+    if(t.weight_type == "fp32" && t.index_type == "int32")
+    {
+        if(a.num_experts > 127)
+        {
+            printf("lds size exceed, only support experts <127 \n");
+            return -1;
+        }
+        if(a.moe_buf_bytes % 16)
+        {
+            printf("buf set size %d unaligned, must be multiple of 16\n", a.moe_buf_bytes);
+            return -1;
+        }
+        using index_t              = ck_tile::index_t;
+        using ms_weight_type       = float;
+        index_t smem_io_unroll_num = ck_tile::integer_divide_ceil(a.tokens * a.topk, 64);
+        switch(smem_io_unroll_num)
+        {
+        case(1): {
+            MOE_SORTING_DISPATCH(1);
+        }
+        case(2): {
+            MOE_SORTING_DISPATCH(2);
+        }
+        case(3): {
+            MOE_SORTING_DISPATCH(3);
+        }
+        case(5): {
+            MOE_SORTING_DISPATCH(5);
+        }
+        case(6): {
+            MOE_SORTING_DISPATCH(6);
+        }
+        case(7): {
+            MOE_SORTING_DISPATCH(7);
+        }
+        case(8): {
+            MOE_SORTING_DISPATCH(8);
+        }
+        case(9): {
+            MOE_SORTING_DISPATCH(9);
+        }
+        case(10): {
+            MOE_SORTING_DISPATCH(10);
+        }
+        case(11): {
+            MOE_SORTING_DISPATCH(11);
+        }
+        default: {
+            MOE_SORTING_DISPATCH(4);
+        }
+        }
+    }
+    return -1;
+}

example/ck_tile/15_fused_moe/main.cpp

+#include <algorithm>
+#include <cstring>
+#include <unordered_set>
+#include <vector>
+#include <set>
+
+#include "ck_tile/host.hpp"
+#include "fused_moe.hpp"
+
+// different threshold for different dtype
+template <typename DataType>
+auto get_elimit()
+{
+    double rtol = 1e-2;
+    double atol = 1e-2;
+    return ck_tile::make_tuple(rtol, atol);
+}
+
+template <>
+auto get_elimit<ck_tile::bf16_t>()
+{
+    double rtol = 1e-2;
+    double atol = 1e-2;
+    return ck_tile::make_tuple(rtol, atol);
+}
+
+// mfma_type, 0:32x32, 1:16x16
+// TODO: padding?
+template <typename T>
+auto shuffle_moe_weight(const ck_tile::HostTensor<T>& t, std::string mfma_dtype, int mfma_type = 0)
+{
+    assert(t.get_lengths().size() == 3);
+    int b_ = t.get_lengths()[0];
+    int n_ = t.get_lengths()[1];
+    int k_ = t.get_lengths()[2];
+    if((mfma_dtype == "bf16" || mfma_dtype == "fp16") && mfma_type == 0)
+    {
+        ck_tile::HostTensor<T> t_view({b_, n_ / 32, 32, k_ / 16, 2, 8});
+        std::copy(t.begin(), t.end(), t_view.begin());
+        return ck_tile::reference_permute(t_view, {0, 1, 3, 4, 2, 5});
+    }
+    else if((mfma_dtype == "bf16" || mfma_dtype == "fp16") && mfma_type == 1)
+    {
+        ck_tile::HostTensor<T> t_view({b_, n_ / 16, 16, k_ / 32, 4, 8});
+        std::copy(t.begin(), t.end(), t_view.begin());
+        return ck_tile::reference_permute(t_view, {0, 1, 3, 4, 2, 5});
+    }
+    else if((mfma_dtype == "int8" || mfma_dtype == "fp8") && mfma_type == 0)
+    {
+        ck_tile::HostTensor<T> t_view({b_, n_ / 32, 32, k_ / 32, 2, 16});
+        std::copy(t.begin(), t.end(), t_view.begin());
+        return ck_tile::reference_permute(t_view, {0, 1, 3, 4, 2, 5});
+    }
+    else if((mfma_dtype == "int8" || mfma_dtype == "fp8") && mfma_type == 1)
+    {
+        ck_tile::HostTensor<T> t_view({b_, n_ / 16, 16, k_ / 64, 4, 16});
+        std::copy(t.begin(), t.end(), t_view.begin());
+        return ck_tile::reference_permute(t_view, {0, 1, 3, 4, 2, 5});
+    }
+    return t;
+}
+
+template <typename IndexType>
+void topid_unique_gen(
+    std::vector<IndexType>& host_tensor, int tokens, int topk, int num_expert, int seed)
+{
+    size_t total_size = topk * tokens;
+    std::srand(seed);
+    std::set<IndexType> unique_set;
+    IndexType current_v;
+    for(size_t i = 0; i < total_size; i++)
+    {
+        if(i % topk == 0)
+        {
+            unique_set.clear();
+        }
+        current_v = std::rand() % num_expert;
+        while(unique_set.find(current_v) != unique_set.end())
+        {
+            current_v = std::rand() % num_expert;
+        }
+        unique_set.insert(current_v);
+        host_tensor[i] = current_v;
+    }
+}
+
+auto create_args(int argc, char* argv[])
+{
+    ck_tile::ArgParser arg_parser;
+    arg_parser.insert("t", "128", "num input tokens")
+        .insert("e", "32", "num of experts")
+        .insert("k", "5", "topk")
+        .insert("h", "8192", "hidden_size of this model")
+        .insert("i", "8192", "intermediate_size between 2 gemms of FFN")
+        .insert("stride", "-1", "stride per row, if -1 then equal to hidden_size")
+        .insert("bm", "32", "blocking factor for sorted tokens")
+        .insert("tp", "8", "tensor parallel size")
+        .insert("v", "1", "cpu validation or not")
+        .insert("kname", "1", "print kernel name or not")
+        .insert("prec_i", "bf16", "input precision")
+        .insert("prec_w", "bf16", "weight precision")
+        .insert("prec_o", "bf16", "output precision")
+        .insert("prec_st", "auto", "token scale data type. auto will set to fp32")
+        .insert("prec_sw", "auto", "weight scale data type. auto will set to fp32")
+        .insert("prec_sq", "auto", "(dynamic) smooth quant data type. auto will set to fp32")
+        .insert("prec_kw", "auto", "topk-weight data type. auto will set to fp32")
+        .insert("fquant", "0", "fused-quant, 0:no, 1:smooth-dynamic-quant, 2:dynamic-quant")
+        .insert(
+            "gate_only", "1", "w0(gate/up) style, 0:gate+up will double interm size, 1:only gate")
+        .insert("api", "0", "benchmark api set: 0:fused-moe(moe-gemm+moe-sorting), 1:moe-gemm")
+        .insert("balance",
+                "0",
+                "if set to 1, will try balance the expert in topk-ids(convenient for testing)")
+        .insert("init",
+                "2",
+                "init method. 0:random stepped float(fast). 1: random uniform, 2:rand normalized"
+                "normalized(slow)")
+        .insert("seed", "11939", "seed used to do random")
+        .insert("warmup", "5", "cold iter")
+        .insert("repeat", "20", "hot iter");
+
+    bool result = arg_parser.parse(argc, argv);
+    return std::make_tuple(result, arg_parser);
+}
+
+// I:input-type, W:weight-type, O:output-type, ST:toke-scale-tpye, SW:weight-scale-type,
+// SQ:smooth-quant-type, KW:topk-weight-type
+template <typename I, typename W, typename O, typename ST, typename SW, typename SQ, typename KW>
+bool run(const ck_tile::ArgParser& arg_parser)
+{
+    ck_tile::index_t tokens            = arg_parser.get_int("t");
+    ck_tile::index_t experts           = arg_parser.get_int("e");
+    ck_tile::index_t topk              = arg_parser.get_int("k");
+    ck_tile::index_t hidden_size       = arg_parser.get_int("h");
+    ck_tile::index_t intermediate_size = arg_parser.get_int("i");
+    ck_tile::index_t stride            = arg_parser.get_int("stride");
+    ck_tile::index_t block_m           = arg_parser.get_int("bm");
+    if(stride < 0)
+        stride = hidden_size;
+    std::string prec_i  = arg_parser.get_str("prec_i");
+    std::string prec_w  = arg_parser.get_str("prec_w");
+    std::string prec_o  = arg_parser.get_str("prec_o");
+    std::string prec_st = arg_parser.get_str("prec_st");
+    std::string prec_sw = arg_parser.get_str("prec_sw");
+    std::string prec_sq = arg_parser.get_str("prec_sq");
+    std::string prec_kw = arg_parser.get_str("prec_kw");
+    prec_st             = (prec_st == "auto") ? "fp32" : prec_st;
+    prec_sw             = (prec_sw == "auto") ? "fp32" : prec_sw;
+    prec_sq             = (prec_sq == "auto") ? "fp32" : prec_sq;
+    prec_kw             = (prec_kw == "auto") ? "fp32" : prec_kw;
+    int kname           = arg_parser.get_int("kname");
+    int do_validation   = arg_parser.get_int("v");
+    int warmup          = arg_parser.get_int("warmup");
+    int repeat          = arg_parser.get_int("repeat");
+    int fused_quant     = arg_parser.get_int("fquant");
+    int gate_only       = arg_parser.get_int("gate_only");
+    int api             = arg_parser.get_int("api");
+    int balance         = arg_parser.get_int("balance");
+    int tp              = arg_parser.get_int("tp");
+    int init            = arg_parser.get_int("init");
+    uint32_t seed       = arg_parser.get_uint32("seed");
+
+    // w0 (Gate+Up or Gate only, N size)
+    ck_tile::index_t shared_intermediate_size_0 = intermediate_size * (gate_only ? 1 : 2) / tp;
+    // w1 (Down, N size)
+    ck_tile::index_t shared_intermediate_size_1 = intermediate_size / tp;
+
+    auto prec_str = [&]() {
+        auto base_str = prec_i;
+        if(prec_i != prec_w)
+            base_str += "x" + prec_w;
+        if(prec_i != prec_o)
+            base_str += "=" + prec_o;
+        if(fused_quant != 0)
+        {
+            base_str += std::string("(") + prec_st + "|" + prec_sw + "|" + prec_sq + ")";
+        }
+        return base_str;
+    }();
+    auto api_str = [&]() {
+        if(api == 0)
+            return std::string("fmoe");
+        else if(api == 1)
+            return std::string("moeg");
+        else if(api == 2)
+            return std::string("moes");
+        return std::string("");
+    }();
+
+    auto stride_str = [&]() {
+        if(stride == hidden_size)
+            return std::string("");
+        else
+            return std::string(", st:") + std::to_string(stride);
+    }();
+
+    std::cout << "[" << api_str << "|" << prec_str << "]"
+              << " t:" << tokens << ", e:" << experts << ", k:" << topk << stride_str
+              << ", hidden:" << hidden_size << ", interm:" << intermediate_size << ", tp:" << tp
+              << ", shrd_interm:" << shared_intermediate_size_0 << "|" << shared_intermediate_size_1
+              << ", go:" << gate_only << ", q:" << fused_quant << std::flush;
+
+    using TypeConfig           = FusedMoeGemmTypeConfig<I, W, O, ST, SW, SQ, KW>;
+    using ADataType            = typename TypeConfig::ADataType;
+    using GDataType            = typename TypeConfig::GDataType;
+    using DDataType            = typename TypeConfig::DDataType;
+    using AccDataType          = typename TypeConfig::AccDataType;
+    using ODataType            = typename TypeConfig::ODataType;
+    using AScaleDataType       = typename TypeConfig::AScaleDataType;
+    using GScaleDataType       = typename TypeConfig::GScaleDataType;
+    using DScaleDataType       = typename TypeConfig::DScaleDataType;
+    using YSmoothScaleDataType = typename TypeConfig::YSmoothScaleDataType;
+    using TopkWeightDataType   = typename TypeConfig::TopkWeightDataType;
+    using IndexDataType        = typename TypeConfig::IndexDataType;
+
+    // host verify
+    ck_tile::HostTensor<ADataType> a_host({tokens, hidden_size}, {stride, 1});
+    ck_tile::HostTensor<GDataType> g_host({experts, shared_intermediate_size_0, hidden_size});
+    ck_tile::HostTensor<DDataType> d_host({experts, hidden_size, shared_intermediate_size_1});
+    ck_tile::HostTensor<ODataType> o_host({tokens, hidden_size}, {stride, 1});
+    ck_tile::HostTensor<AScaleDataType> sa_host({tokens});
+    ck_tile::HostTensor<GScaleDataType> sg_host({shared_intermediate_size_0});
+    ck_tile::HostTensor<DScaleDataType> sd_host({shared_intermediate_size_1});
+    ck_tile::HostTensor<YSmoothScaleDataType> sy_host({shared_intermediate_size_1}); // smooth-quant
+    ck_tile::HostTensor<IndexDataType> topk_ids_host({tokens, topk});                // to be sort
+    ck_tile::HostTensor<TopkWeightDataType> topk_weight_host({tokens, topk});        // to be sort
+
+    int max_num_tokens_padded = topk * tokens + experts * block_m - topk;
+    ck_tile::HostTensor<IndexDataType> sorted_token_ids_host({max_num_tokens_padded});
+    ck_tile::HostTensor<TopkWeightDataType> sorted_weight_host({max_num_tokens_padded});
+    ck_tile::HostTensor<IndexDataType> sorted_expert_ids_host(
+        {(max_num_tokens_padded + block_m - 1) / block_m});
+    ck_tile::HostTensor<IndexDataType> num_sorted_tiles_host({1});
+
+    if(init == 0)
+    {
+        ck_tile::FillStepRange<ADataType>{-.5f, .5f, 0.01f}(a_host);
+        ck_tile::FillStepRange<GDataType>{-.5f, .5f, 0.01f}(g_host);
+        ck_tile::FillStepRange<DDataType, false>{.5f, -.5f, -0.01f}(d_host);
+        ck_tile::FillStepRange<AScaleDataType>{0.f, 1.f, 0.01f}(sa_host);
+        ck_tile::FillStepRange<GScaleDataType>{0.f, 1.f, 0.01f}(sg_host);
+        ck_tile::FillStepRange<DScaleDataType>{0.f, 1.f, 0.01f}(sd_host);
+        ck_tile::FillStepRange<YSmoothScaleDataType>{0.f, 1.f, 0.01f}(sy_host);
+        ck_tile::FillStepRange<TopkWeightDataType>{-.5f, .5f, 0.01f}(topk_weight_host);
+    }
+    else if(init == 1)
+    {
+        ck_tile::FillUniformDistribution<ADataType>{-.5f, .5f, seed, true}(a_host);
+        ck_tile::FillUniformDistribution<GDataType>{-.5f, .5f, seed, true}(g_host);
+        ck_tile::FillUniformDistribution<DDataType>{-.5f, .5f, seed, true}(d_host);
+        ck_tile::FillUniformDistribution<AScaleDataType>{-.5f, .5f, seed, true}(sa_host);
+        ck_tile::FillUniformDistribution<GScaleDataType>{-.5f, .5f, seed, true}(sg_host);
+        ck_tile::FillUniformDistribution<DScaleDataType>{-.5f, .5f, seed, true}(sd_host);
+        ck_tile::FillUniformDistribution<YSmoothScaleDataType>{-.5f, .5f, seed, true}(sy_host);
+        ck_tile::FillUniformDistribution<TopkWeightDataType>{-.5f, .5f, seed, true}(
+            topk_weight_host);
+    }
+    else if(init == 2)
+    {
+        ck_tile::FillNormalDistribution<ADataType>{0.f, 1.f, seed, true}(a_host);
+        ck_tile::FillNormalDistribution<GDataType>{0.f, 1.f, seed, true}(g_host);
+        ck_tile::FillNormalDistribution<DDataType>{0.f, 1.f, seed, true}(d_host);
+        ck_tile::FillNormalDistribution<AScaleDataType>{0.f, 1.f, seed, true}(sa_host);
+        ck_tile::FillNormalDistribution<GScaleDataType>{0.f, 1.f, seed, true}(sg_host);
+        ck_tile::FillNormalDistribution<DScaleDataType>{0.f, 1.f, seed, true}(sd_host);
+        ck_tile::FillNormalDistribution<YSmoothScaleDataType>{0.f, 1.f, seed, true}(sy_host);
+        ck_tile::FillNormalDistribution<TopkWeightDataType>{0.f, 1.f, seed, true}(topk_weight_host);
+    }
+
+    // permute weight
+    ck_tile::HostTensor<GDataType> g_perm_host = shuffle_moe_weight(g_host, prec_w, 1);
+    ck_tile::HostTensor<DDataType> d_perm_host = shuffle_moe_weight(d_host, prec_w, 1);
+
+    // do moe sorting
+    if(balance)
+    {
+        int e_cnt = 0;
+        for(int i = 0; i < static_cast<int>(topk_ids_host.mData.size()); i++)
+        {
+            topk_ids_host.mData[i] = e_cnt;
+            e_cnt++;
+            if(e_cnt >= experts)
+                e_cnt = 0;
+        }
+    }
+    else
+    {
+        topid_unique_gen<IndexDataType>(topk_ids_host.mData, tokens, topk, experts, 11913);
+    }
+
+// leave it here for future debug purpose
+#if 0
+    a_host.loadtxt("../../ater/input_torch.txt");
+
+    topk_ids_host.loadtxt("../../ater/topk_ids_torch.txt", "int");
+    // topk_ids_host.savetxt("topk_ids_2.txt");
+    topk_weight_host.loadtxt("../../ater/topk_weights_torch.txt", "float");
+    std::cout << "------- @@@ " << __LINE__ << std::flush << std::endl;
+
+    g_host.loadtxt("../../ater/w1_torch.txt", "float");
+    std::cout << "------- @@@ " << __LINE__ << std::flush << std::endl;
+    d_host.loadtxt("../../ater/w2_torch.txt", "float");
+    std::cout << "------- @@@ " << __LINE__ << std::flush << std::endl;
+
+    ck_tile::HostTensor<GDataType> g_perm_host = shuffle_moe_weight(g_host, prec_w, 1);
+    std::cout << "------- @@@ " << __LINE__ << std::flush << std::endl;
+    ck_tile::HostTensor<DDataType> d_perm_host = shuffle_moe_weight(d_host, prec_w, 1);
+    std::cout << "------- @@@ " << __LINE__ << std::flush << std::endl;
+#endif
+
+#if 0
+    std::cout << "sorted_token_ids_host:" << sorted_token_ids_host << std::endl;
+    std::cout << "num_sorted_tiles_host:" << num_sorted_tiles_host << std::endl;
+    std::cout << "sorted_expert_ids_host:" << sorted_expert_ids_host << std::endl;
+    std::cout << "topk_weight_host:" << topk_weight_host << std::endl;
+    std::cout << "sorted_weight_host:" << sorted_weight_host << std::endl;
+#endif
+    auto cal_tflops = [&](auto ms) {
+        double flop_gemm_0 =
+            2 * static_cast<double>(tokens) * topk * shared_intermediate_size_0 * hidden_size;
+        double flop_gemm_1 =
+            2 * static_cast<double>(tokens) * topk * shared_intermediate_size_1 * hidden_size;
+        return (flop_gemm_0 + flop_gemm_1) / (static_cast<double>(ms) * 1e-3) / 1e12;
+    };
+
+    // TODO: this method we use expert-by-expert view, just for reference
+    auto cal_tbps = [&](auto ms) {
+        double token_bytes =
+            static_cast<double>(tokens) * topk / experts * hidden_size * sizeof(ADataType);
+        double w0_bytes = static_cast<double>(shared_intermediate_size_0) * experts * hidden_size *
+                          sizeof(GDataType);
+        double w1_bytes = static_cast<double>(shared_intermediate_size_1) * experts * hidden_size *
+                          sizeof(DDataType);
+        double o_bytes =
+            static_cast<double>(tokens) * topk / experts * hidden_size * sizeof(ODataType);
+        double topk_weights_bytes = static_cast<double>(tokens) * topk * sizeof(TopkWeightDataType);
+        // ignore index, they are too small
+
+        return (token_bytes + w0_bytes + w1_bytes + o_bytes + topk_weights_bytes) /
+               (static_cast<double>(ms) * 1e-3) / 1e12;
+    };
+
+    if(api == 0)
+    {
+        ck_tile::DeviceMem a_buf(a_host);
+        ck_tile::DeviceMem g_perm_buf(g_perm_host);
+        ck_tile::DeviceMem d_perm_buf(d_perm_host);
+        ck_tile::DeviceMem sa_buf(sa_host);
+        ck_tile::DeviceMem sg_buf(sg_host);
+        ck_tile::DeviceMem sd_buf(sd_host);
+        ck_tile::DeviceMem sy_buf(sy_host);
+        ck_tile::DeviceMem o_buf(o_host.get_element_space_size_in_bytes());
+
+        ck_tile::DeviceMem topk_ids_buf(topk_ids_host);
+        ck_tile::DeviceMem topk_weight_buf(topk_weight_host);
+
+        ck_tile::DeviceMem sorted_token_ids_buf(
+            sorted_token_ids_host.get_element_space_size_in_bytes());
+        ck_tile::DeviceMem sorted_weight_buf(sorted_weight_host.get_element_space_size_in_bytes());
+        ck_tile::DeviceMem sorted_expert_ids_buf(
+            sorted_expert_ids_host.get_element_space_size_in_bytes());
+        ck_tile::DeviceMem num_sorted_tiles_buf(
+            num_sorted_tiles_host.get_element_space_size_in_bytes());
+
+        fused_moe_traits traits{prec_i,
+                                prec_w,
+                                prec_o,
+                                prec_st,
+                                prec_sw,
+                                prec_sq,
+                                prec_kw,
+                                block_m,
+                                gate_only,
+                                fused_quant};
+
+        fused_moe_args args{a_buf.GetDeviceBuffer(),
+                            fused_quant != 0 ? sa_buf.GetDeviceBuffer() : nullptr,
+                            g_perm_buf.GetDeviceBuffer(),
+                            d_perm_buf.GetDeviceBuffer(),
+                            fused_quant != 0 ? sg_buf.GetDeviceBuffer() : nullptr,
+                            fused_quant != 0 ? sd_buf.GetDeviceBuffer() : nullptr,
+                            fused_quant == 1 ? sy_buf.GetDeviceBuffer() : nullptr,
+                            o_buf.GetDeviceBuffer(),
+                            topk_ids_buf.GetDeviceBuffer(),
+                            topk_weight_buf.GetDeviceBuffer(),
+                            sorted_token_ids_buf.GetDeviceBuffer(),
+                            sorted_weight_buf.GetDeviceBuffer(),
+                            sorted_expert_ids_buf.GetDeviceBuffer(),
+                            num_sorted_tiles_buf.GetDeviceBuffer(),
+                            block_m,
+                            hidden_size,
+                            shared_intermediate_size_0,
+                            tokens,
+                            experts,
+                            topk,
+                            stride};
+        float ave_time = fused_moe(
+            traits, args, ck_tile::stream_config{nullptr, true, kname ? 1 : 0, warmup, repeat});
+
+        if(ave_time < 0)
+        {
+            std::cout << " not supported!" << std::endl << std::flush;
+            return false;
+        }
+
+        // float gb_per_sec = num_byte / 1.E6 / ave_time;
+        std::cout << ", " << ave_time * 1.E3 << " us, " << cal_tflops(ave_time) << " tflops, "
+                  << cal_tbps(ave_time) << " TB/s" << std::flush;
+        bool pass = true;
+
+        if(do_validation)
+        {
+            ck_tile::reference_moe_sorting<TopkWeightDataType, IndexDataType>(
+                topk_ids_host,
+                topk_weight_host,
+                sorted_token_ids_host,
+                sorted_weight_host,
+                sorted_expert_ids_host,
+                num_sorted_tiles_host.mData[0],
+                experts,
+                block_m);
+
+            ck_tile::reference_fused_moe<AccDataType, ck_tile::element_wise::Gelu>(
+                a_host,
+                g_host,
+                d_host,
+                sa_host,
+                sg_host,
+                sd_host,
+                sy_host,
+                o_host,
+                sorted_token_ids_host,
+                sorted_weight_host,
+                sorted_expert_ids_host,
+                num_sorted_tiles_host,
+                topk_ids_host,
+                block_m,
+                tokens,
+                experts,
+                hidden_size,
+                shared_intermediate_size_0,
+                topk,
+                gate_only);
+
+            auto o_dev = o_buf.ToHost<ODataType>();
+            // o_dev.savetxt("gpu-out.txt", "float");
+            auto [rtol, atol] = get_elimit<ADataType>();
+            pass &= ck_tile::check_err(
+                o_dev, o_host, std::string("OUT Error: Incorrect results!"), rtol, atol);
+            std::cout << ", valid:" << (pass ? "y" : "n") << std::flush;
+        }
+        std::cout << std::flush << std::endl;
+        return pass;
+    }
+    else if(api == 1)
+    {
+        ck_tile::reference_moe_sorting<TopkWeightDataType, IndexDataType>(
+            topk_ids_host,
+            topk_weight_host,
+            sorted_token_ids_host,
+            sorted_weight_host,
+            sorted_expert_ids_host,
+            num_sorted_tiles_host.mData[0],
+            experts,
+            block_m);
+
+        // done, preparing GPU buffer
+        ck_tile::DeviceMem a_buf(a_host);
+        ck_tile::DeviceMem g_perm_buf(g_perm_host);
+        ck_tile::DeviceMem d_perm_buf(d_perm_host);
+        ck_tile::DeviceMem sa_buf(sa_host);
+        ck_tile::DeviceMem sg_buf(sg_host);
+        ck_tile::DeviceMem sd_buf(sd_host);
+        ck_tile::DeviceMem sy_buf(sy_host);
+        ck_tile::DeviceMem o_buf(o_host);
+
+        // manually clear output buffer for atomic
+        o_buf.SetZero();
+        //
+
+        ck_tile::DeviceMem sorted_token_ids_buf(sorted_token_ids_host);
+        ck_tile::DeviceMem sorted_weight_buf(sorted_weight_host);
+        ck_tile::DeviceMem sorted_expert_ids_buf(sorted_expert_ids_host);
+        ck_tile::DeviceMem num_sorted_tiles_buf(num_sorted_tiles_host);
+
+        fused_moegemm_traits traits{prec_i,
+                                    prec_w,
+                                    prec_o,
+                                    prec_st,
+                                    prec_sw,
+                                    prec_sq,
+                                    prec_kw,
+                                    block_m,
+                                    gate_only,
+                                    fused_quant};
+
+        fused_moegemm_args args{a_buf.GetDeviceBuffer(),
+                                fused_quant != 0 ? sa_buf.GetDeviceBuffer() : nullptr,
+                                g_perm_buf.GetDeviceBuffer(),
+                                d_perm_buf.GetDeviceBuffer(),
+                                fused_quant != 0 ? sg_buf.GetDeviceBuffer() : nullptr,
+                                fused_quant != 0 ? sd_buf.GetDeviceBuffer() : nullptr,
+                                fused_quant == 1 ? sy_buf.GetDeviceBuffer() : nullptr,
+                                o_buf.GetDeviceBuffer(),
+                                sorted_token_ids_buf.GetDeviceBuffer(),
+                                sorted_weight_buf.GetDeviceBuffer(),
+                                sorted_expert_ids_buf.GetDeviceBuffer(),
+                                num_sorted_tiles_buf.GetDeviceBuffer(),
+                                hidden_size,
+                                shared_intermediate_size_0,
+                                tokens,
+                                experts,
+                                topk,
+                                stride};
+
+        float ave_time = fused_moegemm(
+            traits, args, ck_tile::stream_config{nullptr, true, kname ? 1 : 0, warmup, repeat});
+
+        if(ave_time < 0)
+        {
+            std::cout << " not supported!" << std::endl << std::flush;
+            return false;
+        }
+
+        // float gb_per_sec = num_byte / 1.E6 / ave_time;
+        std::cout << ", " << ave_time * 1.E3 << " us, " << cal_tflops(ave_time) << " tflops, "
+                  << cal_tbps(ave_time) << " TB/s" << std::flush;
+        bool pass = true;
+
+        if(do_validation)
+        {
+            ck_tile::reference_fused_moe<AccDataType, ck_tile::element_wise::Gelu>(
+                a_host,
+                g_host,
+                d_host,
+                sa_host,
+                sg_host,
+                sd_host,
+                sy_host,
+                o_host,
+                sorted_token_ids_host,
+                sorted_weight_host,
+                sorted_expert_ids_host,
+                num_sorted_tiles_host,
+                topk_ids_host,
+                block_m,
+                tokens,
+                experts,
+                hidden_size,
+                shared_intermediate_size_0,
+                topk,
+                gate_only);
+
+            auto o_dev = o_buf.ToHost<ODataType>();
+            // o_dev.savetxt("gpu-out.txt", "float");
+            auto [rtol, atol] = get_elimit<ADataType>();
+            pass &= ck_tile::check_err(
+                o_dev, o_host, std::string("OUT Error: Incorrect results!"), rtol, atol);
+            std::cout << ", valid:" << (pass ? "y" : "n") << std::flush;
+        }
+        std::cout << std::flush << std::endl;
+
+        return pass;
+    }
+    return false;
+}
+
+int main(int argc, char* argv[])
+{
+    auto [result, arg_parser] = create_args(argc, argv);
+    if(!result)
+        return -1;
+
+    std::string prec_i  = arg_parser.get_str("prec_i");
+    std::string prec_w  = arg_parser.get_str("prec_w");
+    std::string prec_o  = arg_parser.get_str("prec_o");
+    std::string prec_st = arg_parser.get_str("prec_st");
+    std::string prec_sw = arg_parser.get_str("prec_sw");
+    std::string prec_sq = arg_parser.get_str("prec_sq");
+    std::string prec_kw = arg_parser.get_str("prec_kw");
+    prec_st             = (prec_st == "auto") ? "fp32" : prec_st;
+    prec_sw             = (prec_sw == "auto") ? "fp32" : prec_sw;
+    prec_sq             = (prec_sq == "auto") ? "fp32" : prec_sq;
+    prec_kw             = (prec_kw == "auto") ? "fp32" : prec_kw;
+
+    // no dynamic quant case
+    if(prec_i == "bf16" && prec_w == "bf16" && prec_o == "bf16" && prec_kw == "fp32")
+    {
+        return run<ck_tile::bf16_t, ck_tile::bf16_t, ck_tile::bf16_t, float, float, float, float>(
+                   arg_parser)
+                   ? 0
+                   : -2;
+    }
+    else if(prec_i == "fp16" && prec_w == "fp16" && prec_o == "fp16" && prec_kw == "fp32")
+    {
+        return run<ck_tile::fp16_t, ck_tile::fp16_t, ck_tile::fp16_t, float, float, float, float>(
+                   arg_parser)
+                   ? 0
+                   : -2;
+    }
+
+    return -3;
+}