diff --git a/cpp/src/neighbors/detail/cagra/search_plan.cuh b/cpp/src/neighbors/detail/cagra/search_plan.cuh index 1b9e518d46..65f8c25e1d 100644 --- a/cpp/src/neighbors/detail/cagra/search_plan.cuh +++ b/cpp/src/neighbors/detail/cagra/search_plan.cuh @@ -201,9 +201,13 @@ struct search_plan_impl : public search_plan_impl_base { uint32_t _max_iterations = max_iterations; if (max_iterations == 0) { if (algo == search_algo::MULTI_CTA) { - constexpr uint32_t mc_itopk_size = 32; - constexpr uint32_t mc_search_width = 1; - _max_iterations = mc_itopk_size / mc_search_width; + constexpr size_t mc_itopk_size = 32; + const size_t minimum_depth = 16; + const auto effective_itopk_size = raft::ceildiv(itopk_size, mc_itopk_size) * mc_itopk_size; + const auto num_ctas = max(search_width, raft::ceildiv(effective_itopk_size, mc_itopk_size)); + + _max_iterations = minimum_depth + raft::ceildiv(mc_itopk_size - minimum_depth, num_ctas); + _max_iterations += raft::ceildiv(static_cast(topk), mc_itopk_size) - 1; } else { _max_iterations = itopk_size / search_width; }