From d4a0a8eec2b6f09a3a4a60187cd14cbdd9b1eda1 Mon Sep 17 00:00:00 2001 From: letaoqin Date: Fri, 13 Dec 2024 09:54:16 +0000 Subject: [PATCH] add gelu and weight --- example/ck_tile/17_fused_moe_general/main.cpp | 20 ++++++-- include/ck_tile/host/host_tensor.hpp | 1 + .../kernel/fused_moegemm_general_kernel.hpp | 28 ++++++++--- .../fused_moegemm_pipeline_general.hpp | 47 ++++++++++++------- 4 files changed, 69 insertions(+), 27 deletions(-) diff --git a/example/ck_tile/17_fused_moe_general/main.cpp b/example/ck_tile/17_fused_moe_general/main.cpp index a8b795698b..bb9391b7db 100644 --- a/example/ck_tile/17_fused_moe_general/main.cpp +++ b/example/ck_tile/17_fused_moe_general/main.cpp @@ -264,6 +264,7 @@ bool run(const ck_tile::ArgParser& arg_parser) {(max_num_tokens_padded + block_m - 1) / block_m}); ck_tile::HostTensor num_sorted_tiles_host({1}); + sorted_token_ids_host.SetValue(max_num_tokens_padded); if(init == 0) { ck_tile::FillStepRange{-.5f, .5f, 0.01f}(a_host); @@ -280,9 +281,6 @@ bool run(const ck_tile::ArgParser& arg_parser) ck_tile::FillUniformDistribution{-.5f, .5f, seed, true}(a_host); ck_tile::FillUniformDistribution{-.5f, .5f, seed, true}(g_host); ck_tile::FillUniformDistribution{-.5f, .5f, seed, true}(d_host); - // ck_tile::FillConstant{1}(a_host); - // ck_tile::FillConstant{1}(g_host); - // ck_tile::FillConstant{1}(d_host); ck_tile::FillUniformDistribution{-.5f, .5f, seed, true}(sa_host); ck_tile::FillUniformDistribution{-.5f, .5f, seed, true}(sg_host); ck_tile::FillUniformDistribution{-.5f, .5f, seed, true}(sd_host); @@ -301,6 +299,18 @@ bool run(const ck_tile::ArgParser& arg_parser) ck_tile::FillNormalDistribution{0.f, 1.f, seed, true}(sy_host); ck_tile::FillNormalDistribution{0.f, 1.f, seed, true}(topk_weight_host); } + else if(init == 3) + { + ck_tile::FillConstant{1}(a_host); + ck_tile::FillConstant{1}(g_host); + ck_tile::FillConstant{1}(d_host); + ck_tile::FillUniformDistribution{-.5f, .5f, seed, true}(sa_host); + ck_tile::FillUniformDistribution{-.5f, .5f, seed, true}(sg_host); + ck_tile::FillUniformDistribution{-.5f, .5f, seed, true}(sd_host); + ck_tile::FillUniformDistribution{-.5f, .5f, seed, true}(sy_host); + ck_tile::FillUniformDistribution{0.0f, 1.0f, seed, true}( + topk_weight_host); + } // permute weight // ck_tile::HostTensor g_perm_host = shuffle_moe_weight(g_host, prec_w, 1); @@ -393,7 +403,7 @@ bool run(const ck_tile::ArgParser& arg_parser) std::cout << sorted_expert_ids_host << std::endl; // std::cout << topk_weight_host << std::endl; - // std::cout << sorted_weight_host << std::endl; + std::cout << sorted_weight_host << std::endl; // done, preparing GPU buffer ck_tile::DeviceMem a_buf(a_host); ck_tile::DeviceMem g_perm_buf(g_host); @@ -490,7 +500,7 @@ bool run(const ck_tile::ArgParser& arg_parser) auto o_dev = o_buf.ToHost(); auto c_dev = c_buf.ToHost(); std::cout << std::endl; - std::cout << o_dev << std::endl; + // std::cout << o_dev << std::endl; // std::cout << c_dev << std::endl; // int count = 0; // std::cout << "["; diff --git a/include/ck_tile/host/host_tensor.hpp b/include/ck_tile/host/host_tensor.hpp index 30af0351ed..aef6fc8fa3 100644 --- a/include/ck_tile/host/host_tensor.hpp +++ b/include/ck_tile/host/host_tensor.hpp @@ -349,6 +349,7 @@ struct HostTensor // void SetZero() { ck_tile::ranges::fill(mData, 0); } void SetZero() { std::fill(mData.begin(), mData.end(), 0); } + void SetValue(int value) { std::fill(mData.begin(), mData.end(), value); } template void ForEach_impl(F&& f, std::vector& idx, size_t rank) diff --git a/include/ck_tile/ops/fused_moe/kernel/fused_moegemm_general_kernel.hpp b/include/ck_tile/ops/fused_moe/kernel/fused_moegemm_general_kernel.hpp index 94f40d0613..b533c1e4e8 100644 --- a/include/ck_tile/ops/fused_moe/kernel/fused_moegemm_general_kernel.hpp +++ b/include/ck_tile/ops/fused_moe/kernel/fused_moegemm_general_kernel.hpp @@ -252,12 +252,12 @@ struct FusedMoeGemmGlKernel index_t idx_n0 = __builtin_amdgcn_readfirstlane(intermediate_tile_id * BlockShape::Block_N0); - const auto a_coord = Pipeline::GetACoord(); // 2d thread offset, [i_row, i_col] - const auto sorted_token_id = a_coord[number<0>{}] + idx_m0; // start block_m - // position + // const auto a_coord = Pipeline::GetACoord(); // 2d thread offset, [i_row, i_col] + // const auto sorted_token_id = a_coord[number<0>{}] + idx_m0; // start block_m + // // position - auto topk_weight = - reinterpret_cast(kargs.sorted_weight_ptr)[sorted_token_id]; + // auto topk_weight = + // reinterpret_cast(kargs.sorted_weight_ptr)[sorted_token_id]; const index_t* sorted_token_ids_ptr = reinterpret_cast(kargs.sorted_token_ids_ptr); @@ -374,12 +374,28 @@ struct FusedMoeGemmGlKernel return o_window_; }(); + const auto w_window = [&]() { + const TopkWeightDataType* w_ptr = reinterpret_cast(kargs.sorted_weight_ptr); + const auto w_view_ = make_naive_tensor_view( + w_ptr, + make_tuple(kargs.max_num_tokens_padded), + make_tuple(1), + number<1>{}, + number<1>{}); + + const auto w_window_ = make_tile_window( + w_view_, + make_tuple(number{}), + {idx_m0}); + return w_window_; + }(); + // do compute yeah Pipeline{}(a_window, g_window, d_window, + w_window, o_window, - topk_weight, smem, kargs.hidden_size, kargs.intermediate_size, diff --git a/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_general.hpp b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_general.hpp index 00ab1c6834..eaa54e5934 100644 --- a/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_general.hpp +++ b/include/ck_tile/ops/fused_moe/pipeline/fused_moegemm_pipeline_general.hpp @@ -89,14 +89,6 @@ struct FusedMoeGemmPipeline_General // return Policy::template GetSmemSize(); } - // this is the thread-offset along row/col - CK_TILE_HOST_DEVICE static auto GetACoord() - { - constexpr auto a_dist = Policy::template MakeGlobalTileDistribution_A(); - const auto a_coord = a_dist.calculate_index(); - return a_coord; - } - template CK_TILE_HOST_DEVICE static void PrintMem(T& tensor, const char* pstr, unsigned int threadid = 0, unsigned int blockid = 0) @@ -129,20 +121,21 @@ struct FusedMoeGemmPipeline_General typename GWindow, typename DWindow, typename OWindow, - typename CWindow> + typename CWindow, + typename WWindow> CK_TILE_DEVICE auto operator()(const AWindow& a_window_, const GWindow& g_window_, const DWindow& d_window_, + const WWindow& w_window_, OWindow& o_window_, - TopkWeightDataType topk_weight, CK_TILE_LDS_ADDR void* smem, index_t hidden_size, index_t /*intermediate_size*/, CWindow& c_window_) { - ignore = topk_weight; ignore = c_window_; ignore = hidden_size; + ignore = w_window_; CK_TILE_LDS_ADDR ADataType* smem_0 = reinterpret_cast(smem); CK_TILE_LDS_ADDR GDataType* smem_1 = reinterpret_cast( smem_0 + GetSmemSizeA() / sizeof(ADataType)); @@ -233,8 +226,8 @@ struct FusedMoeGemmPipeline_General PrintMem(s_acc, "S", 0); #endif // relu - // const auto activation = ck_tile::element_wise::Gelu{}; - // tile_elementwise_inout(activation, s_acc, s_acc); + const auto activation = ck_tile::element_wise::Gelu{}; + tile_elementwise_inout(activation, s_acc, s_acc); // cast data to YDataType auto y_pre = cast_tile(s_acc); @@ -260,6 +253,28 @@ struct FusedMoeGemmPipeline_General constexpr auto gemm_1 = Policy::template GetBlockGemm1(); using OaccBlockTileType = decltype(gemm_1.MakeCBlockTile()); auto o_acc = OaccBlockTileType{}; + + constexpr auto w_dstr = + make_static_tile_distribution(detail::make_reduce_tile_distribution_encoding( + s_acc.get_tile_distribution().get_static_tile_distribution_encoding(), sequence<1>{})); + auto w_global_to_dram_window = make_tile_window( + w_window_.get_bottom_tensor_view(), + make_tuple(number{}), + w_window_.get_window_origin(), + w_dstr); + auto w = load_tile(w_global_to_dram_window); + float weight = type_convert(w.get_thread_buffer()[0]); +#if 0 + constexpr index_t w_buffer_size = decltype(w)::get_thread_buffer_size(); + if(threadIdx.x == 1 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) + { + for(int i = 0; i < w_buffer_size; i++) + { + printf("\n len: %d, w[%d]: %f weight: %f", w_buffer_size, i, type_convert(w.get_thread_buffer()[i]), topk_weight); + } + } +#endif + ignore = w; // y data auto bridge_llds_win = make_tile_window(bridge_lds_view, @@ -308,7 +323,7 @@ struct FusedMoeGemmPipeline_General Policy::template MakeGlobalTileDistribution_O()); auto save_o = [&]() { - if(blockIdx.x == 0 && (blockIdx.y == 0 || blockIdx.y == 1) && blockIdx.z == 0) + //if(blockIdx.x == 0 && (blockIdx.y == 0 || blockIdx.y == 1) && blockIdx.z == 0) { if(threadIdx.x < 64) { @@ -352,8 +367,8 @@ struct FusedMoeGemmPipeline_General gemm_1(o_acc, y, d); // block_sync_lds(); - // tile_elementwise_inout( - // [&topk_weight](auto& x) { x = x * type_convert(topk_weight); }, o_acc); + tile_elementwise_inout( + [&weight](auto& x) { x = x * type_convert(weight); }, o_acc); auto o = cast_tile(o_acc); store_tile(o_alds_win, o); block_sync_lds();