Merge branch 'develop' of https://github.com/ROCm/composable_kernel into update_cka8w8_uc

example/ck_tile/14_moe_smoothquant/instances/moe_smoothquant_fp16_n4096_instance.cpp

@@ -6,9 +6,13 @@
 
 // clang-format off
 //                                                   rm rn tm  tn  vn    pd     2p
-template float moe_smoothquant_<trait_<ck_tile::fp16_t,  1, 2, 1,  256, 8,  true,  false>>(const S&, A);
-template float moe_smoothquant_<trait_<ck_tile::fp16_t,  1, 4, 1,  256, 4,  true,  false>>(const S&, A);
-template float moe_smoothquant_<trait_<ck_tile::fp16_t,  1, 2, 1, 1024, 2,  true,  false>>(const S&, A);
-template float moe_smoothquant_<trait_<ck_tile::fp16_t,  1, 4, 1, 1024, 1,  true,  false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::int8_t, 1, 2, 1,  256, 8,  true,  false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::int8_t, 1, 4, 1,  256, 4,  true,  false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::int8_t, 1, 2, 1, 1024, 2,  true,  false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::int8_t, 1, 4, 1, 1024, 1,  true,  false>>(const S&, A);
 
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::fp8_t, 1, 2, 1,  256, 8,  true,  false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::fp8_t, 1, 4, 1,  256, 4,  true,  false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::fp8_t, 1, 2, 1, 1024, 2,  true,  false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::fp8_t, 1, 4, 1, 1024, 1,  true,  false>>(const S&, A);
 // clang-format on

example/ck_tile/14_moe_smoothquant/instances/moe_smoothquant_fp16_n4096_tp_instance.cpp

@@ -6,9 +6,13 @@
 
 // clang-format off
 //                                                   rm rn tm  tn   vn   pd    2p
-template float moe_smoothquant_<trait_<ck_tile::fp16_t,  1, 2, 1,  256, 8,  true, true>>(const S&, A);
-template float moe_smoothquant_<trait_<ck_tile::fp16_t,  1, 4, 1,  256, 4,  true, true>>(const S&, A);
-template float moe_smoothquant_<trait_<ck_tile::fp16_t,  1, 2, 1, 1024, 2,  true, true>>(const S&, A);
-template float moe_smoothquant_<trait_<ck_tile::fp16_t,  1, 4, 1, 1024, 1,  true, true>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::int8_t, 1, 2, 1,  256, 8,  true, true>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::int8_t, 1, 4, 1,  256, 4,  true, true>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::int8_t, 1, 2, 1, 1024, 2,  true, true>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::int8_t, 1, 4, 1, 1024, 1,  true, true>>(const S&, A);
 
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::fp8_t, 1, 2, 1,  256, 8,  true, true>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::fp8_t, 1, 4, 1,  256, 4,  true, true>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::fp8_t, 1, 2, 1, 1024, 2,  true, true>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::fp8_t, 1, 4, 1, 1024, 1,  true, true>>(const S&, A);
 // clang-format on

example/ck_tile/14_moe_smoothquant/instances/moe_smoothquant_fp16_n512_instance.cpp

@@ -6,8 +6,13 @@
 
 // clang-format off
 //                                                  rm rn  tm  tn  vn   pd    2p
-template float moe_smoothquant_<trait_<ck_tile::fp16_t, 1,  1,  4, 64, 8,  true , false>>(const S&, A);
-template float moe_smoothquant_<trait_<ck_tile::fp16_t, 1,  2,  4, 64, 4,  true , false>>(const S&, A);
-template float moe_smoothquant_<trait_<ck_tile::fp16_t, 1,  4,  4, 64, 2,  true , false>>(const S&, A);
-template float moe_smoothquant_<trait_<ck_tile::fp16_t, 1,  8,  4, 64, 1,  true , false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::int8_t, 1,  1,  4, 64, 8,  true , false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::int8_t, 1,  2,  4, 64, 4,  true , false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::int8_t, 1,  4,  4, 64, 2,  true , false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::int8_t, 1,  8,  4, 64, 1,  true , false>>(const S&, A);
+
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::fp8_t, 1,  1,  4, 64, 8,  true , false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::fp8_t, 1,  2,  4, 64, 4,  true , false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::fp8_t, 1,  4,  4, 64, 2,  true , false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::fp8_t, 1,  8,  4, 64, 1,  true , false>>(const S&, A);
 // clang-format on

example/ck_tile/14_moe_smoothquant/instances/moe_smoothquant_fp16_n64_n128_instance.cpp

@@ -6,7 +6,11 @@
 
 // clang-format off
 //                                                  rm  rn  tm  tn  vn  pd      2p
-template float moe_smoothquant_<trait_<ck_tile::fp16_t, 1,  1,  4,  64, 1,  true, false>>(const S&, A);
-template float moe_smoothquant_<trait_<ck_tile::fp16_t, 1,  1,  4,  64, 2,  true, false>>(const S&, A);
-template float moe_smoothquant_<trait_<ck_tile::fp16_t, 1,  2,  4,  64, 1,  true, false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::int8_t, 1,  1,  4,  64, 1,  true, false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::int8_t, 1,  1,  4,  64, 2,  true, false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::int8_t, 1,  2,  4,  64, 1,  true, false>>(const S&, A);
+
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::fp8_t, 1,  1,  4,  64, 1,  true, false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::fp8_t, 1,  1,  4,  64, 2,  true, false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::fp8_t, 1,  2,  4,  64, 1,  true, false>>(const S&, A);
 // clang-format on

