diff --git a/include/ck/tensor_operation/gpu/device/impl/split_k_utils.hpp b/include/ck/tensor_operation/gpu/device/impl/split_k_utils.hpp index 646591e956..87802fa7c3 100644 --- a/include/ck/tensor_operation/gpu/device/impl/split_k_utils.hpp +++ b/include/ck/tensor_operation/gpu/device/impl/split_k_utils.hpp @@ -32,24 +32,14 @@ inline ck::index_t get_best_occupancy_k_batch_value(int max_occupancy, ck::index static DeviceProperties device_properties; const int max_capacity = max_occupancy * device_properties.num_cu_; - constexpr ck::index_t k_batch_max = 1024; - const auto total_grid_size = grid_size_mn * num_conv_groups; - ck::index_t k_batch = static_cast(max_capacity / std::gcd(max_capacity, total_grid_size)); - if (k_batch > k_batch_max || k_batch == 0) + // Empirically, when we have more than one group, we can oversubscribe by the number of groups. + // This is because the groups are independent and can be processed in parallel. + ck::index_t k_batch = 1; + const auto optimal_split = static_cast(std::floor((1.0 * max_capacity) / (grid_size_mn))); + if (optimal_split > 1) { - // TODO: This could be improved by using Euclidian algorithm to find the optimal k_batch. - auto min_remainder = max_capacity; - for (ck::index_t k = 1; k <= k_batch_max; ++k) - { - const auto remainder = (total_grid_size * k) % max_capacity; - // For equal remainder values, prefer smaller k values. - if (remainder < min_remainder) - { - min_remainder = remainder; - k_batch = k; - } - } - } + k_batch = optimal_split; + } if (ck::EnvIsEnabled(CK_ENV(CK_LOGGING))) {