From bb918bfd04c4752a12f3b892f79e257bed66688a Mon Sep 17 00:00:00 2001
From: huizzhan <huizzhan@amd.com>
Date: Fri, 15 Aug 2025 08:52:55 +0000
Subject: [PATCH] grouped_topk perf test

---
 example/ck_tile/09_topk_softmax/CMakeLists.txt       |  0
 example/ck_tile/09_topk_softmax/README.md            |  0
 example/ck_tile/09_topk_softmax/script/smoke_test.sh |  0
 example/ck_tile/09_topk_softmax/topk_softmax.cpp     | 12 ++++++++++--
 example/ck_tile/09_topk_softmax/topk_softmax_api.cpp |  8 ++++++--
 example/ck_tile/09_topk_softmax/topk_softmax_api.hpp |  0
 6 files changed, 16 insertions(+), 4 deletions(-)
 mode change 100644 => 100755 example/ck_tile/09_topk_softmax/CMakeLists.txt
 mode change 100644 => 100755 example/ck_tile/09_topk_softmax/README.md
 mode change 100644 => 100755 example/ck_tile/09_topk_softmax/script/smoke_test.sh
 mode change 100644 => 100755 example/ck_tile/09_topk_softmax/topk_softmax.cpp
 mode change 100644 => 100755 example/ck_tile/09_topk_softmax/topk_softmax_api.cpp
 mode change 100644 => 100755 example/ck_tile/09_topk_softmax/topk_softmax_api.hpp

diff --git a/example/ck_tile/09_topk_softmax/CMakeLists.txt b/example/ck_tile/09_topk_softmax/CMakeLists.txt
old mode 100644
new mode 100755
diff --git a/example/ck_tile/09_topk_softmax/README.md b/example/ck_tile/09_topk_softmax/README.md
old mode 100644
new mode 100755
diff --git a/example/ck_tile/09_topk_softmax/script/smoke_test.sh b/example/ck_tile/09_topk_softmax/script/smoke_test.sh
old mode 100644
new mode 100755
diff --git a/example/ck_tile/09_topk_softmax/topk_softmax.cpp b/example/ck_tile/09_topk_softmax/topk_softmax.cpp
old mode 100644
new mode 100755
index 6fc25631fd..63fd922b3a
--- a/example/ck_tile/09_topk_softmax/topk_softmax.cpp
+++ b/example/ck_tile/09_topk_softmax/topk_softmax.cpp
@@ -71,6 +71,8 @@ template <typename InputType, typename WeightType, typename IndexType = ck_tile:
 auto reference_topk_softmax(const ck_tile::HostTensor<InputType>& x,
                             ck_tile::HostTensor<WeightType>& y_values,
                             ck_tile::HostTensor<IndexType>& y_indices,
+                            // ck_tile::index_t num_expert_group,
+                            // ck_tile::index_t topk_group,
                             ck_tile::index_t k,
                             ck_tile::index_t dim = -1,
                             bool largest         = true,
@@ -80,14 +82,15 @@ auto reference_topk_softmax(const ck_tile::HostTensor<InputType>& x,
 
     auto y = reference_softmax<InputType, WeightType, WeightType>(x, dim);
     reference_topk(y, y_values, y_indices, k, dim, largest, sorted);
+    // reference_grouped_topk(y, y_values, y_indices, k, num_expert_group, topk_group, dim, largest, sorted);
 }
 
 // different threshold for different dtype
 template <typename DataType>
 auto get_elimit(std::string /*init_method*/)
 {
-    double rtol = 1e-3;
-    double atol = 1e-3;
+    double rtol = 1e-6;
+    double atol = 1e-6;
     return ck_tile::make_tuple(rtol, atol);
 }
 
@@ -152,6 +155,9 @@ bool test_topk_softmax(ck_tile::ArgParser args)
     int warmup              = args.get_int("warmup");
     int repeat              = args.get_int("repeat");
 
+    // int num_expert_group    = 16;
+    // int topk_group          = 2;
+
     if(stride_input < 0)
     {
         stride_input = experts;
@@ -246,6 +252,7 @@ bool test_topk_softmax(ck_tile::ArgParser args)
 
         reference_topk_softmax<InputType, WeightType, IndexType>(
             x_host, value_ref, index_ref, topk);
+            // x_host, value_ref, index_ref, num_expert_group, topk_group, topk);
 
         auto [rtol, atol] = get_elimit<InputType>("");
         for(int i_t = 0; i_t < tokens; i_t++)
@@ -270,6 +277,7 @@ bool test_topk_softmax(ck_tile::ArgParser args)
                                       rtol,
                                       atol);
         }
+        printf("rtol:%f  atol:%f\n", rtol, atol);
     }
 
     printf("valid:%s\n", rtn ? "y" : "n");
diff --git a/example/ck_tile/09_topk_softmax/topk_softmax_api.cpp b/example/ck_tile/09_topk_softmax/topk_softmax_api.cpp
old mode 100644
new mode 100755
index 249a307b81..94453a04b6
--- a/example/ck_tile/09_topk_softmax/topk_softmax_api.cpp
+++ b/example/ck_tile/09_topk_softmax/topk_softmax_api.cpp
@@ -53,6 +53,10 @@ float topk_softmax(topk_softmax_trait t, topk_softmax_kargs a, ck_tile::stream_c
         {
             TOPK_SOFTMAX_DISPATCH(192)
         }
+        else if(t.experts <= 256)
+        {
+            TOPK_SOFTMAX_DISPATCH(256)
+        }
 #else
         if(t.experts <= 128)
         {
@@ -86,9 +90,9 @@ float topk_softmax(topk_softmax_trait t, topk_softmax_kargs a, ck_tile::stream_c
         {
             TOPK_SOFTMAX_DISPATCH(128)
         }
-        else if(t.experts <= 192)
+        else if(t.experts <= 256)
         {
-            TOPK_SOFTMAX_DISPATCH(192)
+            TOPK_SOFTMAX_DISPATCH(256)
         }
 #endif
     }
diff --git a/example/ck_tile/09_topk_softmax/topk_softmax_api.hpp b/example/ck_tile/09_topk_softmax/topk_softmax_api.hpp
old mode 100644
new mode 100755