example/ck_tile/14_moe_smoothquant/instances/moe_smoothquant_fp16_n768_instance.cpp

@@ -6,7 +6,11 @@
 
 // clang-format off
 //                                                  rm  rn  tm  tn  vn  pd       2p
-template float moe_smoothquant_<trait_<ck_tile::fp16_t, 1,  3,  4,  64, 4,  true , false>>(const S&, A);
-template float moe_smoothquant_<trait_<ck_tile::fp16_t, 1,  6,  4,  64, 2,  true , false>>(const S&, A);
-template float moe_smoothquant_<trait_<ck_tile::fp16_t, 1, 12,  4,  64, 1,  true , false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::int8_t, 1,  3,  4,  64, 4,  true , false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::int8_t, 1,  6,  4,  64, 2,  true , false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::int8_t, 1, 12,  4,  64, 1,  true , false>>(const S&, A);
+
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::fp8_t, 1,  3,  4,  64, 4,  true , false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::fp8_t, 1,  6,  4,  64, 2,  true , false>>(const S&, A);
+template float moe_smoothquant_<trait_<ck_tile::fp16_t, ck_tile::fp8_t, 1, 12,  4,  64, 1,  true , false>>(const S&, A);
 // clang-format on

example/ck_tile/14_moe_smoothquant/instances/moe_smoothquant_fwd_api.cpp

@@ -4,7 +4,8 @@
 #include <ck_tile/core.hpp>
 #include "moe_smoothquant.hpp"
 
-template <typename DataType_,
+template <typename InType,
+          typename OutType,
           ck_tile::index_t Repeat_M_,         // each thread repeat along M
           ck_tile::index_t Repeat_N_,         // each thread repeat along N
           ck_tile::index_t ThreadPerBlock_M_, // num threads along M
@@ -12,7 +13,8 @@ template <typename DataType_,
           ck_tile::index_t Vector_N_,         // vector size along N
           bool kPadN_,
           bool kTwoPass_>
-using trait_ = moe_smoothquant_traits_<DataType_,
+using trait_ = moe_smoothquant_traits_<InType,
+                                       OutType,
                                        Repeat_M_,
                                        Repeat_N_,
                                        ThreadPerBlock_M_,
@@ -21,108 +23,108 @@ using trait_ = moe_smoothquant_traits_<DataType_,
                                        kPadN_,
                                        kTwoPass_>;
 
