/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/docs-7.0.1/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm_policy.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/docs-7.0.1/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm_policy.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/docs-7.0.1/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_flatmm_policy.hpp Source File
Go to the documentation of this file.
 // 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/ops/fused_moe/pipeline/fused_moegemm_traits.hpp"
 #include "ck_tile/ops/flatmm.hpp"
 #include "ck_tile/ops/gemm/warp/warp_gemm.hpp"
 #include "ck_tile/ops/gemm/warp/warp_gemm_dispatcher.hpp"
  
 namespace ck_tile {
  
 struct FusedMoeGemmPipelineFlatmmPolicy
 {
     CK_TILE_HOST_DEVICE static constexpr index_t GetAsyncCopyDwords()
     {
         // TODO: always 1 dword
         return 1;
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetAlignment_A()
     {
         // using async
         constexpr index_t copy_bytes = 4 * GetAsyncCopyDwords();
         constexpr index_t data_bytes = sizeof(typename Problem::ADataType);
         static_assert(copy_bytes % data_bytes == 0);
         return copy_bytes / data_bytes;
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetAlignment_G()
     {
         constexpr index_t copy_bytes = [&]() { return 16; }();
         constexpr index_t data_bytes = sizeof(typename Problem::GDataType);
         static_assert(copy_bytes % data_bytes == 0);
         return copy_bytes / data_bytes;
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetAlignment_D()
     {
         constexpr index_t copy_bytes = [&]() { return 16; }();
         constexpr index_t data_bytes = sizeof(typename Problem::DDataType);
         static_assert(copy_bytes % data_bytes == 0);
         return copy_bytes / data_bytes;
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetAlignment_O()
     {
         if constexpr(Problem::Traits::OAtomic == 1)
         {
             // pack fp16/bf16 atomic
             static_assert(sizeof(typename Problem::ODataType) == 2);
             return 2;
         }
         else if constexpr(Problem::Traits::OAtomic == 2)
         {
             // fp32 atomic
             return 1;
         }
         else
         {
             return 16 / sizeof(typename Problem::ODataType);
         }
     }
  
     template <typename DataType_>
     CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPack()
     {
         // TODO: this is for 3d layout
         return 16 / sizeof(remove_cvref_t<DataType_>);
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPack_A()
     {
         return GetSmemKPack<typename Problem::ADataType>();
     }
  
     // used for bridge LDS shuffle
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPack_Y()
     {
         // TODO: this should match mfma layout
         return 16 / sizeof(typename Problem::YDataType);
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize_A()
     {
         constexpr auto a_sld_desc = MakeLdsLoadDesc_A<Problem>();
         constexpr auto a_sst_desc = MakeLdsStoreDesc_A<Problem>();
         static_assert(a_sld_desc.get_element_space_size() == a_sst_desc.get_element_space_size());
         return a_sld_desc.get_element_space_size();
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize_Bridge()
     {
         constexpr auto bridge_sld_desc = MakeBridgeLdsLoadDesc<Problem>();
         constexpr auto bridge_sst_desc = MakeBridgeLdsStoreDesc<Problem>();
         static_assert(bridge_sld_desc.get_element_space_size() ==
                       bridge_sst_desc.get_element_space_size());
         return bridge_sld_desc.get_element_space_size();
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetSmemSize()
     {
         constexpr index_t a_lds      = GetSmemSize_A<Problem>();
         constexpr index_t bridge_lds = GetSmemSize_Bridge<Problem>();
         return max(a_lds, bridge_lds);
     }
  
     template <index_t MPerBlock, index_t KPerBlock, index_t NumWarps, index_t Alignment>
     CK_TILE_HOST_DEVICE static constexpr auto MakeGlobalTileDistribution_SimpleMxK()
     {
         constexpr index_t K_vec = Alignment;
         constexpr index_t K_rem = KPerBlock / K_vec;
  
         if constexpr(get_warp_size() < K_rem)
         {
             static_assert(K_rem % get_warp_size() == 0);
             constexpr index_t K_lan = get_warp_size(); // lane within same wave is along gemm-k
             constexpr index_t K_wav = K_rem / get_warp_size();
             static_assert(K_wav <= NumWarps, "not not support thread has repeat along K yet");
             constexpr index_t M_wav = NumWarps / K_wav;
             static_assert(MPerBlock % M_wav == 0, "this tile size is too small please check");
             constexpr index_t M_rep = MPerBlock / M_wav;
  
             return make_static_tile_distribution(
                 tile_distribution_encoding<
                     sequence<1>,
                     tuple<sequence<M_rep, M_wav>, sequence<K_wav, K_lan, K_vec>>,
                     tuple<sequence<1, 2>, sequence<2>>,
                     tuple<sequence<1, 0>, sequence<1>>,
                     sequence<1, 2>,
                     sequence<0, 2>>{});
         }
         else
         {
             constexpr index_t K_lan = K_rem;
             constexpr index_t M_lan = get_warp_size() / K_lan;
             constexpr index_t M_wav = NumWarps;
             static_assert(MPerBlock % (M_lan * M_wav) == 0,
                           "this tile size is too small please check");
             constexpr index_t M_rep = MPerBlock / (M_lan * M_wav);
             return make_static_tile_distribution(
                 tile_distribution_encoding<
                     sequence<1>,
                     tuple<sequence<M_rep, M_wav, M_lan>, sequence<K_lan, K_vec>>,
                     tuple<sequence<1>, sequence<1, 2>>,
                     tuple<sequence<1>, sequence<2, 0>>,
                     sequence<1, 2>,
                     sequence<0, 1>>{});
         }
     }
  
     // optimized version for async, not same as simple MXK dist(pay attention!!)
     template <index_t MPerBlock, index_t KPerBlock, index_t NumWarps, index_t Alignment>
     CK_TILE_HOST_DEVICE static constexpr auto MakeGlobalTileDistribution_SimpleMxK_Async()
     {
         constexpr index_t K_vec = Alignment;
         constexpr index_t K_rem = KPerBlock / K_vec;
  
         if constexpr(get_warp_size() <= K_rem)
         {
             static_assert(K_rem % get_warp_size() == 0);
             constexpr index_t K_lan = get_warp_size(); // lane within same wave is along gemm-k
             constexpr index_t K_wav = K_rem / get_warp_size();
             static_assert(K_wav <= NumWarps, "do not support thread has repeat along K yet");
             constexpr index_t M_wav = NumWarps / K_wav;
             static_assert(MPerBlock % M_wav == 0, "this tile size is too small please check");
             constexpr index_t M_rep = MPerBlock / M_wav;
             // NOTE: no swap, but hard to avoid LDS bank conflict
             return make_static_tile_distribution(
                 tile_distribution_encoding<
                     sequence<1>,
                     tuple<sequence<M_rep, M_wav>, sequence<K_wav, K_lan, K_vec>>,
                     tuple<sequence<1, 2>, sequence<2>>,
                     tuple<sequence<1, 0>, sequence<1>>,
                     sequence<1, 2>,
                     sequence<0, 2>>{});
         }
         else
         {
             constexpr index_t K_lan = K_rem;
             constexpr index_t M_lan = get_warp_size() / K_lan;
             constexpr index_t M_wav = NumWarps;
             static_assert(MPerBlock % (M_lan * M_wav) == 0,
                           "this tile size is too small please check");
             constexpr index_t M_rep = MPerBlock / (M_lan * M_wav);
             // NOTE: swapped for LDS load bank conflict free
             return make_static_tile_distribution(
                 tile_distribution_encoding<
                     sequence<1>,
                     // Note M_wave(num waves) is the fastest dim, different from sipmle 2d
                     // distribution
                     tuple<sequence<M_rep, M_lan, M_wav>, sequence<K_lan, K_vec>>,
                     tuple<sequence<1>, sequence<1, 2>>,
                     tuple<sequence<2>, sequence<1, 0>>,
                     sequence<1, 2>,
                     sequence<0, 1>>{});
         }
     }
  
     template <index_t WarpPerBlock_N_,
               index_t WarpPerBlock_K_,
               index_t Repeat_N_,
               index_t Repeat_K_,
               index_t WarpSize_,
               index_t Alignment_>
     CK_TILE_HOST_DEVICE static constexpr auto MakeGlobalTileDistribution_Nr_Kr_W()
     {
         return make_static_tile_distribution(
             tile_distribution_encoding<sequence<1>,
                                        tuple<sequence<Repeat_N_, WarpPerBlock_N_>,
                                              sequence<Repeat_K_, WarpPerBlock_K_>,
                                              sequence<WarpSize_, Alignment_>>,
                                        tuple<sequence<1, 2>, sequence<3>>,
                                        tuple<sequence<1, 1>, sequence<0>>,
                                        sequence<1, 2, 3>,
                                        sequence<0, 0, 1>>{});
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto MakeGlobalTileDistribution_A()
     {
         constexpr index_t Block_M_   = Problem::BlockShape::Block_M0;
         constexpr index_t Block_K_   = Problem::BlockShape::Block_K0;
         constexpr index_t NumWarps_  = Problem::BlockShape::NumWarps;
         constexpr index_t Alignment_ = GetAlignment_A<Problem>();
         return MakeGlobalTileDistribution_SimpleMxK_Async<Block_M_,
                                                           Block_K_,
                                                           NumWarps_,
                                                           Alignment_>();
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto MakeGlobalTileDistribution_G()
     {
         constexpr auto PermuteEnum = Problem::Traits::PermuteEnum;
         // constexpr index_t hidden_radio_0 = Problem::Traits::IsGateOnly ? 1 : 2;
         using S_ = typename Problem::BlockShape;
         if constexpr(PermuteEnum == FusedMoeGemmWeightPermuteEnum::b_nr_kr_waveflatten)
         {
             // number<S_::WarpPerBlock_N0>{}.rrr();
             // number<S_::Repeat_N0>{}.eee();
             return MakeGlobalTileDistribution_Nr_Kr_W<S_::WarpPerBlock_N0,
                                                       S_::WarpPerBlock_K0,
                                                       S_::Repeat_N0, 
                                                       S_::Repeat_K0,
                                                       get_warp_size(),
                                                       GetAlignment_G<Problem>()>();
         }
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto MakeGlobalTileDistribution_D()
     {
         constexpr auto PermuteEnum = Problem::Traits::PermuteEnum;
         using S_                   = typename Problem::BlockShape;
         if constexpr(PermuteEnum == FusedMoeGemmWeightPermuteEnum::b_nr_kr_waveflatten)
         {
             return MakeGlobalTileDistribution_Nr_Kr_W<S_::WarpPerBlock_N1,
                                                       S_::WarpPerBlock_K1,
                                                       S_::Repeat_N1,
                                                       S_::Repeat_K1,
                                                       get_warp_size(),
                                                       GetAlignment_D<Problem>()>();
         }
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto MakeGlobalTileDistribution_O()
     {
         using S_       = remove_cvref_t<typename Problem::BlockShape>;
         using WarpGemm = remove_cvref_t<decltype(GetWarpGemm1<Problem>())>;
         // using CDataType = typename WarpGemm::CDataType;
  
         constexpr auto c_block_outer_dstr_encoding =
             tile_distribution_encoding<sequence<>,
                                        tuple<sequence<S_::Repeat_M1, S_::WarpPerBlock_M1>,
                                              sequence<S_::Repeat_N1, S_::WarpPerBlock_N1>>,
                                        tuple<sequence<1, 2>>,
                                        tuple<sequence<1, 1>>,
                                        sequence<1, 2>,
                                        sequence<0, 0>>{};
  
         constexpr auto c_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
             c_block_outer_dstr_encoding, typename WarpGemm::CWarpDstrEncoding{});
         constexpr auto c_block_dstr = make_static_tile_distribution(c_block_dstr_encode);
         return c_block_dstr;
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto MakeLdsStoreDesc_A()
     {
         // A async->LDS
         constexpr index_t Block_M = Problem::BlockShape::Block_M0;
         constexpr index_t Block_K = Problem::BlockShape::Block_K0;
         // constexpr index_t BlockSize = Problem::BlockShape::BlockSize;
         constexpr index_t WarpSize = ck_tile::get_warp_size();
         constexpr index_t NumWarps = Problem::BlockShape::NumWarps;
  
         constexpr index_t KPack   = GetSmemKPack_A<Problem>(); // LDS
         constexpr index_t KVector = GetAlignment_A<Problem>(); // async copy 1 dword
         constexpr index_t KPad    = KPack;                     // pad between warps
  
         static_assert(Block_K % KVector == 0);
         constexpr index_t LanesPerK = Block_K / KVector; // how many thread loading K
         if constexpr(LanesPerK >= WarpSize)
         {
             // need multiple waves to load K
             static_assert(LanesPerK % WarpSize == 0);
             constexpr index_t wavesPerK = LanesPerK / WarpSize;
             if constexpr(wavesPerK > NumWarps)
             {
                 // TODO: need multiple issues along K to load all data
             }
             else
             {
                 constexpr index_t wavesPerM     = NumWarps / wavesPerK;
                 constexpr index_t NumIssues     = Block_M / wavesPerM;
                 constexpr auto lds_block_desc_0 = make_naive_tensor_descriptor(
                     make_tuple(number<NumIssues>{},                             // m0
                                number<wavesPerM>{},                             // m1
                                number<wavesPerK>{},                             // k0
                                number<WarpSize>{},                              // k1
                                number<KVector>{}),                              // k2
                     make_tuple(number<NumWarps*(WarpSize * KVector + KPad)>{},  // m0
                                number<wavesPerK*(WarpSize * KVector + KPad)>{}, // m1
                                number<WarpSize * KVector + KPad>{},             // k0
                                number<KVector>{},                               // k1
                                number<1>{}),                                    // k2
                     number<KVector>{}, // lds store vector(actually no explicit store)
                     number<1>{});
  
                 constexpr auto lds_block_desc_issues_warps_lanes = transform_tensor_descriptor(
                     lds_block_desc_0,
                     make_tuple(
                         make_pass_through_transform(number<NumIssues>{}),
                         make_merge_transform(make_tuple(number<wavesPerM>{}, number<wavesPerK>{})),
                         make_merge_transform(make_tuple(number<WarpSize>{}, number<KVector>{}))),
                     make_tuple(sequence<0>{}, sequence<1, 2>{}, sequence<3, 4>{}),
                     make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}));
  
                 return lds_block_desc_issues_warps_lanes;
             }
         }
         else
         {
             // lanes within a wave load different M but same K
             static_assert(WarpSize % LanesPerK == 0);
             constexpr index_t LaneGroups = WarpSize / LanesPerK; // along m
             constexpr index_t NumIssues  = Block_M / (LaneGroups * NumWarps);
  
             constexpr auto lds_block_desc_0 = make_naive_tensor_descriptor(
                 make_tuple(number<NumIssues>{},                            // m0
                            number<LaneGroups>{},                           // m1
                            number<NumWarps>{},                             // m2
                            number<LanesPerK>{},                            // k0
                            number<KVector>{}),                             // k1
                 make_tuple(number<NumWarps*(WarpSize * KVector + KPad)>{}, // m0
                            number<Block_K>{},                              // m1
                            number<WarpSize * KVector + KPad>{},            // m2
                            number<KVector>{},                              // k0
                            number<1>{}),                                   // k1
                 number<KVector>{}, // lds store vector(actually no explicit store)
                 number<1>{});
  
             constexpr auto lds_block_desc_issues_warps_lanes = transform_tensor_descriptor(
                 lds_block_desc_0,
                 make_tuple(make_pass_through_transform(number<NumIssues>{}),
                            make_pass_through_transform(number<NumWarps>{}),
                            make_merge_transform(make_tuple(
                                number<LaneGroups>{}, number<LanesPerK>{}, number<KVector>{}))),
                 make_tuple(sequence<0>{}, sequence<2>{}, sequence<1, 3, 4>{}),
                 make_tuple(sequence<0>{}, sequence<1>{}, sequence<2>{}));
  
             return lds_block_desc_issues_warps_lanes;
         }
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto MakeLdsLoadDesc_A()
     {
         // A async->LDS
         // Note that, this descriptor is only to construct the layout inside LDS
         // in real Gemm pipeline, ds_read may not follow this pattern
         // (may follow that in tile_distribution)
         // below code is almost the same as SmemStore dist, with difference:
         //  1). modify the GuaranteedLastDimensionVectorLength of naive tensor desc
         //  2). return discriptor is in NxK 2d layout
         constexpr index_t Block_M = Problem::BlockShape::Block_M0;
         constexpr index_t Block_K = Problem::BlockShape::Block_K0;
         // constexpr index_t BlockSize = Problem::BlockShape::BlockSize;
         constexpr index_t WarpSize = ck_tile::get_warp_size();
         constexpr index_t NumWarps = Problem::BlockShape::NumWarps;
  
         constexpr index_t KPack   = GetSmemKPack_A<Problem>(); // LDS
         constexpr index_t KVector = GetAlignment_A<Problem>(); // async copy 1 dword
         constexpr index_t KPad    = KPack;                     // pad between warps
  
         static_assert(Block_K % KVector == 0);
         constexpr index_t LanesPerK = Block_K / KVector; // how many thread loading K
         if constexpr(LanesPerK >= WarpSize)
         {
             // need multiple waves to load K
             static_assert(LanesPerK % WarpSize == 0);
             constexpr index_t wavesPerK = LanesPerK / WarpSize;
             if constexpr(wavesPerK >= NumWarps)
             {
                 // TODO: need multiple issues along K to load all data
             }
             else
             {
                 constexpr index_t wavesPerM     = NumWarps / wavesPerK;
                 constexpr index_t NumIssues     = Block_M / wavesPerM;
                 constexpr auto lds_block_desc_0 = make_naive_tensor_descriptor(
                     make_tuple(number<NumIssues>{},                             // m0
                                number<wavesPerM>{},                             // m1
                                number<wavesPerK>{},                             // k0
                                number<WarpSize>{},                              // k1
                                number<KVector>{}),                              // k2
                     make_tuple(number<NumWarps*(WarpSize * KVector + KPad)>{},  // m0
                                number<wavesPerK*(WarpSize * KVector + KPad)>{}, // m1
                                number<WarpSize * KVector + KPad>{},             // k0
                                number<KVector>{},                               // k1
                                number<1>{}),                                    // k2
                     number<KPack>{},                                            // lds load vector
                     number<1>{});
  
                 constexpr auto lds_desc_m_k = transform_tensor_descriptor(
                     lds_block_desc_0,
                     make_tuple(
                         make_merge_transform(make_tuple(number<NumIssues>{}, number<wavesPerM>{})),
                         make_merge_transform(make_tuple(
                             number<wavesPerK>{}, number<WarpSize>{}, number<KVector>{}))),
                     make_tuple(sequence<0, 1>{}, sequence<2, 3, 4>{}),
                     make_tuple(sequence<0>{}, sequence<1>{}));
  
                 return lds_desc_m_k;
             }
         }
         else
         {
             // lanes within a wave load different M but same K
             static_assert(WarpSize % LanesPerK == 0);
             constexpr index_t LaneGroups = WarpSize / LanesPerK; // along m
             constexpr index_t NumIssues  = Block_M / (LaneGroups * NumWarps);
  
             constexpr auto lds_block_desc_0 = make_naive_tensor_descriptor(
                 make_tuple(number<NumIssues>{},                            // m0
                            number<LaneGroups>{},                           // m1
                            number<NumWarps>{},                             // m2
                            number<LanesPerK>{},                            // k0
                            number<KVector>{}),                             // k1
                 make_tuple(number<NumWarps*(WarpSize * KVector + KPad)>{}, // m0
                            number<Block_K>{},                              // m1
                            number<WarpSize * KVector + KPad>{},            // m2
                            number<KVector>{},                              // k0
                            number<1>{}),                                   // k1
                 number<KPack>{},                                           // lds load vector
                 number<1>{});
  
             constexpr auto lds_desc_m_k = transform_tensor_descriptor(
                 lds_block_desc_0,
                 make_tuple(
                     make_merge_transform(
                         make_tuple(number<NumIssues>{}, number<LaneGroups>{}, number<NumWarps>{})),
                     make_merge_transform(make_tuple(number<LanesPerK>{}, number<KVector>{}))),
                 make_tuple(sequence<0, 1, 2>{}, sequence<3, 4>{}),
                 make_tuple(sequence<0>{}, sequence<1>{}));
  
             return lds_desc_m_k;
         }
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto MakeBridgeLdsLoadDesc()
     {
         constexpr index_t Block_M = Problem::BlockShape::Block_M0;
         constexpr index_t Block_N = Problem::BlockShape::Block_N0;
  
         constexpr index_t KVector = GetSmemKPack_Y<Problem>(); // async copy 1 dword
         constexpr index_t KPad    = 0;                         // pad between warps
  
         constexpr auto desc =
             make_naive_tensor_descriptor(make_tuple(number<Block_M>{}, number<Block_N>{}),
                                          make_tuple(number<Block_N + KPad>{}, number<1>{}),
                                          number<KVector>{},
                                          number<1>{});
         return desc;
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto MakeBridgeLdsStoreDesc()
     {
         constexpr index_t Block_M = Problem::BlockShape::Block_M0;
         constexpr index_t Block_N = Problem::BlockShape::Block_N0;
  
         constexpr index_t KVector = GetSmemKPack_Y<Problem>(); // async copy 1 dword
         constexpr index_t KPad    = 0; // KVector;                   // pad between warps
  
         constexpr auto desc =
             make_naive_tensor_descriptor(make_tuple(number<Block_M>{}, number<Block_N>{}),
                                          make_tuple(number<Block_N + KPad>{}, number<1>{}),
                                          number<KVector>{},
                                          number<1>{});
         return desc;
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto MakeBridgeLdsStoreForUKDesc()
     {
         constexpr index_t WarpPerBlock_N = Problem::BlockShape::WarpPerBlock_N0;
         constexpr index_t Repeat_N       = Problem::BlockShape::Repeat_N0;
         constexpr index_t Repeat_M       = Problem::BlockShape::Repeat_M0;
  
         constexpr index_t kAMLane     = 16;
         constexpr index_t kABKLane    = 4;
         constexpr index_t kABKPerLane = 4;
  
         constexpr index_t KPack = kABKPerLane;
  
         constexpr auto lds_block_desc_0 = make_naive_tensor_descriptor(
             make_tuple(number<Repeat_M>{},                                               // m
                        number<Repeat_N>{},                                               // n
                        number<WarpPerBlock_N>{},                                         // n
                        number<kABKLane>{},                                               // n
                        number<kAMLane>{},                                                // m
                        number<KPack>{}),                                                 // n
             make_tuple(number<Repeat_N * WarpPerBlock_N * kABKLane * kAMLane * KPack>{}, //  m
                        number<WarpPerBlock_N * kABKLane * kAMLane * KPack>{},            //  n
                        number<kABKLane * kAMLane * KPack>{},                             //  n
                        number<kAMLane * KPack>{},                                        //  n
                        number<KPack>{},                                                  //  m
                        number<1>{}),                                                     //  n
             number<KPack>{}, // lds store vector(actually no explicit store)
             number<1>{});
  
         constexpr auto desc = transform_tensor_descriptor(
             lds_block_desc_0,
             make_tuple(make_merge_transform(make_tuple(number<Repeat_M>{}, number<kAMLane>{})),
                        make_merge_transform(make_tuple(number<Repeat_N>{},
                                                        number<WarpPerBlock_N>{},
                                                        number<kABKLane>{},
                                                        number<KPack>{}))),
             make_tuple(sequence<0, 4>{}, sequence<1, 2, 3, 5>{}),
             make_tuple(sequence<0>{}, sequence<1>{}));
  
         return desc;
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetWarpGemm0()
     {
         using S_ = typename Problem::BlockShape;
         // A is vgpr, B is agpr. But since we transposed, so also need swap this
         // TODO: this is ugly
         constexpr auto wg_ctrl = WGAttrCtlEnum::Raw_avv;
         // TODO: ugly
         if constexpr(std::is_same_v<typename Problem::ADataType, ck_tile::bf16_t> &&
                      std::is_same_v<typename Problem::GDataType, ck_tile::bf16_t> &&
                      S_::Warp_M0 == 32 && S_::Warp_N0 == 32 && S_::Warp_K0 == 16)
         {
             return WarpGemmImpl<WarpGemmAtrributeMfmaIterateKAndTransposedCDistribution_SwizzleB<
                 WarpGemmAttributeMfmaImplBf16Bf16F32M32N32K8<wg_ctrl>,
                 2>>{};
         }
         else if constexpr(std::is_same_v<typename Problem::ADataType, ck_tile::int8_t> &&
                           std::is_same_v<typename Problem::GDataType, ck_tile::int8_t> &&
                           S_::Warp_M0 == 32 && S_::Warp_N0 == 32 && S_::Warp_K0 == 32)
         {
             return WarpGemmImpl<WarpGemmAtrributeMfmaIterateKAndTransposedCDistribution_SwizzleB<
                 WarpGemmAttributeMfmaImpl_i32_32x32x16_i8<wg_ctrl>,
                 2>>{};
         }
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetSequencer_0()
     {
         // this function return seq<...> used to identify gld/sld/valu... inside mfma sequence
         // the purpose is to hide thoes instructions under mfma
         // every value inside seq<...> is a mask, indicating a specific operation
         using S_                = typename Problem::BlockShape;
         constexpr index_t SLD_A = static_cast<index_t>(FusedMoeGemmPipelineSequencerEnum::SLD_A);
         constexpr index_t GLD_A = static_cast<index_t>(FusedMoeGemmPipelineSequencerEnum::GLD_A);
         constexpr index_t GLD_B = static_cast<index_t>(FusedMoeGemmPipelineSequencerEnum::GLD_B);
         if constexpr(std::is_same_v<typename Problem::YDataType, ck_tile::bf16_t> &&
                      std::is_same_v<typename Problem::DDataType, ck_tile::bf16_t> &&
                      S_::Warp_M0 == 32 && S_::Warp_N0 == 32 && S_::Warp_K0 == 16 &&
                      S_::Block_M0 == 32 && S_::Block_N0 == 512 && S_::Block_K0 == 128 &&
                      S_::Block_N1 == 128)
         {
             // Total 64 instructions, 32 buffer-load-dwordx4 gld_b, 8x buffer-load-dwordx1-async
             // gld_a 8x ds_read_b128 sld_a total 64 slot :)
             // clang-format off
             constexpr auto seq_all =
                     //       0       1       2        3       4      5        6       7
                    sequence<GLD_B,  GLD_A,  GLD_B,  GLD_A,  GLD_B,  GLD_A,  GLD_B,  GLD_A,    // 0
                             GLD_B,  GLD_A,  GLD_B,  GLD_A,  GLD_B,  GLD_A,  GLD_B,  GLD_A,    // 1
                             GLD_B,  SLD_A,  GLD_B,  SLD_A,  GLD_B,  SLD_A,  GLD_B,  SLD_A,    // 2
                             GLD_B,  SLD_A,  GLD_B,  SLD_A,  GLD_B,  SLD_A,  GLD_B,  SLD_A,    // 3
                             GLD_B,      0,  GLD_B,      0,  GLD_B,      0,  GLD_B,      0,    // 4
                             GLD_B,      0,  GLD_B,      0,  GLD_B,      0,  GLD_B,      0,    // 5
                             GLD_B,      0,  GLD_B,      0,  GLD_B,      0,  GLD_B,      0,    // 6
                             GLD_B,      0,  GLD_B,      0,  GLD_B,      0,  GLD_B,      0>{}; // 7
             return seq_all;
             // clang-format on
         }
         else if constexpr(std::is_same_v<typename Problem::YDataType, ck_tile::bf16_t> &&
                           std::is_same_v<typename Problem::DDataType, ck_tile::bf16_t> &&
                           S_::Warp_M0 == 32 && S_::Warp_N0 == 32 && S_::Warp_K0 == 16 &&
                           S_::Block_M0 == 32 && S_::Block_N0 == 256 && S_::Block_K0 == 128 &&
                           S_::Block_N1 == 128)
         {
             // Total 32 instructions, 16 buffer-load-dwordx4 gld_b, 8x buffer-load-dwordx1-async
             // gld_a 8x ds_read_b128 sld_a total 64 slot :)
             // clang-format off
             constexpr auto seq_all =
                     //       0       1       2        3       4      5        6       7
                    sequence<GLD_B,  GLD_A,  GLD_B,  GLD_A,  GLD_B,  GLD_A,  GLD_B,  GLD_A,    // 0
                             GLD_B,  GLD_A,  GLD_B,  GLD_A,  GLD_B,  GLD_A,  GLD_B,  GLD_A,    // 1
                             GLD_B,  SLD_A,  GLD_B,  SLD_A,  GLD_B,  SLD_A,  GLD_B,  SLD_A,    // 2
                             GLD_B,  SLD_A,  GLD_B,  SLD_A,  GLD_B,  SLD_A,  GLD_B,  SLD_A>{};    // 3
             return seq_all;
             // clang-format on
         }
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetSequencer_1()
     {
         // this function return seq<...> used to identify gld/sld/valu... inside mfma sequence
         // the purpose is to hide thoes instructions under mfma
         // every value inside seq<...> is a mask, indicating a specific operation
         using S_                = typename Problem::BlockShape;
         constexpr index_t GLD_B = static_cast<index_t>(FusedMoeGemmPipelineSequencerEnum::GLD_B);
         constexpr index_t GST_O = static_cast<index_t>(FusedMoeGemmPipelineSequencerEnum::GST_O);
         if constexpr(std::is_same_v<typename Problem::YDataType, ck_tile::bf16_t> &&
                      std::is_same_v<typename Problem::DDataType, ck_tile::bf16_t> &&
                      S_::Warp_M1 == 32 && S_::Warp_N1 == 32 && S_::Warp_K1 == 16 &&
                      S_::Block_M0 == 32 && S_::Block_N0 == 512 && S_::Block_K0 == 128 &&
                      S_::Block_N1 == 128)
         {
             // Total 64 instructions, 32 buffer-load-dwordx4 gld_b, 8x buffer-load-dwordx1-async
             // gld_a 8x ds_read_b128 sld_a total 64 slot :)
             // clang-format off
             constexpr auto seq_all =
                     //       0       1       2        3       4      5        6       7
                    sequence<GLD_B,  GST_O,  GLD_B,  GST_O,  GLD_B,  GST_O,  GLD_B,  GST_O,    // 0
                             GLD_B,  GST_O,  GLD_B,  GST_O,  GLD_B,  GST_O,  GLD_B,  GST_O,    // 1
                             GLD_B,      0,  GLD_B,      0,  GLD_B,      0,  GLD_B,      0,    // 2
                             GLD_B,      0,  GLD_B,      0,  GLD_B,      0,  GLD_B,      0,    // 3
                             GLD_B,      0,  GLD_B,      0,  GLD_B,      0,  GLD_B,      0,    // 4
                             GLD_B,      0,  GLD_B,      0,  GLD_B,      0,  GLD_B,      0,    // 5
                             GLD_B,      0,  GLD_B,      0,  GLD_B,      0,  GLD_B,      0,    // 6
                             GLD_B,      0,  GLD_B,      0,  GLD_B,      0,  GLD_B,      0>{}; // 7
             return seq_all;
             // clang-format on
         }
         else if constexpr(std::is_same_v<typename Problem::YDataType, ck_tile::bf16_t> &&
                           std::is_same_v<typename Problem::DDataType, ck_tile::bf16_t> &&
                           S_::Warp_M1 == 32 && S_::Warp_N1 == 32 && S_::Warp_K1 == 16 &&
                           S_::Block_M0 == 32 && S_::Block_N0 == 256 && S_::Block_K0 == 128 &&
                           S_::Block_N1 == 128)
         {
             // Total 64 instructions, 32 buffer-load-dwordx4 gld_b, 8x buffer-load-dwordx1-async
             // gld_a 8x ds_read_b128 sld_a total 64 slot :)
             // clang-format off
             constexpr auto seq_all =
                     //       0       1       2        3       4      5        6       7
                    sequence<GLD_B,  GST_O,  GLD_B,  GST_O,  GLD_B,  GST_O,  GLD_B,  GST_O,    // 0
                             GLD_B,  GST_O,  GLD_B,  GST_O,  GLD_B,  GST_O,  GLD_B,  GST_O,    // 1
                             GLD_B,      0,  GLD_B,      0,  GLD_B,      0,  GLD_B,      0,    // 2
                             GLD_B,      0,  GLD_B,      0,  GLD_B,      0,  GLD_B,      0>{};    // 3
             return seq_all;
             // clang-format on
         }
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetWarpGemm1()
     {
         using S_               = typename Problem::BlockShape;
         constexpr auto wg_ctrl = WGAttrCtlEnum::Raw_avv;
         // TODO: ugly
         if constexpr(std::is_same_v<typename Problem::YDataType, ck_tile::bf16_t> &&
                      std::is_same_v<typename Problem::DDataType, ck_tile::bf16_t> &&
                      S_::Warp_M0 == 32 && S_::Warp_N0 == 32 && S_::Warp_K0 == 16)
         {
             return WarpGemmImpl<WarpGemmAtrributeMfmaIterateKAndTransposedCDistribution_SwizzleB<
                 WarpGemmAttributeMfmaImplBf16Bf16F32M32N32K8<wg_ctrl>,
                 2>>{};
         }
         else if constexpr(std::is_same_v<typename Problem::YDataType, ck_tile::int8_t> &&
                           std::is_same_v<typename Problem::DDataType, ck_tile::int8_t> &&
                           S_::Warp_M0 == 32 && S_::Warp_N0 == 32 && S_::Warp_K0 == 32)
         {
             return WarpGemmImpl<WarpGemmAtrributeMfmaIterateKAndTransposedCDistribution_SwizzleB<
                 WarpGemmAttributeMfmaImpl_i32_32x32x16_i8<wg_ctrl>,
                 2>>{};
         }
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto MakeCBlockTile_Gemm0()
     {
         using S_        = remove_cvref_t<typename Problem::BlockShape>;
         using WarpGemm  = remove_cvref_t<decltype(GetWarpGemm0<Problem>())>;
         using CDataType = typename WarpGemm::CDataType;
  
         constexpr auto c_block_outer_dstr_encoding =
             tile_distribution_encoding<sequence<>,
                                        tuple<sequence<S_::Repeat_M0, S_::WarpPerBlock_M0>,
                                              sequence<S_::Repeat_N0, S_::WarpPerBlock_N0>>,
                                        tuple<sequence<1, 2>>,
                                        tuple<sequence<1, 1>>,
                                        sequence<1, 2>,
                                        sequence<0, 0>>{};
  
         constexpr auto c_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
             c_block_outer_dstr_encoding, typename WarpGemm::CWarpDstrEncoding{});
         constexpr auto c_block_dstr = make_static_tile_distribution(c_block_dstr_encode);
         auto c_block_tensor         = make_static_distributed_tensor<CDataType>(c_block_dstr);
         return c_block_tensor;
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto MakeCBlockTile_Gemm1()
     {
         using S_        = remove_cvref_t<typename Problem::BlockShape>;
         using WarpGemm  = remove_cvref_t<decltype(GetWarpGemm1<Problem>())>;
         using CDataType = typename WarpGemm::CDataType;
  
         constexpr auto c_block_outer_dstr_encoding =
             tile_distribution_encoding<sequence<>,
                                        tuple<sequence<S_::Repeat_M1, S_::WarpPerBlock_M1>,
                                              sequence<S_::Repeat_N1, S_::WarpPerBlock_N1>>,
                                        tuple<sequence<1, 2>>,
                                        tuple<sequence<1, 1>>,
                                        sequence<1, 2>,
                                        sequence<0, 0>>{};
  
         constexpr auto c_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
             c_block_outer_dstr_encoding, typename WarpGemm::CWarpDstrEncoding{});
         constexpr auto c_block_dstr = make_static_tile_distribution(c_block_dstr_encode);
         auto c_block_tensor         = make_static_distributed_tensor<CDataType>(c_block_dstr);
         return c_block_tensor;
     }
  
     // this is used as A matrix for 2nd gemm
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto MakeYTileDistribution()
     {
         using S_       = remove_cvref_t<typename Problem::BlockShape>;
         using WarpGemm = remove_cvref_t<decltype(GetWarpGemm1<Problem>())>;
  
         // TODO: all waves a along different N, but same M
         constexpr auto y_outer_dstr_enc =
             tile_distribution_encoding<sequence<S_::WarpPerBlock_M1>,
                                        tuple<sequence<S_::Repeat_M1>, sequence<S_::Repeat_K1>>,
                                        tuple<sequence<0>>,
                                        tuple<sequence<0>>,
                                        sequence<1, 2>,
                                        sequence<0, 0>>{};
  
         constexpr auto y_block_dstr_encode = detail::make_embed_tile_distribution_encoding(
             y_outer_dstr_enc, typename WarpGemm::AWarpDstrEncoding{});
         constexpr auto y_block_dstr = make_static_tile_distribution(y_block_dstr_encode);
         return y_block_dstr;
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto MakeYBlockTile()
     {
         constexpr auto y_block_dstr = MakeYTileDistribution<Problem>();
         auto y_block_tensor =
             make_static_distributed_tensor<typename Problem::YDataType>(y_block_dstr);
         return y_block_tensor;
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetUK_0()
     {
         using S_ = typename Problem::BlockShape;
         if constexpr(std::is_same_v<typename Problem::ADataType, ck_tile::bf16_t> &&
                      std::is_same_v<typename Problem::GDataType, ck_tile::bf16_t> &&
                      S_::Block_M0 == 32 && S_::Block_N0 == 512 && S_::Block_K0 == 128 &&
                      S_::Warp_M0 == 16 && S_::Warp_N0 == 16 && S_::Warp_K0 == 32)
         {
             return Flatmm_32x512x128_1x4x1_16x16x32_BF16{};
         }
         else if constexpr(std::is_same_v<typename Problem::ADataType, ck_tile::fp16_t> &&
                           std::is_same_v<typename Problem::GDataType, ck_tile::fp16_t> &&
                           S_::Block_M0 == 32 && S_::Block_N0 == 512 && S_::Block_K0 == 128 &&
                           S_::Warp_M0 == 16 && S_::Warp_N0 == 16 && S_::Warp_K0 == 32)
         {
             return Flatmm_32x512x128_1x4x1_16x16x32_FP16{};
         }
     }
  
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetUK_1()
     {
         using S_ = typename Problem::BlockShape;
         using T_ = typename Problem::Traits;
         if constexpr(std::is_same_v<typename Problem::YDataType, ck_tile::bf16_t> &&
                      std::is_same_v<typename Problem::DDataType, ck_tile::bf16_t> &&
                      std::is_same_v<typename Problem::TopkWeightDataType, float> &&
                      S_::Block_M1 == 32 && S_::Block_N1 == 128 && S_::Block_K1 == 512 &&
                      S_::Warp_M0 == 16 && S_::Warp_N0 == 16 && S_::Warp_K0 == 32 &&
                      T_::PipeInterleave == false)
         {
             return FlatmmSn_32x128x512_1x4x1_16x16x32_BF16{};
             // return FlatmmSn_32x128x512_1x4x1_16x16x32_BF16_itl{};
         }
         else if constexpr(std::is_same_v<typename Problem::YDataType, ck_tile::fp16_t> &&
                           std::is_same_v<typename Problem::DDataType, ck_tile::fp16_t> &&
                           std::is_same_v<typename Problem::TopkWeightDataType, float> &&
                           S_::Block_M1 == 32 && S_::Block_N1 == 128 && S_::Block_K1 == 512 &&
                           S_::Warp_M0 == 16 && S_::Warp_N0 == 16 && S_::Warp_K0 == 32 &&
                           T_::PipeInterleave == false)
         {
             return FlatmmSn_32x128x512_1x4x1_16x16x32_FP16{};
             // return FlatmmSn_32x128x512_1x4x1_16x16x32_FP16_itl{};
         }
         else if constexpr(std::is_same_v<typename Problem::YDataType, ck_tile::bf16_t> &&
                           std::is_same_v<typename Problem::DDataType, ck_tile::bf16_t> &&
                           std::is_same_v<typename Problem::TopkWeightDataType, float> &&
                           S_::Block_M1 == 32 && S_::Block_N1 == 128 && S_::Block_K1 == 512 &&
                           S_::Warp_M0 == 16 && S_::Warp_N0 == 16 && S_::Warp_K0 == 32 &&
                           T_::PipeInterleave == true)
         {
             // return FlatmmSn_32x128x512_1x4x1_16x16x32_FP16{};
             return FlatmmSn_32x128x512_1x4x1_16x16x32_BF16_itl{};
         }
         else if constexpr(std::is_same_v<typename Problem::YDataType, ck_tile::fp16_t> &&
                           std::is_same_v<typename Problem::DDataType, ck_tile::fp16_t> &&
                           std::is_same_v<typename Problem::TopkWeightDataType, float> &&
                           S_::Block_M1 == 32 && S_::Block_N1 == 128 && S_::Block_K1 == 512 &&
                           S_::Warp_M0 == 16 && S_::Warp_N0 == 16 && S_::Warp_K0 == 32 &&
                           T_::PipeInterleave == true)
         {
             // return FlatmmSn_32x128x512_1x4x1_16x16x32_FP16{};
             return FlatmmSn_32x128x512_1x4x1_16x16x32_FP16_itl{};
         }
     }
 };
 } // namespace ck_tile