ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-18 01:28:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use simple best occupancy model to calculate the split-K.

Browse Source
This commit is contained in:
Ville Pietilä
2025-07-10 11:09:19 +00:00
parent 62b08dd1bb
commit 611d6ac82d
1 changed files with 7 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
include/ck/tensor_operation/gpu/device/impl/split_k_utils.hpp
View File

@@ -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<ck::index_t>(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<ck::index_t>(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)))
{