From 611d6ac82d14e8f2a6431a2f511d721338499bb4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ville=20Pietil=C3=A4?= <>
Date: Thu, 10 Jul 2025 11:09:19 +0000
Subject: [PATCH] Use simple best occupancy model to calculate the split-K.

---
 .../gpu/device/impl/split_k_utils.hpp         | 24 ++++++-------------
 1 file changed, 7 insertions(+), 17 deletions(-)

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<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)))
     {