-template <typename data_type>
+template <typename in_type, typename out_type>
 float moe_smoothquant_dispatch(moe_smoothquant_traits /*t*/,
                                moe_smoothquant_args a,
                                const ck_tile::stream_config& s)
 {
     float r = -1;
     // clang-format off
-    //                                         rm  rn  tm  tn  vn   pd    2p
+    //                                                    rm  rn  tm  tn  vn   pd    2p
     if(a.hidden_size <= 64) {
-            r = moe_smoothquant_<trait_<data_type, 1,  1,  4,  64, 1,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type, 1,  1,  4,  64, 1,  true, false>>(s, a);
     }
     else if(a.hidden_size <= 128) {
         if (a.hidden_size % 2 == 0)
-            r = moe_smoothquant_<trait_<data_type, 1,  1,  4,  64, 2,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type, 1,  1,  4,  64, 2,  true, false>>(s, a);
         else
-            r = moe_smoothquant_<trait_<data_type, 1,  2,  4,  64, 1,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type, 1,  2,  4,  64, 1,  true, false>>(s, a);
     }
     else if(a.hidden_size <= 256) {
         if (a.hidden_size % 4 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 1,  4,  64, 4,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 1,  4,  64, 4,  true, false>>(s, a);
         else if (a.hidden_size % 2 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 2,  4,  64, 2,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 2,  4,  64, 2,  true, false>>(s, a);
         else
-            r = moe_smoothquant_<trait_<data_type,  1, 4,  4,  64, 1,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 4,  4,  64, 1,  true, false>>(s, a);
     }
     else if(a.hidden_size <= 512) {
         if (a.hidden_size % 8 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 1,  4,  64, 8,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 1,  4,  64, 8,  true, false>>(s, a);
         else if (a.hidden_size % 4 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 2,  4,  64, 4,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 2,  4,  64, 4,  true, false>>(s, a);
         else if (a.hidden_size % 2 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 4,  4,  64, 2,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 4,  4,  64, 2,  true, false>>(s, a);
         else
-            r = moe_smoothquant_<trait_<data_type,  1, 8,  4,  64, 1,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 8,  4,  64, 1,  true, false>>(s, a);
     }
     else if(a.hidden_size <= 768) {
         if (a.hidden_size % 4 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 3,  4,  64, 4,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 3,  4,  64, 4,  true, false>>(s, a);
         else if (a.hidden_size % 2 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 6,  4,  64, 2,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 6,  4,  64, 2,  true, false>>(s, a);
         else
-            r = moe_smoothquant_<trait_<data_type,  1,12,  4,  64, 1,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1,12,  4,  64, 1,  true, false>>(s, a);
     }
     else if(a.hidden_size <= 1024) {
         if (a.hidden_size % 8 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 1, 2,  128, 8,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 1, 2,  128, 8,  true, false>>(s, a);
         else if (a.hidden_size % 4 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 2, 2,  128, 4,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 2, 2,  128, 4,  true, false>>(s, a);
         else if (a.hidden_size % 2 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 4, 2,  128, 2,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 4, 2,  128, 2,  true, false>>(s, a);
         else
-            r = moe_smoothquant_<trait_<data_type,  1, 4, 1,  256, 1,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 4, 1,  256, 1,  true, false>>(s, a);
     }
     else if(a.hidden_size <= 1536) {
         if (a.hidden_size % 8 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 3, 4,   64, 8,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 3, 4,   64, 8,  true, false>>(s, a);
         else if (a.hidden_size % 4 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 3, 2,  128, 4,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 3, 2,  128, 4,  true, false>>(s, a);
         else if (a.hidden_size % 2 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 3, 1,  256, 2,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 3, 1,  256, 2,  true, false>>(s, a);
         else
-            r = moe_smoothquant_<trait_<data_type,  1, 6, 1,  256, 1,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 6, 1,  256, 1,  true, false>>(s, a);
     }
     else if(a.hidden_size <= 2048) {
         if (a.hidden_size % 8 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 1, 1,  256, 8,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 1, 1,  256, 8,  true, false>>(s, a);
         else if (a.hidden_size % 4 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 2, 1,  256, 4,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 2, 1,  256, 4,  true, false>>(s, a);
         else if (a.hidden_size % 2 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 4, 1,  256, 2,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 4, 1,  256, 2,  true, false>>(s, a);
         else
-            r = moe_smoothquant_<trait_<data_type,  1, 8, 1,  256, 1,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 8, 1,  256, 1,  true, false>>(s, a);
     }
     else if(a.hidden_size <= 3072) {
         if (a.hidden_size % 8 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 3, 1,  128, 8,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 3, 1,  128, 8,  true, false>>(s, a);
         else if (a.hidden_size % 4 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 3, 1,  256, 4,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 3, 1,  256, 4,  true, false>>(s, a);
         else if (a.hidden_size % 2 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 6, 1,  256, 2,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 6, 1,  256, 2,  true, false>>(s, a);
         else
-            r = moe_smoothquant_<trait_<data_type,  1, 3, 1, 1024, 1,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 3, 1, 1024, 1,  true, false>>(s, a);
     }
     else if(a.hidden_size <= 4096) {
         if (a.hidden_size % 8 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 2, 1,  256, 8,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 2, 1,  256, 8,  true, false>>(s, a);
         else if (a.hidden_size % 4 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 4, 1,  256, 4,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 4, 1,  256, 4,  true, false>>(s, a);
         else if (a.hidden_size % 2 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 2, 1, 1024, 2,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 2, 1, 1024, 2,  true, false>>(s, a);
         else
-            r = moe_smoothquant_<trait_<data_type,  1, 4, 1, 1024, 1,  true, false>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 4, 1, 1024, 1,  true, false>>(s, a);
     }
     else if(a.hidden_size > 4096) {
         if (a.hidden_size % 8 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 2, 1,  256, 8,  true, true>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 2, 1,  256, 8,  true, true>>(s, a);
         else if (a.hidden_size % 4 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 4, 1,  256, 4,  true, true>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 4, 1,  256, 4,  true, true>>(s, a);
         else if (a.hidden_size % 2 == 0)
-            r = moe_smoothquant_<trait_<data_type,  1, 2, 1, 1024, 2,  true, true>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 2, 1, 1024, 2,  true, true>>(s, a);
         else
-            r = moe_smoothquant_<trait_<data_type,  1, 4, 1, 1024, 1,  true, true>>(s, a);
+            r = moe_smoothquant_<trait_<in_type, out_type,  1, 4, 1, 1024, 1,  true, true>>(s, a);
     }
     return r;
     // clang-format on
@@ -132,13 +134,21 @@ float moe_smoothquant(moe_smoothquant_traits t,
                       moe_smoothquant_args a,
                       const ck_tile::stream_config& s)
 {
-    if(t.data_type.compare("fp16") == 0)
+    if(t.in_type.compare("fp16") == 0 && t.out_type == "int8")
     {
-        return moe_smoothquant_dispatch<ck_tile::fp16_t>(t, a, s);
+        return moe_smoothquant_dispatch<ck_tile::fp16_t, ck_tile::int8_t>(t, a, s);
     }
-    else if(t.data_type.compare("bf16") == 0)
+    else if(t.in_type.compare("fp16") == 0 && t.out_type == "fp8")
     {
-        return moe_smoothquant_dispatch<ck_tile::bf16_t>(t, a, s);
+        return moe_smoothquant_dispatch<ck_tile::fp16_t, ck_tile::fp8_t>(t, a, s);
+    }
+    else if(t.in_type.compare("bf16") == 0 && t.out_type == "int8")
+    {
+        return moe_smoothquant_dispatch<ck_tile::bf16_t, ck_tile::int8_t>(t, a, s);
+    }
+    else if(t.in_type.compare("bf16") == 0 && t.out_type == "fp8")
+    {
+        return moe_smoothquant_dispatch<ck_tile::bf16_t, ck_tile::fp8_t>(t, a, s);
     }
     else
         throw std::runtime_error("Without supported instances!");

example/ck_tile/14_moe_smoothquant/instances/moe_smoothquant_instance_common.hpp

 
 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
 
 #include <ck_tile/core.hpp>
 #include "moe_smoothquant.hpp"
@@ -11,7 +11,8 @@
 using S = ck_tile::stream_config;
 using A = moe_smoothquant_args;
 
-template <typename DataType_,
+template <typename InputType_,
+          typename OutputType_,
           ck_tile::index_t Repeat_M_,         // each thread repeat along M
           ck_tile::index_t Repeat_N_,         // each thread repeat along N
           ck_tile::index_t ThreadPerBlock_M_, // num threads along M
@@ -19,7 +20,8 @@ template <typename DataType_,
           ck_tile::index_t Vector_N_,         // vector size along N
           bool kPadN_,
           bool kTwoPass_>
-using trait_ = moe_smoothquant_traits_<DataType_,
+using trait_ = moe_smoothquant_traits_<InputType_,
+                                       OutputType_,
                                        Repeat_M_,
                                        Repeat_N_,
                                        ThreadPerBlock_M_,
@@ -31,14 +33,15 @@ using trait_ = moe_smoothquant_traits_<DataType_,
 template <typename Traits_>
 float moe_smoothquant_(const S& s, A a)
 {
-    using DataType = typename Traits_::DataType;
+    using InputType  = typename Traits_::InputType;
+    using OutputType = typename Traits_::OutputType;
 
     using PipelineProblem = ck_tile::SmoothquantPipelineProblem<
-        typename MoeSmoothquantTypeConfig<DataType>::XDataType,
-        typename MoeSmoothquantTypeConfig<DataType>::XScaleDataType,
-        typename MoeSmoothquantTypeConfig<DataType>::ComputeDataType,
-        typename MoeSmoothquantTypeConfig<DataType>::YScaleDataType,
-        typename MoeSmoothquantTypeConfig<DataType>::QYDataType,
+        typename MoeSmoothquantTypeConfig<InputType, OutputType>::XDataType,
+        typename MoeSmoothquantTypeConfig<InputType, OutputType>::SmoothScaleDataType,
+        typename MoeSmoothquantTypeConfig<InputType, OutputType>::ComputeDataType,
+        typename MoeSmoothquantTypeConfig<InputType, OutputType>::YScaleDataType,
+        typename MoeSmoothquantTypeConfig<InputType, OutputType>::QYDataType,
         typename Traits_::Shape,
         Traits_::kPadN,
         Traits_::kTwoPass>;

example/ck_tile/14_moe_smoothquant/moe_smoothquant.cpp

@@ -63,7 +63,8 @@ auto create_args(int argc, char* argv[])
         .insert("stride", "-1", "stride per row, if -1 then equal to hidden_size")
         .insert("v", "1", "cpu validation or not")
         .insert("kname", "1", "print kernel name or not")
-        .insert("prec", "fp16", "precision")
+        .insert("prec_i", "fp16", "input precision, fp16/bf16")
+        .insert("prec_o", "int8", "precision, int8/fp8")
         .insert("warmup", "5", "cold iter")
         .insert("repeat", "20", "hot iter");
 
@@ -71,7 +72,7 @@ auto create_args(int argc, char* argv[])
     return std::make_tuple(result, arg_parser);
 }
 
-template <typename DataType>
+template <typename InputType, typename OutputType>
 bool run(const ck_tile::ArgParser& arg_parser)
 {
     ck_tile::index_t tokens      = arg_parser.get_int("t");
@@ -81,7 +82,8 @@ bool run(const ck_tile::ArgParser& arg_parser)
         stride = hidden_size;
     ck_tile::index_t experts = arg_parser.get_int("e");
     ck_tile::index_t topk    = arg_parser.get_int("k");
-    std::string data_type    = arg_parser.get_str("prec");
+    std::string prec_i       = arg_parser.get_str("prec_i");
+    std::string prec_o       = arg_parser.get_str("prec_o");
     int kname                = arg_parser.get_int("kname");
     int do_validation        = arg_parser.get_int("v");
     int warmup               = arg_parser.get_int("warmup");
@@ -89,17 +91,17 @@ bool run(const ck_tile::ArgParser& arg_parser)
 
     assert(stride >= hidden_size);
 
-    using TypeConfig = MoeSmoothquantTypeConfig<DataType>;
+    using TypeConfig = MoeSmoothquantTypeConfig<InputType, OutputType>;
 
-    using XDataType       = typename TypeConfig::XDataType;
-    using XScaleDataType  = typename TypeConfig::XScaleDataType;
-    using YScaleDataType  = typename TypeConfig::YScaleDataType;
-    using QYDataType      = typename TypeConfig::QYDataType;
-    using ComputeDataType = typename TypeConfig::ComputeDataType;
+    using XDataType           = typename TypeConfig::XDataType;
+    using SmoothScaleDataType = typename TypeConfig::SmoothScaleDataType;
+    using YScaleDataType      = typename TypeConfig::YScaleDataType;
+    using QYDataType          = typename TypeConfig::QYDataType;
+    using ComputeDataType     = typename TypeConfig::ComputeDataType;
 
     // host verify
     ck_tile::HostTensor<XDataType> x_host({tokens, hidden_size}, {stride, 1});
-    ck_tile::HostTensor<XScaleDataType> xscale_host({experts * hidden_size});
+    ck_tile::HostTensor<SmoothScaleDataType> smscale_host({experts * hidden_size});
     ck_tile::HostTensor<ck_tile::index_t> topk_ids_host({tokens, topk});
 
     ck_tile::HostTensor<YScaleDataType> yscale_host_ref({topk * tokens}, {1});
@@ -110,26 +112,26 @@ bool run(const ck_tile::ArgParser& arg_parser)
 
     topid_unique_gen<ck_tile::index_t>(topk_ids_host.mData, tokens, topk, experts, 11937);
     ck_tile::FillUniformDistribution<XDataType>{-.5f, .5f}(x_host);
-    ck_tile::FillUniformDistribution<XScaleDataType>{1e-3, .5f}(xscale_host);
+    ck_tile::FillUniformDistribution<SmoothScaleDataType>{1e-3, .5f}(smscale_host);
 
     ck_tile::DeviceMem x_buf(x_host.get_element_space_size_in_bytes());
-    ck_tile::DeviceMem xscale_buf(xscale_host.get_element_space_size_in_bytes());
+    ck_tile::DeviceMem smscale_buf(smscale_host.get_element_space_size_in_bytes());
     ck_tile::DeviceMem topk_ids_buf(topk_ids_host.get_element_space_size_in_bytes());
     ck_tile::DeviceMem yscale_buf(yscale_host_dev.get_element_space_size_in_bytes());
     ck_tile::DeviceMem qy_buf(qy_host_dev.get_element_space_size_in_bytes());
 
     x_buf.ToDevice(x_host.data());
-    xscale_buf.ToDevice(xscale_host.data());
+    smscale_buf.ToDevice(smscale_host.data());
     topk_ids_buf.ToDevice(topk_ids_host.data());
 
-    std::cout << "[" << data_type << "]"
+    std::cout << "[" << prec_i << "-" << prec_o << "]"
               << " tokens:" << tokens << ", hidden_size:" << hidden_size << ", stride:" << stride
               << ", experts:" << experts << ", topk:" << topk << std::flush;
 
-    moe_smoothquant_traits traits{data_type};
+    moe_smoothquant_traits traits{prec_i, prec_o};
 
     moe_smoothquant_args args{x_buf.GetDeviceBuffer(),
-                              xscale_buf.GetDeviceBuffer(),
+                              smscale_buf.GetDeviceBuffer(),
                               topk_ids_buf.GetDeviceBuffer(),
                               yscale_buf.GetDeviceBuffer(),
                               qy_buf.GetDeviceBuffer(),
@@ -143,9 +145,10 @@ bool run(const ck_tile::ArgParser& arg_parser)
     float ave_time = moe_smoothquant(
         traits, args, ck_tile::stream_config{nullptr, true, kname ? 1 : 0, warmup, repeat});
 
-    std::size_t num_byte =
-        sizeof(XDataType) * tokens * hidden_size + sizeof(XScaleDataType) * topk * hidden_size +
-        sizeof(YScaleDataType) * topk * tokens + sizeof(QYDataType) * topk * tokens * hidden_size;
+    std::size_t num_byte = sizeof(XDataType) * tokens * hidden_size +
+                           sizeof(SmoothScaleDataType) * topk * hidden_size +
+                           sizeof(YScaleDataType) * topk * tokens +
+                           sizeof(QYDataType) * topk * tokens * hidden_size;
 
     float gb_per_sec = num_byte / 1.E6 / ave_time;
     std::cout << ", " << ave_time * 1.E3 << " us, " << gb_per_sec << " GB/s" << std::flush;
@@ -165,11 +168,11 @@ bool run(const ck_tile::ArgParser& arg_parser)
 
                     for(int i_h = 0; i_h < hidden_size; ++i_h)
                     {
-                        auto v_xscale = ck_tile::type_convert<ComputeDataType>(
-                            xscale_host(i_expert * hidden_size + i_h));
+                        auto v_smscale = ck_tile::type_convert<ComputeDataType>(
+                            smscale_host(i_expert * hidden_size + i_h));
                         auto v_x = ck_tile::type_convert<ComputeDataType>(x_host(i_token, i_h));
-                        // y_host(i_token * topk + i_topk, i_h) = v_x * v_xscale;
-                        y_host(i_topk * tokens + i_token, i_h) = v_x * v_xscale;
+                        // y_host(i_token * topk + i_topk, i_h) = v_x * v_smscale;
+                        y_host(i_topk * tokens + i_token, i_h) = v_x * v_smscale;
                     }
                 }
             };
@@ -250,14 +253,23 @@ int main(int argc, char* argv[])
     if(!result)
         return -1;
 
-    const std::string data_type = arg_parser.get_str("prec");
-    if(data_type == "fp16")
+    const std::string prec_i = arg_parser.get_str("prec_i");
+    const std::string prec_o = arg_parser.get_str("prec_o");
+    if(prec_i == "fp16" && prec_o == "int8")
+    {
+        return run<ck_tile::half_t, ck_tile::int8_t>(arg_parser) ? 0 : -2;
+    }
+    else if(prec_i == "fp16" && prec_o == "fp8")
+    {
+        return run<ck_tile::half_t, ck_tile::fp8_t>(arg_parser) ? 0 : -2;
+    }
+    else if(prec_i == "bf16" && prec_o == "int8")
     {
-        return run<ck_tile::half_t>(arg_parser) ? 0 : -2;
+        return run<ck_tile::bf16_t, ck_tile::int8_t>(arg_parser) ? 0 : -2;
     }
-    else if(data_type == "bf16")
+    else if(prec_i == "bf16" && prec_o == "fp8")
     {
-        return run<ck_tile::bf16_t>(arg_parser) ? 0 : -2;
+        return run<ck_tile::bf16_t, ck_tile::fp8_t>(arg_parser) ? 0 : -2;
     }
 
     return -3;

example/ck_tile/14_moe_smoothquant/moe_smoothquant.hpp

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
@@ -8,27 +8,14 @@
 #include "ck_tile/ops/smoothquant.hpp"
 #include <string>
 
-template <typename DataType>
-struct MoeSmoothquantTypeConfig;
-
-template <>
-struct MoeSmoothquantTypeConfig<ck_tile::half_t>
-{
-    using XDataType       = ck_tile::half_t;
-    using XScaleDataType  = float;
-    using YScaleDataType  = float;
-    using QYDataType      = ck_tile::int8_t;
-    using ComputeDataType = float;
-};
-
-template <>
-struct MoeSmoothquantTypeConfig<ck_tile::bf16_t>
+template <typename InputType, typename OutputType>
+struct MoeSmoothquantTypeConfig
 {
-    using XDataType       = ck_tile::bf16_t;
-    using XScaleDataType  = float;
-    using YScaleDataType  = float;
-    using QYDataType      = ck_tile::int8_t;
-    using ComputeDataType = float;
+    using XDataType           = InputType;
+    using SmoothScaleDataType = float;
+    using YScaleDataType      = float;
+    using QYDataType          = OutputType;
+    using ComputeDataType     = float;
 };
 
 // runtime args
@@ -37,7 +24,8 @@ struct moe_smoothquant_args : public ck_tile::MoeSmoothquantHostArgs
 };
 
 // this is used to pattern-match internl kernel implementation, not to instantiate kernel
-template <typename DataType_,
+template <typename InputType_,
+          typename OutputType_,
           ck_tile::index_t Repeat_M_,         // each thread repeat along M
           ck_tile::index_t Repeat_N_,         // each thread repeat along N
           ck_tile::index_t ThreadPerBlock_M_, // num threads along M
@@ -47,7 +35,8 @@ template <typename DataType_,
           bool kTwoPass_>
 struct moe_smoothquant_traits_
 {
-    using DataType = ck_tile::remove_cvref_t<DataType_>;
+    using InputType  = ck_tile::remove_cvref_t<InputType_>;
+    using OutputType = ck_tile::remove_cvref_t<OutputType_>;
 
     static constexpr bool is_warp_per_row = ThreadPerBlock_N_ <= warpSize;
     static_assert((ThreadPerBlock_M_ * ThreadPerBlock_N_) % warpSize == 0);
@@ -108,7 +97,8 @@ float moe_smoothquant_(const ck_tile::stream_config& s, moe_smoothquant_args a);
 // This is the public API, will be generated by script
 struct moe_smoothquant_traits
 {
-    std::string data_type;
+    std::string in_type;  // input type
+    std::string out_type; // output type
 };
 
 float moe_smoothquant(moe_smoothquant_traits, moe_smoothquant_args, const ck_tile::stream_config&);

example/ck_tile/14_moe_smoothquant/script/smoke_test.sh

@@ -2,29 +2,31 @@
 EXE=build/bin/tile_example_moe_smoothquant
 
 for pr_i in "fp16" "bf16" ; do
-$EXE -prec=$pr_i -t=99  -h=13
-$EXE -prec=$pr_i -t=17  -h=16
-$EXE -prec=$pr_i -t=1   -h=100
-$EXE -prec=$pr_i -t=4   -h=128
-$EXE -prec=$pr_i -t=80  -h=127
-$EXE -prec=$pr_i -t=22  -h=255 -stride=256
-$EXE -prec=$pr_i -t=7   -h=599
-$EXE -prec=$pr_i -t=19  -h=512
-$EXE -prec=$pr_i -t=33  -h=313 -stride=1000
-$EXE -prec=$pr_i -t=11  -h=510
-$EXE -prec=$pr_i -t=171 -h=676 -stride=818
-$EXE -prec=$pr_i -t=91  -h=636
-$EXE -prec=$pr_i -t=12  -h=768 -stride=800
-$EXE -prec=$pr_i -t=100 -h=766 -stride=812
-$EXE -prec=$pr_i -t=31  -h=1024
-$EXE -prec=$pr_i -t=64  -h=1000 -stride=1004
-$EXE -prec=$pr_i -t=8   -h=1501
-$EXE -prec=$pr_i -t=3   -h=1826
-$EXE -prec=$pr_i -t=5   -h=2040
-$EXE -prec=$pr_i -t=7   -h=2734
-$EXE -prec=$pr_i -t=1   -h=3182
-$EXE -prec=$pr_i -t=9   -h=4096
-$EXE -prec=$pr_i -t=3   -h=8192
-$EXE -prec=$pr_i -t=1   -h=10547
-$EXE -prec=$pr_i -t=3   -h=17134
+for pr_o in "int8" "fp8" ; do
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=99  -h=13
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=17  -h=16
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=1   -h=100
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=4   -h=128
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=80  -h=127
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=22  -h=255 -stride=256
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=7   -h=599
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=19  -h=512
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=33  -h=313 -stride=1000
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=11  -h=510
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=171 -h=676 -stride=818
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=91  -h=636
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=12  -h=768 -stride=800
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=100 -h=766 -stride=812
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=31  -h=1024
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=64  -h=1000 -stride=1004
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=8   -h=1501
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=3   -h=1826
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=5   -h=2040
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=7   -h=2734
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=1   -h=3182
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=9   -h=4096
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=3   -h=8192
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=1   -h=10547
+$EXE -prec_i=$pr_i -prec_o=$pr_o -t=3   -h=17134
+done
 done

example/ck_tile/15_fused_moe/README.md

@@ -8,6 +8,9 @@ The benifit of this fused-moe:
 * much less kernel instance, easy to maintain
 
 # Implementation and feature support
+## NOTES:
+currently gate+up in fp16 case will very easily cause accumulator overflow the fp16 max(65504), hence result in INF. Please use BF16 for gate+up case, API side will have no check for this.
+
 ## 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)
 

example/ck_tile/15_fused_moe/fused_moe.hpp

@@ -26,7 +26,7 @@ struct fused_moe_args
 
     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 intermediate_size; // n / TP, for Gate. and Up, Down is also this value
     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?
@@ -45,7 +45,8 @@ struct fused_moe_traits
     std::string prec_sq; // smooth quant scale
     std::string prec_kw; // topk-weight data type
     int block_m;
-    int gate_only;
+    int activation;  // 0:gelu, 1:silu
+    int gate_only;   // 0:g1u0, 1:g1u1
     int fused_quant; // 0:no-sweep, 1:smooth-dynamic-quant, 2:dynamic-quant
 };
 

example/ck_tile/15_fused_moe/fused_moegemm.hpp

@@ -77,7 +77,8 @@ struct fused_moegemm_traits
     std::string prec_sq; // smooth quant scale
     std::string prec_kw; // topk-weight data type
     int block_m;
-    int gate_only;
+    int activation;  // 0:gelu, 1:silu
+    int gate_only;   // 0:g1u0, 1:g1u1
     int fused_quant; // 0:no-sweep, 1:smooth-dynamic-quant, 2:dynamic-quant
 };
 

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

@@ -41,6 +41,7 @@ float fused_moe(fused_moe_traits t, fused_moe_args a, const ck_tile::stream_conf
                                    t.prec_sq,
                                    t.prec_kw,
                                    t.block_m,
+                                   t.activation,
                                    t.gate_only,
                                    t.fused_quant};
     auto a1 = fused_moegemm_args{

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

@@ -17,15 +17,67 @@ float fused_moegemm(fused_moegemm_traits t, fused_moegemm_args a, const ck_tile:
     // 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)
+       t.prec_sw == "fp32" && t.prec_sq == "fp32" && t.prec_kw == "fp32" && t.block_m == 32 && t.gate_only == 1 && t.activation == 0)
     {
-        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>;
+        constexpr ck_tile::index_t act_ = 0;
+        constexpr ck_tile::index_t go_  = 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>, act_, go_, 0>;
+        r = fused_moegemm_<t_>(s, a);
+    }
+    else 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 == 0 && t.activation == 0)
+    {
+        constexpr ck_tile::index_t act_ = 0;
+        constexpr ck_tile::index_t go_  = 0;
+        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>, act_, go_, 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 && t.activation == 0)
+    {
+        constexpr ck_tile::index_t act_ = 0;
+        constexpr ck_tile::index_t go_  = 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>, act_, go_, 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 == 0 && t.activation == 0)
+    {
+        constexpr ck_tile::index_t act_ = 0;
+        constexpr ck_tile::index_t go_  = 0;
+        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>, act_, go_, 0>;
+        r = fused_moegemm_<t_>(s, a);
+    }
+    else 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 && t.activation == 1)
+    {
+        constexpr ck_tile::index_t act_ = 1;
+        constexpr ck_tile::index_t go_  = 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>, act_, go_, 0>;
+        r = fused_moegemm_<t_>(s, a);
+    }
+    else 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 == 0 && t.activation == 1)
+    {
+        constexpr ck_tile::index_t act_ = 1;
+        constexpr ck_tile::index_t go_  = 0;
+        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>, act_, go_, 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 && t.activation == 1)
+    {
+        constexpr ck_tile::index_t act_ = 1;
+        constexpr ck_tile::index_t go_  = 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>, act_, go_, 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)
+       t.prec_sw == "fp32" && t.prec_sq == "fp32" && t.prec_kw == "fp32" && t.block_m == 32 && t.gate_only == 0 && t.activation == 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>;
+        constexpr ck_tile::index_t act_ = 1;
+        constexpr ck_tile::index_t go_  = 0;
+        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>, act_, go_, 0>;
         r = fused_moegemm_<t_>(s, a);
     }
     // clang-format on

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

@@ -21,21 +21,31 @@ float fused_moegemm_(const ck_tile::stream_config& s, fused_moegemm_args a)
                                                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>;
+
+    constexpr auto get_activation_ = []() {
+        if constexpr(Ts_::Activation == 0)
+        {
+            return ck_tile::element_wise::FastGeluAsm{};
+        }
+        else
+            return ck_tile::element_wise::Silu{};
+    };
+    using f_act_ = ck_tile::remove_cvref_t<decltype(get_activation_())>;
+
+    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,
+                                                           f_act_, // TODO: hardcoded
+                                                           f_shape,
+                                                           f_traits>;
 
     // using f_pipeline    = ck_tile::FusedMoeGemmPipeline_FlatmmEx<f_problem>;
     using f_pipeline    = ck_tile::FusedMoeGemmPipeline_FlatmmUk<f_problem>;

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

@@ -15,7 +15,8 @@ template <typename I,
           typename KW,
           typename BlockTIle_, // seq<b_token, b_interm, b_hidden, b_down>
           typename WarpPerBlock_,
-          typename WarpTile_, // seq<*,*,*>, used to select mfma
+          typename WarpTile_,               // seq<*,*,*>, used to select mfma
+          ck_tile::index_t Activation_ = 0, // 0: Gelu 1: Silu
           ck_tile::index_t GateOnly_   = 0,
           ck_tile::index_t FusedQuant_ = 0>
 struct fmoe_ // traits, ugly name, only used for internal
@@ -44,10 +45,11 @@ struct fmoe_ // traits, ugly name, only used for internal
     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 BlockTile_1    = ck_tile::sequence<BT_, BD_, BI_>;
     using WarpPerBlock_1 = ck_tile::remove_cvref_t<WarpPerBlock_>;
     using WarpTile_1     = ck_tile::remove_cvref_t<WarpTile_>;
 
+    static constexpr ck_tile::index_t Activation = Activation_; // 0: Gelu 1: Silu
     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

@@ -8,7 +8,18 @@
 
 // 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>
+    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>, 0, 0, 0>
 >(const ck_tile::stream_config& s, fused_moegemm_args a);
 
+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>, 0, 1, 0>
+>(const ck_tile::stream_config& s, fused_moegemm_args a);
+
+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, 0>
+>(const ck_tile::stream_config& s, fused_moegemm_args a);
+
+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, 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

@@ -8,7 +8,19 @@
 
 // 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>
+    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>, 0, 0, 0>
+>(const ck_tile::stream_config& s, fused_moegemm_args a);
+
+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>, 0, 1, 0>
+>(const ck_tile::stream_config& s, fused_moegemm_args a);
+
+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, 0>
+>(const ck_tile::stream_config& s, fused_moegemm_args a);
+
+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, 1, 0>
 >(const ck_tile::stream_config& s, fused_moegemm_args a);
 
 // clang-format on

example/ck_tile/15_fused_moe/main.cpp

@@ -108,12 +108,14 @@ auto create_args(int argc, char* argv[])
         .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("act", "0", "activation after first gemm. 0:gelu, 1:silu")
         .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"
+                "1",
+                "init method. 0:random stepped float(fast). 1: random uniform[-0.5, 0.5], 2:rand "
+                "normalized[0, 1]"
                 "normalized(slow)")
         .insert("seed", "11939", "seed used to do random")
         .insert("warmup", "5", "cold iter")
@@ -135,6 +137,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
     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");
+    ck_tile::index_t activation        = arg_parser.get_int("act");
     if(stride < 0)
         stride = hidden_size;
     std::string prec_i  = arg_parser.get_str("prec_i");
@@ -194,11 +197,14 @@ bool run(const ck_tile::ArgParser& arg_parser)
             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;
+    std::cout
+        << "[" << api_str << "|" << prec_str << "]"
+        << " t:" << tokens << ", e:" << experts << ", k:" << topk << stride_str
+        << ", hidden:" << hidden_size << ", interm:" << intermediate_size << ", tp:" << tp
+        << ", act:"
+        << activation
+        // << ", shrd_interm:" << shared_intermediate_size_0 << "|" << shared_intermediate_size_1
+        << (gate_only ? ", g1u0" : ", g1u1") << ", q:" << fused_quant << std::flush;
 
     using TypeConfig           = FusedMoeGemmTypeConfig<I, W, O, ST, SW, SQ, KW>;
     using ADataType            = typename TypeConfig::ADataType;
@@ -370,6 +376,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
                                 prec_sq,
                                 prec_kw,
                                 block_m,
+                                activation,
                                 gate_only,
                                 fused_quant};
 
@@ -389,7 +396,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
                             num_sorted_tiles_buf.GetDeviceBuffer(),
                             block_m,
                             hidden_size,
-                            shared_intermediate_size_0,
+                            intermediate_size / tp,
                             tokens,
                             experts,
                             topk,
@@ -408,6 +415,28 @@ bool run(const ck_tile::ArgParser& arg_parser)
                   << cal_tbps(ave_time) << " TB/s" << std::flush;
         bool pass = true;
 
+#define CPU_FUSED_MOE(act_type_)                                                 \
+    ck_tile::reference_fused_moe<AccDataType, act_type_>(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,            \
+                                                         intermediate_size / tp, \
+                                                         topk,                   \
+                                                         gate_only)
+
         if(do_validation)
         {
             ck_tile::reference_moe_sorting<TopkWeightDataType, IndexDataType>(
@@ -419,28 +448,14 @@ bool run(const ck_tile::ArgParser& arg_parser)
                 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);
+            if(activation == 0)
+            {
+                CPU_FUSED_MOE(ck_tile::element_wise::Gelu);
+            }
+            else
+            {
+                CPU_FUSED_MOE(ck_tile::element_wise::Silu);
+            }
 
             auto o_dev = o_buf.ToHost<ODataType>();
             // o_dev.savetxt("gpu-out.txt", "float");
@@ -491,6 +506,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
                                     prec_sq,
                                     prec_kw,
                                     block_m,
+                                    activation,
                                     gate_only,
                                     fused_quant};
 
@@ -507,7 +523,7 @@ bool run(const ck_tile::ArgParser& arg_parser)
                                 sorted_expert_ids_buf.GetDeviceBuffer(),
                                 num_sorted_tiles_buf.GetDeviceBuffer(),
                                 hidden_size,
-                                shared_intermediate_size_0,
+                                intermediate_size / tp,
                                 tokens,
                                 experts,
                                 topk,
@@ -529,27 +545,14 @@ bool run(const ck_tile::ArgParser& arg_parser)
 
         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);
+            if(activation == 0)
+            {
+                CPU_FUSED_MOE(ck_tile::element_wise::Gelu);
+            }
+            else
+            {
+                CPU_FUSED_MOE(ck_tile::element_wise::Silu);
+            }
 
             auto o_dev = o_buf.ToHost<ODataType>();
             // o_dev.savetxt("gpu-out.txt", "float");