diff --git a/example/ck_tile/03_gemm/universal_gemm.cpp b/example/ck_tile/03_gemm/universal_gemm.cpp
index e6a2811918..b60a3b274b 100644
--- a/example/ck_tile/03_gemm/universal_gemm.cpp
+++ b/example/ck_tile/03_gemm/universal_gemm.cpp
@@ -12,6 +12,19 @@
 #include "ck_tile/host.hpp"
 #include "gemm_utils.hpp"
 
+template <typename Pipeline, ck_tile::TailNumber TN>
+void try_run(ck_tile::TailNumber tn)
+{
+    if constexpr(Pipeline::PrefetchStages > static_cast<int>(TN))
+    {
+        if(tn == TN)
+        {
+            RunSplitk(ck_tile::bool_constant<true>{},
+                      ck_tile::integral_constant<ck_tile::TailNumber, TN>{});
+        }
+    }
+}
+
 template <typename ADataType,
           typename BDataType,
           typename AccDataType,
@@ -164,7 +177,6 @@ float gemm_calc(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config&
             throw std::runtime_error(err.str());
         }
 #elif(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_MEMORY)
-        // Tail pipeline One to Seven
         if(tail_num == ck_tile::TailNumber::One)
         {
             RunSplitk(ck_tile::bool_constant<true>{},
@@ -176,60 +188,17 @@ float gemm_calc(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config&
                       ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Full>{});
         }
 
-        if constexpr(BaseGemmPipeline::PrefetchStages > 2)
-        {
-            if(tail_num == ck_tile::TailNumber::Two)
-            {
-                RunSplitk(
-                    ck_tile::bool_constant<true>{},
-                    ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Two>{});
-            }
-        }
-        if constexpr(BaseGemmPipeline::PrefetchStages > 3)
-        {
-            if(tail_num == ck_tile::TailNumber::Three)
-            {
-                RunSplitk(
-                    ck_tile::bool_constant<true>{},
-                    ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Three>{});
-            }
-        }
-        if constexpr(BaseGemmPipeline::PrefetchStages > 4)
-        {
-            if(tail_num == ck_tile::TailNumber::Four)
-            {
-                RunSplitk(
-                    ck_tile::bool_constant<true>{},
-                    ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Four>{});
-            }
-        }
-        if constexpr(BaseGemmPipeline::PrefetchStages > 5)
-        {
-            if(tail_num == ck_tile::TailNumber::Five)
-            {
-                RunSplitk(
-                    ck_tile::bool_constant<true>{},
-                    ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Five>{});
-            }
-        }
-        if constexpr(BaseGemmPipeline::PrefetchStages > 6)
-        {
-            if(tail_num == ck_tile::TailNumber::Six)
-            {
-                RunSplitk(
-                    ck_tile::bool_constant<true>{},
-                    ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Six>{});
-            }
-        }
-        if constexpr(BaseGemmPipeline::PrefetchStages > 7)
-        {
-            if(tail_num == ck_tile::TailNumber::Seven)
-            {
-                RunSplitk(
-                    ck_tile::bool_constant<true>{},
-                    ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Seven>{});
-            }
-        }
+        auto check_tail = [&](auto... TNs) {
+            (try_run<BaseGemmPipeline, decltype(TNs)::value>(tail_num), ...);
+        };
+
+        check_tail(ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Two>{},
+                   ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Three>{},
+                   ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Four>{},
+                   ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Five>{},
+                   ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Six>{},
+                   ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Seven>{});
+
 #elif(CK_TILE_PIPELINE_DEFAULT == CK_TILE_PIPELINE_COMPUTE_V4)
         if(tail_num == ck_tile::TailNumber::Three)
         {
@@ -259,7 +228,7 @@ float gemm_calc(const ck_tile::GemmHostArgs& args, const ck_tile::stream_config&
         else if(tail_num == ck_tile::TailNumber::Even)
         {
             RunSplitk(ck_tile::bool_constant<false>{},
-                      ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Odd>{});
+                      ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Even>{});
         }
         else
         {
diff --git a/test/ck_tile/gemm/test_gemm_pipeline_util.hpp b/test/ck_tile/gemm/test_gemm_pipeline_util.hpp
index 0329f16416..85742cb3de 100644
--- a/test/ck_tile/gemm/test_gemm_pipeline_util.hpp
+++ b/test/ck_tile/gemm/test_gemm_pipeline_util.hpp
@@ -63,6 +63,19 @@ struct GemmPipelineTypeSelector<GemmPipelineType::CompV4, Problem>
     using pipeline      = ck_tile::GemmPipelineAgBgCrCompV4<Problem>;
 };
 
+template <typename Pipeline, ck_tile::TailNumber TN>
+void try_run(ck_tile::TailNumber tn)
+{
+    if constexpr(Pipeline::PrefetchStages > static_cast<int>(TN))
+    {
+        if(tn == TN)
+        {
+            RunSplitk(ck_tile::bool_constant<true>{},
+                      ck_tile::integral_constant<ck_tile::TailNumber, TN>{});
+        }
+    }
+}
+
 template <typename Tuple>
 class TestCkTileGemmPipeline : public ::testing::Test
 {
@@ -251,60 +264,17 @@ class TestCkTileGemmPipeline : public ::testing::Test
                                                          ck_tile::TailNumber::Full>{});
                 }
 
-                if constexpr(BaseGemmPipeline::PrefetchStages > 2)
-                {
-                    if(tail_num == ck_tile::TailNumber::Two)
-                    {
-                        RunSplitk(ck_tile::bool_constant<true>{},
-                                  ck_tile::integral_constant<ck_tile::TailNumber,
-                                                             ck_tile::TailNumber::Two>{});
-                    }
-                }
-                if constexpr(BaseGemmPipeline::PrefetchStages > 3)
-                {
-                    if(tail_num == ck_tile::TailNumber::Three)
-                    {
-                        RunSplitk(ck_tile::bool_constant<true>{},
-                                  ck_tile::integral_constant<ck_tile::TailNumber,
-                                                             ck_tile::TailNumber::Three>{});
-                    }
-                }
-                if constexpr(BaseGemmPipeline::PrefetchStages > 4)
-                {
-                    if(tail_num == ck_tile::TailNumber::Four)
-                    {
-                        RunSplitk(ck_tile::bool_constant<true>{},
-                                  ck_tile::integral_constant<ck_tile::TailNumber,
-                                                             ck_tile::TailNumber::Four>{});
-                    }
-                }
-                if constexpr(BaseGemmPipeline::PrefetchStages > 5)
-                {
-                    if(tail_num == ck_tile::TailNumber::Five)
-                    {
-                        RunSplitk(ck_tile::bool_constant<true>{},
-                                  ck_tile::integral_constant<ck_tile::TailNumber,
-                                                             ck_tile::TailNumber::Five>{});
-                    }
-                }
-                if constexpr(BaseGemmPipeline::PrefetchStages > 6)
-                {
-                    if(tail_num == ck_tile::TailNumber::Six)
-                    {
-                        RunSplitk(ck_tile::bool_constant<true>{},
-                                  ck_tile::integral_constant<ck_tile::TailNumber,
-                                                             ck_tile::TailNumber::Six>{});
-                    }
-                }
-                if constexpr(BaseGemmPipeline::PrefetchStages > 7)
-                {
-                    if(tail_num == ck_tile::TailNumber::Seven)
-                    {
-                        RunSplitk(ck_tile::bool_constant<true>{},
-                                  ck_tile::integral_constant<ck_tile::TailNumber,
-                                                             ck_tile::TailNumber::Seven>{});
-                    }
-                }
+                auto check_tail = [&](auto... TNs) {
+                    (try_run<BaseGemmPipeline, decltype(TNs)::value>(tail_num), ...);
+                };
+
+                check_tail(
+                    ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Two>{},
+                    ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Three>{},
+                    ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Four>{},
+                    ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Five>{},
+                    ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Six>{},
+                    ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Seven>{});
             }
 
             if constexpr(PipelineType == GemmPipelineType::CompV4)
