[CK_TILE] Fix mock token id, support g1u1/g1u0 through same inline co…

…de block (#1808)

* fix mock token id

* prepare host for g1u1

* reformat inline-asm

* restructure uk_0

* restructure gate_up

* done

* change default to init=1

* update readme

* fix a bug in interleave pipeline

* rcp for silu

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