From f7ce8aa4f7ff9ff45f6e9be223ae89c9cb6d93c0 Mon Sep 17 00:00:00 2001 From: huizzhan Date: Mon, 11 Aug 2025 10:54:12 +0000 Subject: [PATCH] topk debug3 --- ...grouped_topk_pipeline_agmem_bgmem_creg.hpp | 102 ++++++--------- .../gemm_softmax_grouped_topk.cpp | 120 ++++++++++++------ .../gemm_softmax_grouped_topk.hpp | 8 +- .../grid_gemm_softmax_grouped_topk.hpp | 39 +++++- .../reference_gemm_softmax_grouped_topk.hpp | 59 +++++++++ 5 files changed, 218 insertions(+), 110 deletions(-) diff --git a/example/ck_tile/39_gemm_softmax_grouped_topk/block_gemm_softmax_grouped_topk_pipeline_agmem_bgmem_creg.hpp b/example/ck_tile/39_gemm_softmax_grouped_topk/block_gemm_softmax_grouped_topk_pipeline_agmem_bgmem_creg.hpp index 1a3d52b9bb..9c176bb4ab 100755 --- a/example/ck_tile/39_gemm_softmax_grouped_topk/block_gemm_softmax_grouped_topk_pipeline_agmem_bgmem_creg.hpp +++ b/example/ck_tile/39_gemm_softmax_grouped_topk/block_gemm_softmax_grouped_topk_pipeline_agmem_bgmem_creg.hpp @@ -38,8 +38,8 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg // for topk computing struct ArgmaxPacket { - WeightType arg; - IndexType value; + WeightType value; + IndexType arg; }; using BlockGemm = remove_cvref_t())>; @@ -208,13 +208,15 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg } #endif - template - CK_TILE_HOST_DEVICE void operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + // template + template + CK_TILE_HOST_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, const BDramBlockWindowTmp& b_dram_block_window_tmp, - ValueBlockTile& value_block_tile, - IndexBlockTile& index_block_tile, index_t num_loop, void* p_smem) const + // DebugBlockTile& debug_block_tile, + // ValueBlockTile& value_block_tile, + // IndexBlockTile& index_block_tile, { static_assert( std::is_same_v> && @@ -437,9 +439,8 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg auto p_compute = make_static_distributed_tensor(c_block_tile.get_tile_distribution()); - auto wgid = blockIdx.x + blockIdx.y * gridDim.x + gridDim.x * gridDim.y * blockIdx.z; - auto tid = (threadIdx.z * (blockDim.x * blockDim.y)) + (threadIdx.y * blockDim.x) + threadIdx.x; - int i = 0, j = 0; + auto debug_block_tile = + make_static_distributed_tensor(p_compute.get_tile_distribution()); constexpr auto p_spans = decltype(p_compute)::get_distributed_spans(); @@ -450,13 +451,8 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg constexpr auto i_j_idx = make_tuple(idx0, idx1); p_compute(i_j_idx) = exp(c_block_tile[i_j_idx] - m_local[i_idx]); - if (wgid == 0 && tid ==0 && i == 0) { - // if (i == 0 && j == 0) { - printf("============c_block_tile(i %d, j %d): %f===============\n", i, j, c_block_tile(i_j_idx)); - } - j++; + }); - i++; }); // rowsum for p_compute{i, j} @@ -466,21 +462,14 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg block_tile_reduce_sync(rowsum_p, f_sum); // softmax = p_compute{i, j} / rowsum_p - i = 0; - j = 0; sweep_tile_span(p_spans[number<0>{}], [&](auto idx0) { constexpr auto i_idx = make_tuple(idx0); sweep_tile_span(p_spans[number<1>{}], [&](auto idx1) { constexpr auto i_j_idx = make_tuple(idx0, idx1); p_compute(i_j_idx) = p_compute[i_j_idx] / rowsum_p[i_idx]; - if (wgid == 0 && tid ==0 && i == 0) { - // if (i == 0 && j == 0) { - printf("============p_compute(i %d, j %d): %f===============\n", i, j, p_compute(i_j_idx)); - } - j++; + // debug_block_tile(i_j_idx) = p_compute(i_j_idx); }); - i++; }); // apply topk for softmax output @@ -489,81 +478,66 @@ struct BlockGemmSoftmaxGroupedTopkPipelineAGmemBGmemCReg // argmax for topk const auto f_argmax = [](ArgmaxPacket e0, ArgmaxPacket e1) { - return e0.arg > e1.arg ? e0 : e1; + return e0.value > e1.value ? e0 : e1; }; - printf("==========================topk: %d====================================\n", topk); - - for(index_t i_k = 0; i_k < topk; i_k++) + for(index_t i_k = 0; i_k < 1; i_k++) { - printf("==========================i_k: %d====================================\n", i_k); - constexpr auto span_2d = decltype(p_compute)::get_distributed_spans(); + constexpr auto p_compute_spans = decltype(p_compute)::get_distributed_spans(); auto packet = [&]() { auto tmp = make_static_distributed_tensor(p_compute.get_tile_distribution()); - i = 0; - j = 0; - sweep_tile_span(span_2d[number<0>{}], [&](auto idx0) { - sweep_tile_span(span_2d[number<1>{}], [&](auto idx1) { + sweep_tile_span(p_compute_spans[number<0>{}], [&](auto idx0) { + sweep_tile_span(p_compute_spans[number<1>{}], [&](auto idx1) { const auto tile_idx = get_x_indices_from_distributed_indices( tmp.get_tile_distribution(), make_tuple(idx0, idx1)); constexpr auto i_j_idx = make_tuple(idx0, idx1); ArgmaxPacket t; - t.arg = x_tmp(i_j_idx); // !!! we reference p_compute here - t.value = tile_idx.at(number<1>{}); + t.value = x_tmp(i_j_idx); // !!! we reference p_compute here + t.arg = tile_idx.at(number<1>{}); tmp(i_j_idx) = t; - // if (wgid == 0 && tid ==0 && i == 0) { - // printf("=========p_compute(i %d, j %d)- t.arg: %f t.arg: %d=======\n", i, j, t.arg, t.value); - // } - j++; }); - i++; }); return tmp; }(); - i = 0; - j = 0; - sweep_tile_span(span_2d[number<0>{}], [&](auto idx0) { - sweep_tile_span(span_2d[number<1>{}], [&](auto idx1) { - constexpr auto i_j_idx = make_tuple(idx0, idx1); - if (wgid == 0 && tid ==0 && i == 0) { - printf("====packet(i %d, j %d)- t.arg: %f t.arg: %d====\n", i, j, packet(i_j_idx).arg, packet(i_j_idx).value); - } - j++; - }); - i++; - }); - auto argmax_init = ArgmaxPacket{-numeric::infinity(), 0}; auto r = block_tile_reduce(packet, sequence<1>{}, f_argmax, argmax_init); + block_tile_reduce_xor_sync(r, f_argmax); - // auto value_block_tile = make_static_distributed_tensor(dst_dist); - // auto index_block_tile = make_static_distributed_tensor(dst_dist); - - sweep_tile_span(span_2d[number<0>{}], [&](auto idx0) { - sweep_tile_span(span_2d[number<1>{}], [&](auto idx1) { + // constexpr auto value_spans = decltype(value_block_tile)::get_distributed_spans(); + + sweep_tile_span(p_compute_spans[number<0>{}], [&](auto idx0) { + constexpr auto i_idx = make_tuple(idx0); + sweep_tile_span(p_compute_spans[number<1>{}], [&](auto idx1) { + const auto tile_idx = get_x_indices_from_distributed_indices( + p_compute.get_tile_distribution(), make_tuple(idx0, idx1)); + // auto row_id = tile_idx.at(number<0>{}); + auto col_id = tile_idx.at(number<1>{}); constexpr auto i_j_idx = make_tuple(idx0, idx1); - ArgmaxPacket tmp = r(i_j_idx); - value_block_tile(i_j_idx) = tmp.arg; - index_block_tile(i_j_idx) = tmp.value; + ArgmaxPacket tmp = r(i_idx); + // debug_block_tile(i_j_idx) = (col_id == i_k) ? tmp.value: 0; + debug_block_tile(i_j_idx) = (col_id == i_k) ? tmp.arg: 0; + // value_block_tile(i_j_idx) = tmp.value; + // index_block_tile(i_j_idx) = tmp.arg; }); }); // update value - sweep_tile_span(span_2d[number<0>{}], [&](auto idx0) { - sweep_tile_span(span_2d[number<1>{}], [&](auto idx1) { + sweep_tile_span(p_compute_spans[number<0>{}], [&](auto idx0) { + sweep_tile_span(p_compute_spans[number<1>{}], [&](auto idx1) { const auto tile_idx = get_x_indices_from_distributed_indices( p_compute.get_tile_distribution(), make_tuple(idx0, idx1)); auto col_id = tile_idx.at(number<1>{}); constexpr auto i_j_idx = make_tuple(idx0, idx1); - x_tmp(i_j_idx) = (col_id == r(i_j_idx).value) ? -numeric::infinity() + x_tmp(i_j_idx) = (col_id == r(i_j_idx).arg) ? -numeric::infinity() : x_tmp(i_j_idx); }); }); } + return debug_block_tile; } }; diff --git a/example/ck_tile/39_gemm_softmax_grouped_topk/gemm_softmax_grouped_topk.cpp b/example/ck_tile/39_gemm_softmax_grouped_topk/gemm_softmax_grouped_topk.cpp index 6a607bceb8..918c32ad2d 100755 --- a/example/ck_tile/39_gemm_softmax_grouped_topk/gemm_softmax_grouped_topk.cpp +++ b/example/ck_tile/39_gemm_softmax_grouped_topk/gemm_softmax_grouped_topk.cpp @@ -74,7 +74,7 @@ int main(int argc, char* argv[]) ck_tile::index_t M = 3328; ck_tile::index_t N = 4096; ck_tile::index_t K = 4096; - ck_tile::index_t topk = 16; + ck_tile::index_t topk = 8; if(argc == 2) { @@ -143,8 +143,8 @@ int main(int argc, char* argv[]) const auto b_lengths = std::array{N, K}; const auto b_strides = std::array{Ldb, 1}; - // const auto c_lengths = std::array{M, N}; - // const auto c_strides = std::array{Ldc, 1}; + const auto debug_lengths = std::array{M, N}; + const auto debug_strides = std::array{N, 1}; const auto out_lengths = std::array{M, topk}; const auto out_strides = std::array{Ldout, 1}; @@ -152,7 +152,7 @@ int main(int argc, char* argv[]) // host verify ck_tile::HostTensor a_host(a_lengths, a_strides); ck_tile::HostTensor b_host(b_lengths, b_strides); - // ck_tile::HostTensor c_host_dev(c_lengths, c_strides); + ck_tile::HostTensor debug_host_dev(debug_lengths, debug_strides); ck_tile::HostTensor value_host_dev(out_lengths, out_strides); ck_tile::HostTensor index_host_dev(out_lengths, out_strides); @@ -184,7 +184,7 @@ int main(int argc, char* argv[]) ck_tile::DeviceMem a_buf(a_host.get_element_space_size_in_bytes()); ck_tile::DeviceMem b_buf(b_host.get_element_space_size_in_bytes()); - // ck_tile::DeviceMem c_buf(c_host_dev.get_element_space_size_in_bytes()); + ck_tile::DeviceMem debug_buf(debug_host_dev.get_element_space_size_in_bytes()); ck_tile::DeviceMem value_buf(value_host_dev.get_element_space_size_in_bytes()); ck_tile::DeviceMem index_buf(index_host_dev.get_element_space_size_in_bytes()); @@ -207,7 +207,7 @@ int main(int argc, char* argv[]) constexpr ck_tile::index_t kGemmKPerBlock = 16; #endif constexpr ck_tile::index_t kGemmNPerBlock = 256; - constexpr ck_tile::index_t kGemmTopKPerBlock = 16; + constexpr ck_tile::index_t kGemmTopKPerBlock = 8; ck_tile::index_t kGridSize = (M / kGemmMPerBlock) * (N / kGemmNPerBlock); @@ -241,6 +241,7 @@ int main(int argc, char* argv[]) 0, static_cast(a_buf.GetDeviceBuffer()), static_cast(b_buf.GetDeviceBuffer()), + static_cast(debug_buf.GetDeviceBuffer()), static_cast(value_buf.GetDeviceBuffer()), static_cast(index_buf.GetDeviceBuffer()), M, @@ -255,61 +256,96 @@ int main(int argc, char* argv[]) bool rtn = true; if(verification) { + ck_tile::HostTensor debug_ref({M, N}, {N, 1}); ck_tile::HostTensor value_ref(out_lengths, out_strides); ck_tile::HostTensor index_ref(out_lengths, out_strides); - reference_basic_gemm_softmax_grouped_topk( - a_host, b_host, value_ref, index_ref, topk); + // reference_basic_gemm_softmax_grouped_topk( + // a_host, b_host, value_ref, index_ref, topk); + debug_ref = reference_basic_gemm_softmax(a_host, b_host); + debug_buf.FromDevice(debug_host_dev.mData.data()); value_buf.FromDevice(value_host_dev.mData.data()); index_buf.FromDevice(index_host_dev.mData.data()); - // pass &= ck_tile::check_err(c_host_dev, c_host_ref); + // rtn &= ck_tile::check_err(debug_host_dev, debug_ref); + // for(std::size_t i = 0; i < debug_ref.size(); ++i) { + // const double o = *std::next(std::begin(debug_host_dev), i); + // const double r = *std::next(std::begin(debug_ref), i); + // std::cout << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl; + // } + // std::cout << "valid:" << (rtn ? "y" : "n") << std::endl; + const ck_tile::index_t tokens = M; - auto [rtol, atol] = get_elimit(""); + auto [rtol, atol] = get_elimit(""); for(int i_t = 0; i_t < tokens; i_t++) { auto s_begin = std::vector{static_cast(i_t), static_cast(0)}; auto s_end = std::vector{static_cast(i_t + 1), static_cast(topk)}; - auto s_value_host = value_host_dev.slice(s_begin, s_end); - auto s_value_ref = value_ref.slice(s_begin, s_end); - rtn &= ck_tile::check_err(s_value_host, - s_value_ref, + auto s_debug_host = debug_host_dev.slice(s_begin, s_end); + auto s_debug_ref = debug_ref.slice(s_begin, s_end); + rtn &= ck_tile::check_err(s_debug_host, + s_debug_ref, std::string("[") + std::to_string(i_t) + std::string("] Value Error:"), rtol, atol); - // for(std::size_t i = 0; i < s_value_ref.size(); ++i) { - // const double o = *std::next(std::begin(s_value_host), i); - // const double r = *std::next(std::begin(s_value_ref), i); - // std::cout << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl; - // } - auto s_index_host = index_host_dev.slice(s_begin, s_end); - auto s_index_ref = index_ref.slice(s_begin, s_end); - rtn &= ck_tile::check_err(s_index_host, - s_index_ref, - std::string("[") + std::to_string(i_t) + - std::string("] Index Error:"), - rtol, - atol); - // for(std::size_t i = 0; i < s_index_ref.size(); ++i) { - // const double o = *std::next(std::begin(s_index_host), i); - // const double r = *std::next(std::begin(s_index_ref), i); - // std::cout << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl; - // } + printf("row [%d]\n", i_t); + for(std::size_t i = 0; i < s_debug_ref.size(); ++i) { + // const double o = *std::next(std::begin(s_debug_host), i); + const double r = *std::next(std::begin(s_debug_ref), i); + printf("ref[%zu]:%f ", i, r); + // std::cout << i_t << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl; + } + printf("\n"); + for(std::size_t i = 0; i < s_debug_ref.size(); ++i) { + const double o = *std::next(std::begin(s_debug_host), i); + // const double r = *std::next(std::begin(s_debug_ref), i); + printf("out[%zu]:%f ", i, o); + // std::cout << i_t << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl; + } + printf("\n"); } std::cout << "valid:" << (rtn ? "y" : "n") << std::endl; } + // const ck_tile::index_t tokens = M; + // auto [rtol, atol] = get_elimit(""); + // for(int i_t = 0; i_t < tokens; i_t++) + // { + // auto s_begin = std::vector{static_cast(i_t), static_cast(0)}; + // auto s_end = + // std::vector{static_cast(i_t + 1), static_cast(topk)}; + // auto s_value_host = value_host_dev.slice(s_begin, s_end); + // auto s_value_ref = value_ref.slice(s_begin, s_end); + // rtn &= ck_tile::check_err(s_value_host, + // s_value_ref, + // std::string("[") + std::to_string(i_t) + + // std::string("] Value Error:"), + // rtol, + // atol); + // // for(std::size_t i = 0; i < s_value_ref.size(); ++i) { + // // const double o = *std::next(std::begin(s_value_host), i); + // // const double r = *std::next(std::begin(s_value_ref), i); + // // std::cout << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl; + // // } + // auto s_index_host = index_host_dev.slice(s_begin, s_end); + // auto s_index_ref = index_ref.slice(s_begin, s_end); + // rtn &= ck_tile::check_err(s_index_host, + // s_index_ref, + // std::string("[") + std::to_string(i_t) + + // std::string("] Index Error:"), + // rtol, + // atol); + // // for(std::size_t i = 0; i < s_index_ref.size(); ++i) { + // // const double o = *std::next(std::begin(s_index_host), i); + // // const double r = *std::next(std::begin(s_index_ref), i); + // // std::cout << " out[" << i << "] != ref[" << i << "]: " << o << " != " << r << std::endl; + // // } + // } + // std::cout << "valid:" << (rtn ? "y" : "n") << std::endl; + // } - std::size_t flop = std::size_t(2) * M * N * K; - std::size_t num_btype = - sizeof(ADataType) * M * K + sizeof(BDataType) * K * N + sizeof(CDataType) * M * N; - - float tflops = static_cast(flop) / 1.E9 / ave_time; - - float gb_per_sec = num_btype / 1.E6 / ave_time; - - std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s" - << std::endl; + std::cout << "Perf: " << ave_time << " ms, " << std::endl; return rtn; + // return !rtn; } diff --git a/example/ck_tile/39_gemm_softmax_grouped_topk/gemm_softmax_grouped_topk.hpp b/example/ck_tile/39_gemm_softmax_grouped_topk/gemm_softmax_grouped_topk.hpp index 9f5d5bf436..cbfc8a433f 100755 --- a/example/ck_tile/39_gemm_softmax_grouped_topk/gemm_softmax_grouped_topk.hpp +++ b/example/ck_tile/39_gemm_softmax_grouped_topk/gemm_softmax_grouped_topk.hpp @@ -179,6 +179,7 @@ struct Gemm CK_TILE_DEVICE void operator()(const ADataType* p_a, const BDataType* p_b, + WeightType* p_debug, WeightType* p_value, IndexType* p_index, const index_t M, @@ -200,6 +201,11 @@ struct Gemm p_b, make_tuple(N, K), make_tuple(Ldb, 1), number{}, number<1>{}); }(); + const auto debug_dram = [&] { + return make_naive_tensor_view( + p_debug, make_tuple(M, N), make_tuple(N, 1), number{}, number<1>{}); + }(); + const auto value_dram = [&] { return make_naive_tensor_view( p_value, make_tuple(M, topK), make_tuple(Ldout, 1), number{}, number<1>{}); @@ -210,7 +216,7 @@ struct Gemm p_index, make_tuple(M, topK), make_tuple(Ldout, 1), number{}, number<1>{}); }(); - GridGemm{}(a_dram, b_dram, value_dram, index_dram, c_element_func); + GridGemm{}(a_dram, b_dram, debug_dram, value_dram, index_dram, c_element_func); } }; diff --git a/example/ck_tile/39_gemm_softmax_grouped_topk/grid_gemm_softmax_grouped_topk.hpp b/example/ck_tile/39_gemm_softmax_grouped_topk/grid_gemm_softmax_grouped_topk.hpp index b02ef993d0..a744abf954 100755 --- a/example/ck_tile/39_gemm_softmax_grouped_topk/grid_gemm_softmax_grouped_topk.hpp +++ b/example/ck_tile/39_gemm_softmax_grouped_topk/grid_gemm_softmax_grouped_topk.hpp @@ -23,9 +23,10 @@ struct GridGemm static constexpr auto topk = Policy::kTopKPerBlock; static constexpr auto kBlockSize = Policy::kBlockSize; - template + template CK_TILE_DEVICE void operator()(const AGridTensorView& a_grid, const BGridTensorView& b_grid, + DebugGridTensorView& debug_grid, ValueGridTensorView& value_grid, IndexGridTensorView& index_grid, const CElementFunction& c_element_func) const @@ -70,7 +71,7 @@ struct GridGemm // constexpr index_t kMPerBlock = Problem::BlockGemmShape::kM; // constexpr index_t kKPerBlock = Problem::BlockGemmShape::kK; - constexpr index_t K1 = 16 / sizeof( WeightType); + constexpr index_t K1 = 16 / sizeof(WeightType); constexpr index_t K0 = topk / K1; constexpr index_t M2 = get_warp_size() / K0; // coalesce reading for each blocks @@ -109,9 +110,39 @@ struct GridGemm tile_elementwise_inout([](auto& value) { value = 0; }, value_block_tile); tile_elementwise_inout([](auto& index) { index = 0; }, index_block_tile); - block_gemm_pipeline(a_block_window, b_block_window, value_block_tile, index_block_tile, K / kKPerBlock, p_smem_char); + // constexpr index_t debugK1 = 16 / sizeof(WeightType); + // constexpr index_t debugK0 = kNPerBlock / debugK1; + // constexpr index_t debugM2 = get_warp_size() / debugK0; + // // coalesce reading for each blocks + // constexpr index_t debugM1 = kBlockSize / get_warp_size(); + // constexpr index_t debugM0 = kMPerBlock / (debugM2 * debugM1); + + auto debug_window = make_tile_window( + debug_grid, make_tuple(number{}, number{}), {iM, iN}); + + // auto debug_window = make_tile_window( + // debug_grid, make_tuple(number{}, number{}), {iM, iN}, + // make_static_tile_distribution( + // tile_distribution_encoding, + // tuple, sequence>, + // tuple, sequence<1, 2>>, + // tuple, sequence<2, 0>>, + // sequence<1, 2>, + // sequence<0, 1>>{})); + + // using DebugBlockTileDistr = decltype(debug_window.get_tile_distribution()); + // using DebugBlockTile = decltype(make_static_distributed_tensor(DebugBlockTileDistr{})); + // DebugBlockTile debug_block_tile; + // tile_elementwise_inout([](auto& debug) { debug = 0; }, debug_block_tile); + + // block_gemm_pipeline(a_block_window, b_block_window, debug_block_tile, value_block_tile, index_block_tile, K / kKPerBlock, p_smem_char); + const auto debug_block_tile = block_gemm_pipeline(a_block_window, b_block_window, K / kKPerBlock, p_smem_char); // cast DataType and apply CElementFunction + const auto debug_cast_block_tile = tile_elementwise_in( + [&](const auto& debug) { return c_element_func(type_convert(debug)); }, + debug_block_tile); + const auto value_cast_block_tile = tile_elementwise_in( [&](const auto& value) { return c_element_func(type_convert(value)); }, value_block_tile); @@ -120,6 +151,8 @@ struct GridGemm [&](const auto& index) { return c_element_func(type_convert(index)); }, index_block_tile); + + store_tile(debug_window, debug_cast_block_tile); store_tile(value_window, value_cast_block_tile); store_tile(index_window, index_cast_block_tile); } diff --git a/example/ck_tile/39_gemm_softmax_grouped_topk/reference_gemm_softmax_grouped_topk.hpp b/example/ck_tile/39_gemm_softmax_grouped_topk/reference_gemm_softmax_grouped_topk.hpp index 2e78644d7d..58545d9fa7 100755 --- a/example/ck_tile/39_gemm_softmax_grouped_topk/reference_gemm_softmax_grouped_topk.hpp +++ b/example/ck_tile/39_gemm_softmax_grouped_topk/reference_gemm_softmax_grouped_topk.hpp @@ -66,3 +66,62 @@ void reference_basic_gemm_softmax_grouped_topk(const ck_tile::HostTensor +auto reference_basic_gemm_softmax(const ck_tile::HostTensor& a_m_k, + const ck_tile::HostTensor& b_n_k) +{ + const int M = a_m_k.mDesc.get_lengths()[0]; + const int N = b_n_k.mDesc.get_lengths()[0]; + const int K = b_n_k.mDesc.get_lengths()[1]; + ck_tile::HostTensor c_m_n({M, N}, {N, 1}); + + auto f = [&](auto m) { + for(int n = 0; n < N; ++n) + { + AccDataType v_acc = 0; + + for(int k = 0; k < K; ++k) + { + ADataType v_a = a_m_k(m, k); + BDataType v_b = b_n_k(n, k); + v_acc += ck_tile::type_convert(v_a) * + ck_tile::type_convert(v_b); + } + + c_m_n(m, n) = ck_tile::type_convert(v_acc); + } + // reference softmax + AccDataType v_max = std::numeric_limits::lowest(); + + // max + for(int n = 0; n < N; ++n) + { + const AccDataType v_c = c_m_n(m, n); + v_max = v_max < v_c ? v_c : v_max; + } + + AccDataType v_exp_sum = 0; + + // sum + for(int n = 0; n < N; ++n) + { + const AccDataType v_c = c_m_n(m, n); + v_exp_sum += ck_tile::exp(v_c - v_max); + } + + // elementwise + for(int n = 0; n < N; ++n) + { + const AccDataType v_c = c_m_n(m, n); + c_m_n(m, n) = ck_tile::exp(v_c - v_max) / v_exp_sum; + } + }; + + ck_tile::make_ParallelTensorFunctor(f, c_m_n.mDesc.get_lengths()[0])( + std::thread::hardware_concurrency()); + + // reference_topk(c_m_n, y_values, y_indices, topk); + // reference_grouped_topk(c_m_n, y_values, y_indices, topk, num_expert_group, topk_group, dim, largest, sorted); + return c_m_n; +}