diff --git a/tile_engine/ops/gemm/configs/instance_combination.json b/tile_engine/ops/gemm/configs/instance_combination.json
index 66dbdafa11..53197ada6c 100644
--- a/tile_engine/ops/gemm/configs/instance_combination.json
+++ b/tile_engine/ops/gemm/configs/instance_combination.json
@@ -19,7 +19,7 @@
       "values": [256]
     },
     "tile_k": {
-      "values": [64, 32]
+      "values": [32]
     },
     "warp_m": {
       "values": [2]
diff --git a/tile_engine/ops/gemm/gemm_instance_builder.py b/tile_engine/ops/gemm/gemm_instance_builder.py
index a748c35feb..3839523e3d 100755
--- a/tile_engine/ops/gemm/gemm_instance_builder.py
+++ b/tile_engine/ops/gemm/gemm_instance_builder.py
@@ -37,7 +37,9 @@ DEFAULT_EPILOGUE = """
                                                                       WarpTileM,
                                                                       WarpTileN,
                                                                       WarpTileK,
-                                                                      UniversalGemmProblem::TransposeC>>;
+                                                                      UniversalGemmProblem::TransposeC,
+                                                                      true,
+                                                                      memory_operation>>;
 """
 
 CSHUFFLE_EPILOGUE = """
@@ -55,22 +57,23 @@ CSHUFFLE_EPILOGUE = """
                                                              WarpTileM,
                                                              WarpTileN,
                                                              WarpTileK,
-                                                             UniversalGemmProblem::TransposeC>>;
+                                                             UniversalGemmProblem::TransposeC,
+                                                             memory_operation>>;
 """
 HOT_LOOP_FALSE = """
             if(tail_num == ck_tile::TailNumber::Full)
             {
-                Run(ck_tile::bool_constant<false>{},
+                RunSplitk(ck_tile::bool_constant<false>{},
                     ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Full>{});
             }
             else if(tail_num == ck_tile::TailNumber::Odd)
             {
-                Run(ck_tile::bool_constant<false>{},
+                RunSplitk(ck_tile::bool_constant<false>{},
                     ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Odd>{});
             }
             else if(tail_num == ck_tile::TailNumber::Even)
             {
-                Run(ck_tile::bool_constant<false>{},
+                RunSplitk(ck_tile::bool_constant<false>{},
                     ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Even>{});
             }
             else
@@ -79,68 +82,43 @@ HOT_LOOP_FALSE = """
             }  
 """
 RUN_MEM = """
-            if(tail_num == ck_tile::TailNumber::One)
-            {
-                Run(ck_tile::bool_constant<true>{},
+            // Handle One and Full cases directly
+            if (tail_num == ck_tile::TailNumber::One) {
+                RunSplitk(ck_tile::bool_constant<true>{},
                     ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::One>{});
-            }
-            else if(tail_num == ck_tile::TailNumber::Full)
-            {
-                Run(ck_tile::bool_constant<true>{},
+            } else if (tail_num == ck_tile::TailNumber::Full) {
+                RunSplitk(ck_tile::bool_constant<true>{},
                     ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Full>{});
             }
+            // Variadic call using fold expression
+            auto check_tail = [&](auto... TNs) {
+                (try_run< BaseGemmPipeline, decltype(TNs)::value>(tail_num), ...);
+            };
 
-            if constexpr(BaseGemmPipeline::PrefetchStages > 2)
-            {
-                if(tail_num == ck_tile::TailNumber::Two)
-                {
-                    Run(ck_tile::bool_constant<true>{},
-                        ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Two>{});
-                }
-        
-                if(tail_num == ck_tile::TailNumber::Three)
-                {
-                    Run(ck_tile::bool_constant<true>{},
-                        ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Three>{});
-                }
-                if(tail_num == ck_tile::TailNumber::Four)
-                {
-                    Run(ck_tile::bool_constant<true>{},
-                        ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Four>{});
-                }
-                if(tail_num == ck_tile::TailNumber::Five)
-                {
-                    Run(ck_tile::bool_constant<true>{},
-                        ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Five>{});
-                }
-                if(tail_num == ck_tile::TailNumber::Six)
-                {
-                    Run(ck_tile::bool_constant<true>{},
-                        ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Six>{});
-                }
-                if(tail_num == ck_tile::TailNumber::Seven)
-                {
-                    Run(ck_tile::bool_constant<true>{},
-                        ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Seven>{});
-                }
-                throw std::runtime_error("The tile number is wrong! It should not exceed the prefetch stage numbers");
-            }
+            check_tail(
+                ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Two>{},
+                ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Three>{},
+                ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Four>{},
+                ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Five>{},
+                ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Six>{},
+                ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Seven>{}
+            );
 """
 
 RUN_COMPV3 = """
             if(tail_num == ck_tile::TailNumber::Full)
             {
-                Run(ck_tile::bool_constant<true>{},
+                RunSplitk(ck_tile::bool_constant<true>{},
                     ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Full>{});
             }
             else if(tail_num == ck_tile::TailNumber::Odd)
             {
-                Run(ck_tile::bool_constant<true>{},
+                RunSplitk(ck_tile::bool_constant<true>{},
                     ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Odd>{});
             }
             else if(tail_num == ck_tile::TailNumber::Even)
             {
-                Run(ck_tile::bool_constant<true>{},
+                RunSplitk(ck_tile::bool_constant<true>{},
                     ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Even>{});
             }
             else
@@ -152,12 +130,12 @@ RUN_COMPV3 = """
 RUN_COMPV4 = """
             if(tail_num == ck_tile::TailNumber::Three)
             {
-                Run(ck_tile::bool_constant<true>{},
+                RunSplitk(ck_tile::bool_constant<true>{},
                     ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Three>{});
             }
             else
             {
-                Run(ck_tile::bool_constant<true>{},
+                RunSplitk(ck_tile::bool_constant<true>{},
                     ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Two>{});
             }
 """
@@ -347,6 +325,15 @@ namespace {group_name} {{
                                kPadM: bool, kPadN: bool, kPadK: bool) -> str:
         """Generate kernel struct template"""
         return f"""
+template <typename Pipeline, ck_tile::TailNumber TN>
+void try_run(ck_tile::TailNumber tn) {{
+    if constexpr (Pipeline::PrefetchStages > static_cast<int>(TN)) {{
+        if (tn == TN) {{
+            RunSplitk(ck_tile::bool_constant<true>{{}},
+                ck_tile::integral_constant<ck_tile::TailNumber, TN>{{}});
+        }}
+    }}
+}}
 template <int TileM, int TileN, int TileK,
           int WarpM, int WarpN, int WarpK,
           int WarpTileM, int WarpTileN, int WarpTileK,
@@ -355,7 +342,7 @@ struct GemmKernel {{
     static constexpr bool kPadM = {BOOL_MAP(kPadM)};
     static constexpr bool kPadN = {BOOL_MAP(kPadN)};
     static constexpr bool kPadK = {BOOL_MAP(kPadK)};
-
+   
     static float launch(ck_tile::GemmHostArgs& args, const ck_tile::stream_config& s) {{
         static constexpr bool permuteA = false;
         static constexpr bool permuteB = false;
@@ -399,10 +386,11 @@ struct GemmKernel {{
 
         float ave_time{{0}};
 
-        const auto Run = [&](const auto has_hot_loop_, const auto tail_number_) {{
+        const auto Run = [&](const auto has_hot_loop_, const auto tail_number_, const auto memory_operation_) {{
             constexpr bool has_hot_loop_v = has_hot_loop_.value;
             constexpr auto tail_number_v  = tail_number_.value;
             constexpr auto scheduler      = {SCHEDULER_MAP[scheduler]};
+            constexpr auto memory_operation = memory_operation_.value;
 
             using UniversalGemmProblem = 
                 ck_tile::UniversalGemmPipelineProblem<ADataType,
@@ -442,6 +430,20 @@ struct GemmKernel {{
 
         }};
 
+        const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {{
+            if(args.k_batch == 1) {{
+                Run(has_hot_loop_,
+                    tail_number_,
+                    ck_tile::integral_constant<ck_tile::memory_operation_enum,
+                                            ck_tile::memory_operation_enum::set>{{}});
+            }} else {{
+                Run(has_hot_loop_,
+                    tail_number_,
+                    ck_tile::integral_constant<ck_tile::memory_operation_enum,
+                                            ck_tile::memory_operation_enum::atomic_add>{{}});
+            }}
+        }};
+
         if(has_hot_loop) {{
             {HOT_LOOP_TRUE[pipeline]}
         }} else {{
@@ -450,6 +452,7 @@ struct GemmKernel {{
 
         return ave_time;
     }}
+    
     static std::string get_name() {{
         return std::string("GemmKernel<Bllktile: ") + std::to_string(TileM) + "x" + std::to_string(TileN) + "x" + std::to_string(TileK) + ", " +
                 "WaveMap: " + std::to_string(WarpM) + "x" + std::to_string(WarpN) + "x" + std::to_string(WarpK) + ", " +