diff --git a/example/ck_tile/39_gemm_softmax/block_gemm_pipeline_agmem_bgmem_creg.hpp b/example/ck_tile/39_gemm_softmax/block_gemm_pipeline_agmem_bgmem_creg.hpp index d975befe2b..b56169373e 100755 --- a/example/ck_tile/39_gemm_softmax/block_gemm_pipeline_agmem_bgmem_creg.hpp +++ b/example/ck_tile/39_gemm_softmax/block_gemm_pipeline_agmem_bgmem_creg.hpp @@ -23,6 +23,8 @@ struct BlockGemmSoftmaxPipelineAGmemBGmemCReg using ADataType = remove_cvref_t; using BDataType = remove_cvref_t; using CDataType = remove_cvref_t; + using WeightType = remove_cvref_t; + using IndexType = remove_cvref_t; using ComputeDataType = float; using BlockGemmShape = remove_cvref_t; @@ -31,6 +33,14 @@ struct BlockGemmSoftmaxPipelineAGmemBGmemCReg static constexpr index_t kMPerBlock = BlockGemmShape::kM; static constexpr index_t kNPerBlock = BlockGemmShape::kN; static constexpr index_t kKPerBlock = BlockGemmShape::kK; + static constexpr index_t topk = BlockGemmShape::kTopK; + + // for topk computing + struct ArgmaxPacket + { + WeightType arg; + IndexType value; + }; using BlockGemm = remove_cvref_t())>; @@ -198,9 +208,11 @@ struct BlockGemmSoftmaxPipelineAGmemBGmemCReg } #endif - template - CK_TILE_HOST_DEVICE auto operator()(const ADramBlockWindowTmp& a_dram_block_window_tmp, + template + CK_TILE_HOST_DEVICE void 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 { @@ -453,11 +465,70 @@ struct BlockGemmSoftmaxPipelineAGmemBGmemCReg p_compute(i_j_idx) = p_compute[i_j_idx] / rowsum_p[i_idx]; }); }); - // CDramBlockWindowTmp c_dram_block_window_tmp = c_dram_block_window; + + // apply topk for softmax output + auto x_tmp = p_compute; + // constexpr auto dst_dist = BlockGemmPipelineAGmemBGmemCRegDefaultPolicy::MakeOutputDistribution(); - // store_tile(c_dram_block_window_tmp, type_convert(p_compute)); + // argmax for topk + const auto f_argmax = [](ArgmaxPacket e0, ArgmaxPacket e1) { + return e0.arg > e1.arg ? e0 : e1; + }; - return p_compute; + for(index_t i_k = 0; i_k < topk; i_k++) + { + constexpr auto span_2d = decltype(p_compute)::get_distributed_spans(); + auto packet = [&]() { + auto tmp = make_static_distributed_tensor(p_compute.get_tile_distribution()); + + sweep_tile_span(span_2d[number<0>{}], [&](auto idx0) { + sweep_tile_span(span_2d[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>{}); + tmp(i_j_idx) = t; + }); + }); + return tmp; + }(); + + 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); + + // Initialize value_block_tile and index_block_tile + tile_elementwise_inout([](auto& value) { value = 0; }, value_block_tile); + tile_elementwise_inout([](auto& index) { index = 0; }, index_block_tile); + + 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); + ArgmaxPacket tmp = r(i_j_idx); + value_block_tile(i_j_idx) = tmp.arg; + index_block_tile(i_j_idx) = tmp.value; + }); + }); + + // update value + sweep_tile_span(span_2d[number<0>{}], [&](auto idx0) { + sweep_tile_span(span_2d[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); + }); + }); + } } }; diff --git a/example/ck_tile/39_gemm_softmax/block_gemm_pipeline_agmem_bgmem_creg_default_policy.hpp b/example/ck_tile/39_gemm_softmax/block_gemm_pipeline_agmem_bgmem_creg_default_policy.hpp index 30bdcd679f..9958aaceb7 100755 --- a/example/ck_tile/39_gemm_softmax/block_gemm_pipeline_agmem_bgmem_creg_default_policy.hpp +++ b/example/ck_tile/39_gemm_softmax/block_gemm_pipeline_agmem_bgmem_creg_default_policy.hpp @@ -261,6 +261,39 @@ struct BlockGemmPipelineAGmemBGmemCRegDefaultPolicy sequence<0, 1>>{}); } + // template + // CK_TILE_HOST_DEVICE static constexpr auto MakeOutputDistribution() + // { + // using WeightType = remove_cvref_t; + + // constexpr index_t kBlockSize = Problem::kBlockSize; + + // constexpr index_t kMPerBlock = Problem::BlockGemmShape::kM; + // constexpr index_t topk = Problem::BlockGemmShape::kTopK; + + // 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 + // constexpr index_t M1 = kBlockSize / get_warp_size(); + // constexpr index_t M0 = kMPerBlock / (M2 * M1); + + // // return make_static_tile_distribution( + // // tile_distribution_encoding, + // // tuple, sequence>, + // // tuple, sequence<1, 2>>, + // // tuple, sequence<2, 0>>, + // // sequence<1, 2>, + // // sequence<0, 1>>{}); + // return make_static_tile_distribution( + // tile_distribution_encoding, + // tuple, sequence<4, 4>>, + // tuple, sequence<1, 2>>, + // tuple, sequence<2, 0>>, + // sequence<1, 2>, + // sequence<0, 1>>{}); + // } + #if defined(ENABLE_INSTRUCTION_SCH) static constexpr auto I0 = number<0>{}; static constexpr auto I1 = number<1>{}; diff --git a/example/ck_tile/39_gemm_softmax/gemm.hpp b/example/ck_tile/39_gemm_softmax/gemm.hpp index 0667eaba12..a58533eefb 100755 --- a/example/ck_tile/39_gemm_softmax/gemm.hpp +++ b/example/ck_tile/39_gemm_softmax/gemm.hpp @@ -18,6 +18,8 @@ template struct GridGemmProblem { @@ -25,21 +27,26 @@ struct GridGemmProblem using BDataType = BDataType_; using AccDataType = AccDataType_; using CDataType = CDataType_; + using WeightType = WeightType_; + using IndexType = IndexType_; using CElementFunction = CElementFunction_; }; -template +template struct TileGemmShape { static constexpr index_t kM = kMPerTile; static constexpr index_t kN = kNPerTile; static constexpr index_t kK = kKPerTile; + static constexpr index_t kTopK = TopkKPerTile; }; template struct BlockGemmPipelineProblem @@ -47,6 +54,8 @@ struct BlockGemmPipelineProblem using ADataType = remove_cvref_t; using BDataType = remove_cvref_t; using CDataType = remove_cvref_t; + using WeightType = remove_cvref_t; + using IndexType = remove_cvref_t; using BlockGemmShape = remove_cvref_t; static constexpr index_t kBlockSize = kBlockSize_; @@ -57,18 +66,21 @@ template + index_t kKPerBlock_, + index_t kTopKPerBlock_> struct Gemm { using GridGemmProblem = - GridGemmProblem; + GridGemmProblem; struct GridGemmPolicy { @@ -76,6 +88,7 @@ struct Gemm static constexpr index_t kMPerBlock = kMPerBlock_; static constexpr index_t kNPerBlock = kNPerBlock_; static constexpr index_t kKPerBlock = kKPerBlock_; + static constexpr index_t kTopKPerBlock = kTopKPerBlock_; template CK_TILE_HOST_DEVICE static constexpr auto MakeBlock2TileMap(index_t M0, index_t N0) @@ -154,8 +167,10 @@ struct Gemm BlockGemmPipelineProblem>; + TileGemmShape>; return BlockGemmSoftmaxPipelineAGmemBGmemCReg{}; } }; @@ -164,13 +179,15 @@ struct Gemm CK_TILE_DEVICE void operator()(const ADataType* p_a, const BDataType* p_b, - CDataType* p_c, + WeightType* p_value, + IndexType* p_index, const index_t M, const index_t N, const index_t K, + const index_t topK, const index_t Lda, const index_t Ldb, - const index_t Ldc, + const index_t Ldout, const CElementFunction& c_element_func) const { const auto a_dram = [&] { @@ -183,12 +200,17 @@ struct Gemm p_b, make_tuple(N, K), make_tuple(Ldb, 1), number{}, number<1>{}); }(); - const auto c_dram = [&] { + const auto value_dram = [&] { return make_naive_tensor_view( - p_c, make_tuple(M, N), make_tuple(Ldc, 1), number{}, number<1>{}); + p_value, make_tuple(M, topK), make_tuple(Ldout, 1), number{}, number<1>{}); }(); - GridGemm{}(a_dram, b_dram, c_dram, c_element_func); + const auto index_dram = [&] { + return make_naive_tensor_view( + 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); } }; diff --git a/example/ck_tile/39_gemm_softmax/gemm_softmax.cpp b/example/ck_tile/39_gemm_softmax/gemm_softmax.cpp index 4215ad107c..2adbb907f8 100755 --- a/example/ck_tile/39_gemm_softmax/gemm_softmax.cpp +++ b/example/ck_tile/39_gemm_softmax/gemm_softmax.cpp @@ -26,28 +26,66 @@ struct CElementFunction } }; +// different threshold for different dtype +template +auto get_elimit(std::string /*init_method*/) +{ + double rtol = 1e-3; + double atol = 1e-3; + return ck_tile::make_tuple(rtol, atol); +} + +template <> +auto get_elimit(std::string /*init_method*/) +{ + double rtol = 1e-2; + double atol = 1e-2; + return ck_tile::make_tuple(rtol, atol); +} + +template <> +auto get_elimit(std::string init_method) +{ + if(init_method == "ui" || init_method == "ni") + { + unsigned max_rounding_point_distance = 0; + double atol = 2e-3; + return ck_tile::make_tuple(max_rounding_point_distance, atol); + } + else + { + unsigned max_rounding_point_distance = 1; + double atol = 0.0625; + return ck_tile::make_tuple(max_rounding_point_distance, atol); + } +} + int main(int argc, char* argv[]) { using ADataType = ck_tile::half_t; using BDataType = ck_tile::half_t; using AccDataType = float; using CDataType = ck_tile::half_t; + using WeightType = float; + using IndexType = ck_tile::index_t; ck_tile::index_t verification = 0; ck_tile::index_t M = 3328; ck_tile::index_t N = 4096; ck_tile::index_t K = 4096; + ck_tile::index_t topk = 16; if(argc == 2) { verification = std::stoi(argv[1]); } - if(argc == 5) + if(argc == 6) { verification = std::stoi(argv[1]); M = std::stoi(argv[2]); N = std::stoi(argv[3]); K = std::stoi(argv[4]); + topk = std::stoi(argv[5]); } #if defined(KERNEL_A) @@ -95,7 +133,8 @@ int main(int argc, char* argv[]) const ck_tile::index_t Lda = K; const ck_tile::index_t Ldb = K; - const ck_tile::index_t Ldc = N; + // const ck_tile::index_t Ldc = N; + const ck_tile::index_t Ldout = topk; const auto a_lengths = std::array{M, K}; const auto a_strides = std::array{Lda, 1}; @@ -103,20 +142,27 @@ 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 c_lengths = std::array{M, N}; + // const auto c_strides = std::array{Ldc, 1}; + + const auto out_lengths = std::array{M, topk}; + const auto out_strides = std::array{Ldout, 1}; // 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 c_host_dev(c_lengths, c_strides); + ck_tile::HostTensor value_host_dev(out_lengths, out_strides); + ck_tile::HostTensor index_host_dev(out_lengths, out_strides); ck_tile::FillUniformDistributionIntegerValue{-5.f, 5.f}(a_host); ck_tile::FillUniformDistributionIntegerValue{-5.f, 5.f}(b_host); 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 c_buf(c_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()); a_buf.ToDevice(a_host.mData.data()); b_buf.ToDevice(b_host.mData.data()); @@ -124,7 +170,8 @@ int main(int argc, char* argv[]) // Alignment constexpr ck_tile::index_t kAAlignment = 8; constexpr ck_tile::index_t kBAlignment = 8; - constexpr ck_tile::index_t kCAlignment = 8; + // constexpr ck_tile::index_t kCAlignment = 8; + constexpr ck_tile::index_t kOutAlignment = 8; constexpr ck_tile::index_t kBlockSize = 256; @@ -136,6 +183,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; ck_tile::index_t kGridSize = (M / kGemmMPerBlock) * (N / kGemmNPerBlock); @@ -149,14 +197,17 @@ int main(int argc, char* argv[]) BDataType, AccDataType, CDataType, + WeightType, + IndexType, CElementFunction, kAAlignment, kBAlignment, - kCAlignment, + kOutAlignment, kBlockSize, kGemmMPerBlock, kGemmNPerBlock, - kGemmKPerBlock>; + kGemmKPerBlock, + kGemmTopKPerBlock>; float ave_time = ck_tile::launch_kernel(ck_tile::stream_config{nullptr, true, 0, 5, 1000}, ck_tile::make_kernel( @@ -166,25 +217,61 @@ int main(int argc, char* argv[]) 0, static_cast(a_buf.GetDeviceBuffer()), static_cast(b_buf.GetDeviceBuffer()), - static_cast(c_buf.GetDeviceBuffer()), + static_cast(value_buf.GetDeviceBuffer()), + static_cast(index_buf.GetDeviceBuffer()), M, N, K, + topk, Lda, Ldb, - Ldc, + Ldout, CElementFunction{})); - auto pass = true; - + // auto pass = true; + bool rtn = true; if(verification) { // reference gemm - ck_tile::HostTensor c_host_ref(c_lengths, c_strides); - reference_basic_gemm_softmax( - a_host, b_host, c_host_ref); - c_buf.FromDevice(c_host_dev.mData.data()); - pass &= ck_tile::check_err(c_host_dev, c_host_ref); - std::cout << "valid:" << (pass ? "y" : "n") << std::endl; + // ck_tile::HostTensor c_host_ref(c_lengths, c_strides); + // reference_basic_gemm_softmax( + // a_host, b_host, c_host_ref); + // c_buf.FromDevice(c_host_dev.mData.data()); + // pass &= ck_tile::check_err(c_host_dev, c_host_ref); + // std::cout << "valid:" << (pass ? "y" : "n") << std::endl; + + ck_tile::HostTensor value_ref(out_lengths, out_strides); + ck_tile::HostTensor index_ref(out_lengths, out_strides); + reference_basic_gemm_softmax_topk( + a_host, b_host, value_ref, index_ref, topk); + 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); + 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); + 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); + } + std::cout << "valid:" << (rtn ? "y" : "n") << std::endl; } std::size_t flop = std::size_t(2) * M * N * K; @@ -198,5 +285,6 @@ int main(int argc, char* argv[]) std::cout << "Perf: " << ave_time << " ms, " << tflops << " TFlops, " << gb_per_sec << " GB/s" << std::endl; - return !pass; + // return !pass; + return rtn; } diff --git a/example/ck_tile/39_gemm_softmax/grid_gemm.hpp b/example/ck_tile/39_gemm_softmax/grid_gemm.hpp index 7946da263e..04dbbc8dc0 100755 --- a/example/ck_tile/39_gemm_softmax/grid_gemm.hpp +++ b/example/ck_tile/39_gemm_softmax/grid_gemm.hpp @@ -11,6 +11,8 @@ struct GridGemm using ADataType = typename Problem::ADataType; using BDataType = typename Problem::BDataType; using CDataType = typename Problem::CDataType; + using WeightType = typename Problem::WeightType; + using IndexType = typename Problem::IndexType; using AccDataType = typename Problem::AccDataType; using ComputeDataType = float; using CElementFunction = typename Problem::CElementFunction; @@ -18,15 +20,17 @@ struct GridGemm static constexpr auto kMPerBlock = Policy::kMPerBlock; static constexpr auto kNPerBlock = Policy::kNPerBlock; static constexpr auto kKPerBlock = Policy::kKPerBlock; + static constexpr auto topk = Policy::kTopKPerBlock; - template + template CK_TILE_DEVICE void operator()(const AGridTensorView& a_grid, const BGridTensorView& b_grid, - CGridTensorView& c_grid, + ValueGridTensorView& value_grid, + IndexGridTensorView& index_grid, const CElementFunction& c_element_func) const { const auto M = a_grid.get_tensor_descriptor().get_length(number<0>{}); - const auto N = c_grid.get_tensor_descriptor().get_length(number<1>{}); + const auto N = b_grid.get_tensor_descriptor().get_length(number<0>{}); const auto K = a_grid.get_tensor_descriptor().get_length(number<1>{}); // divide problem @@ -56,22 +60,52 @@ struct GridGemm __shared__ char p_smem_char[block_gemm_pipeline.GetStaticLdsSize()]; - // store C - auto c_window = make_tile_window( - c_grid, make_tuple(number{}, number{}), {iM, iN}); + // // store C + // auto c_window = make_tile_window( + // c_grid, make_tuple(number{}, number{}), {iM, iN}); - // block_gemm_pipeline(a_block_window, b_block_window, c_window, K / kKPerBlock, p_smem_char, c_element_func); + // store value and index + auto value_window = make_tile_window( + value_grid, make_tuple(number{}, number{}), {iM, iN}, + make_static_tile_distribution( + tile_distribution_encoding, + tuple, sequence<4, 4>>, + tuple, sequence<1, 2>>, + tuple, sequence<2, 0>>, + sequence<1, 2>, + sequence<0, 1>>{})); + auto index_window = make_tile_window( + index_grid, make_tuple(number{}, number{}), {iM, iN}, + make_static_tile_distribution( + tile_distribution_encoding, + tuple, sequence<4, 4>>, + tuple, sequence<1, 2>>, + tuple, sequence<2, 0>>, + sequence<1, 2>, + sequence<0, 1>>{})); - const auto acc_block_tile = - block_gemm_pipeline(a_block_window, b_block_window, K / kKPerBlock, p_smem_char); + using ValueBlockTileDistr = decltype(value_window.get_tile_distribution()); + using IndexBlockTileDistr = decltype(index_window.get_tile_distribution()); - // cast to CDataType and apply CElementFunction - const auto c_block_tile = tile_elementwise_in( - [&](const auto& acc) { return c_element_func(type_convert(acc)); }, - acc_block_tile); + using ValueBlockTile = decltype(make_static_distributed_tensor(ValueBlockTileDistr{})); + using IndexBlockTile = decltype(make_static_distributed_tensor(IndexBlockTileDistr{})); + ValueBlockTile value_block_tile; + IndexBlockTile index_block_tile; - store_tile(c_window, c_block_tile); + block_gemm_pipeline(a_block_window, b_block_window, value_block_tile, index_block_tile, K / kKPerBlock, p_smem_char); + + // cast DataType and apply CElementFunction + const auto value_cast_block_tile = tile_elementwise_in( + [&](const auto& value) { return c_element_func(type_convert(value)); }, + value_block_tile); + + // const auto index_cast_block_tile = tile_elementwise_in( + // [&](const auto& index) { return c_element_func(type_convert(index)); }, + // index_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/reference_gemm.hpp b/example/ck_tile/39_gemm_softmax/reference_gemm.hpp index 0c6bbecfc1..fd01780a53 100755 --- a/example/ck_tile/39_gemm_softmax/reference_gemm.hpp +++ b/example/ck_tile/39_gemm_softmax/reference_gemm.hpp @@ -6,13 +6,17 @@ #include "ck_tile/core.hpp" #include "ck_tile/host/host_tensor.hpp" -template -void reference_basic_gemm_softmax(const ck_tile::HostTensor& a_m_k, +template +void reference_basic_gemm_softmax_topk(const ck_tile::HostTensor& a_m_k, const ck_tile::HostTensor& b_n_k, - ck_tile::HostTensor& c_m_n) + ck_tile::HostTensor& y_values, + ck_tile::HostTensor& y_indices, + ck_tile::index_t topk) { + 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) @@ -62,4 +66,6 @@ void reference_basic_gemm_softmax(const ck_tile::HostTensor& a_m_k, 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); }