[CK-Tile] Refactor base pipeline usage (#3251)

* initial poc

* factor out common parts in operator()

* cv4

* rest of the universal gemm pipelines

* fix test

* remove boilerplate from tile engine

* fix example

* fix example

* format

* fix tests build for gemm

* remove base pipeline codegen from gemm instance builder

* unify v3 logic with the rest of universal gemm pipelines

* fix build for multi abd test

* fix test gemm multi d

* fix build for weight preshuffle

* fix grouped gemm test

* fix grouped gemm multi d test

* fix grouped gemm preshuffle

* fix grouped gemm example except for quant

* fix gemm preshuffle

* fix splitk 2 stage example

* fix batched gemm example

* fix multid example

* fix multiabd example

* fix batched gemm test

* fixup

* fix examples build

* fix grouped gemm test build

* fix smoke builder

example/ck_tile/20_grouped_convolution/grouped_convolution_backward_data_invoker.hpp

@@ -57,43 +57,9 @@ struct GroupedConvolutionBackwardDataInvoker
             GroupedConvTraitsType::FixedGemmParams::UseStructuredSparsity,
             GroupedConvTraitsType::FixedGemmParams::Persistent,
             ConvConfig::NumWaveGroups>;
+        constexpr auto scheduler = ConvConfig::Scheduler;
 
-        using GemmPipelineProblem = ck_tile::GemmPipelineProblem<
-            OutDataType,
-            WeiDataType,
-            AccDataType,
-            GemmShape,
-            typename GroupedConvTraitsType::template GroupedConvImplicitGemmTraitsBwdData<
-                ConvConfig::NumWaveGroups>,
-            ck_tile::element_wise::PassThrough,
-            ck_tile::element_wise::PassThrough,
-            InDataType,
-            GroupedConvTraitsType::FixedGemmParams::FixedVectorSize,
-            GroupedConvTraitsType::VectorSizeA,
-            GroupedConvTraitsType::VectorSizeB>;
-
-        using BaseGemmPipeline = typename PipelineTypeTraits<
-            ConvConfig::Pipeline>::template UniversalGemmPipeline<GemmPipelineProblem>;
-
-        const ck_tile::index_t gemm_k =
-            args.K_ * std::accumulate(args.filter_spatial_lengths_.begin(),
-                                      args.filter_spatial_lengths_.end(),
-                                      1,
-                                      std::multiplies<ck_tile::index_t>());
-
-        const ck_tile::index_t k_grain     = args.k_batch * ConvConfig::K_Tile;
-        const ck_tile::index_t K_split     = (gemm_k + k_grain - 1) / k_grain * ConvConfig::K_Tile;
-        const ck_tile::index_t num_loop    = TilePartitioner::GetLoopNum(K_split);
-        const bool has_hot_loop            = BaseGemmPipeline::BlockHasHotloop(num_loop);
-        const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
-        float ave_time{0};
-
-        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        = ConvConfig::Scheduler;
+        const auto Run = [&](const auto memory_operation_) {
             constexpr auto memory_operation = memory_operation_.value;
 
             using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem<
@@ -103,8 +69,6 @@ struct GroupedConvolutionBackwardDataInvoker
                 GemmShape,
                 GemmUniversalTraits,
                 scheduler,
-                has_hot_loop_v,
-                tail_number_v,
                 ck_tile::element_wise::PassThrough,
                 ck_tile::element_wise::PassThrough,
                 InDataType,
@@ -170,26 +134,19 @@ struct GroupedConvolutionBackwardDataInvoker
                     kargs.in_ptr, 0, args.template GetInputByte<InDataType>(), s.stream_id_));
             };
 
-            ave_time = ck_tile::launch_kernel_time_mask(
+            return ck_tile::launch_kernel_time_mask(
                 s,
                 preprocess,
                 ck_tile::make_kernel<ConvConfig::kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
-
-            return ave_time;
         };
 
-        const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {
-            if(args.k_batch == 1)
-            {
-                Run(has_hot_loop_, tail_number_, MemoryOpSet{});
-            }
-            else
-            {
-                Run(has_hot_loop_, tail_number_, MemoryOpAtomicAdd{});
-            }
-        };
-
-        BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
-        return ave_time;
+        if(args.k_batch == 1)
+        {
+            return Run(MemoryOpSet{});
+        }
+        else
+        {
+            return Run(MemoryOpAtomicAdd{});
+        }
     }
 };

example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_invoker.hpp

@@ -57,43 +57,9 @@ struct GroupedConvolutionBackwardWeightInvoker
             GroupedConvTraitsType::FixedGemmParams::UseStructuredSparsity,
             GroupedConvTraitsType::FixedGemmParams::Persistent,
             ConvConfig::NumWaveGroups>;
+        constexpr auto scheduler = ConvConfig::Scheduler;
 
-        using GemmPipelineProblem = ck_tile::GemmPipelineProblem<
-            OutDataType,
-            InDataType,
-            AccDataType,
-            GemmShape,
-            typename GroupedConvTraitsType::template GroupedConvImplicitGemmTraitsBwdWeight<
-                ConvConfig::NumWaveGroups>,
-            ck_tile::element_wise::PassThrough,
-            ck_tile::element_wise::PassThrough,
-            WeiDataType,
-            GroupedConvTraitsType::FixedGemmParams::FixedVectorSize,
-            GroupedConvTraitsType::VectorSizeA,
-            GroupedConvTraitsType::VectorSizeB>;
-
-        using BaseGemmPipeline = typename PipelineTypeTraits<
-            ConvConfig::Pipeline>::template UniversalGemmPipeline<GemmPipelineProblem>;
-
-        const ck_tile::index_t gemm_k =
-            args.N_ * std::accumulate(args.output_spatial_lengths_.begin(),
-                                      args.output_spatial_lengths_.end(),
-                                      1,
-                                      std::multiplies<ck_tile::index_t>());
-
-        const ck_tile::index_t k_grain     = args.k_batch * ConvConfig::K_Tile;
-        const ck_tile::index_t K_split     = (gemm_k + k_grain - 1) / k_grain * ConvConfig::K_Tile;
-        const ck_tile::index_t num_loop    = TilePartitioner::GetLoopNum(K_split);
-        const bool has_hot_loop            = BaseGemmPipeline::BlockHasHotloop(num_loop);
-        const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
-        float ave_time{0};
-
-        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        = ConvConfig::Scheduler;
+        const auto Run = [&](const auto memory_operation_) {
             constexpr auto memory_operation = memory_operation_.value;
 
             using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem<
@@ -103,8 +69,6 @@ struct GroupedConvolutionBackwardWeightInvoker
                 GemmShape,
                 GemmUniversalTraits,
                 scheduler,
-                has_hot_loop_v,
-                tail_number_v,
                 ck_tile::element_wise::PassThrough,
                 ck_tile::element_wise::PassThrough,
                 WeiDataType,
@@ -176,26 +140,19 @@ struct GroupedConvolutionBackwardWeightInvoker
                 }
             };
 
-            ave_time = ck_tile::launch_kernel_time_mask(
+            return ck_tile::launch_kernel_time_mask(
                 s,
                 preprocess,
                 ck_tile::make_kernel<ConvConfig::kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
-
-            return ave_time;
         };
 
-        const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {
-            if(args.k_batch == 1)
-            {
-                Run(has_hot_loop_, tail_number_, MemoryOpSet{});
-            }
-            else
-            {
-                Run(has_hot_loop_, tail_number_, MemoryOpAtomicAdd{});
-            }
-        };
-
-        BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
-        return ave_time;
+        if(args.k_batch == 1)
+        {
+            return Run(MemoryOpSet{});
+        }
+        else
+        {
+            return Run(MemoryOpAtomicAdd{});
+        }
     }
 };

example/ck_tile/20_grouped_convolution/grouped_convolution_backward_weight_two_stage_invoker.hpp

@@ -60,42 +60,9 @@ struct GroupedConvolutionBackwardWeightTwoStageInvoker
             GroupedConvTraitsType::FixedGemmParams::Persistent,
             ConvConfig::NumWaveGroups>;
 
-        using GemmPipelineProblem = ck_tile::GemmPipelineProblem<
-            OutDataType,
-            InDataType,
-            AccDataType,
-            GemmShape,
-            typename GroupedConvTraitsType::template GroupedConvImplicitGemmTraitsBwdWeight<
-                ConvConfig::NumWaveGroups>,
-            ck_tile::element_wise::PassThrough,
-            ck_tile::element_wise::PassThrough,
-            WeiDataType,
-            GroupedConvTraitsType::FixedGemmParams::FixedVectorSize,
-            GroupedConvTraitsType::VectorSizeA,
-            GroupedConvTraitsType::VectorSizeB>;
+        constexpr auto scheduler = ConvConfig::Scheduler;
 
-        using BaseGemmPipeline = typename PipelineTypeTraits<
-            ConvConfig::Pipeline>::template UniversalGemmPipeline<GemmPipelineProblem>;
-
-        const ck_tile::index_t gemm_k =
-            args.N_ * std::accumulate(args.output_spatial_lengths_.begin(),
-                                      args.output_spatial_lengths_.end(),
-                                      1,
-                                      std::multiplies<ck_tile::index_t>());
-
-        const ck_tile::index_t k_grain     = args.k_batch * ConvConfig::K_Tile;
-        const ck_tile::index_t K_split     = (gemm_k + k_grain - 1) / k_grain * ConvConfig::K_Tile;
-        const ck_tile::index_t num_loop    = TilePartitioner::GetLoopNum(K_split);
-        const bool has_hot_loop            = BaseGemmPipeline::BlockHasHotloop(num_loop);
-        const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
-        float ave_time{0};
-
-        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        = ConvConfig::Scheduler;
+        const auto Run = [&](const auto memory_operation_) {
             constexpr auto memory_operation = memory_operation_.value;
 
             using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem<
@@ -105,8 +72,6 @@ struct GroupedConvolutionBackwardWeightTwoStageInvoker
                 GemmShape,
                 GemmUniversalTraits,
                 scheduler,
-                has_hot_loop_v,
-                tail_number_v,
                 ck_tile::element_wise::PassThrough,
                 ck_tile::element_wise::PassThrough,
                 WeiDataType,
@@ -209,7 +174,6 @@ struct GroupedConvolutionBackwardWeightTwoStageInvoker
             {
                 std::cout << "Launching kernel with args: " << Kernel::GetName() << '\n'
                           << "shape: " << GemmShape::GetName() << '\n'
-                          << "problem: " << GemmPipelineProblem::GetName() << '\n'
                           << "pipeline: " << GemmPipeline::GetName() << '\n'
                           << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}"
                           << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z
@@ -228,7 +192,7 @@ struct GroupedConvolutionBackwardWeightTwoStageInvoker
                                        s.stream_id_));
             };
 
-            ave_time = ck_tile::launch_kernel_time_mask(
+            return ck_tile::launch_kernel_time_mask(
                 s,
                 preprocess,
                 ck_tile::make_kernel<ConvConfig::kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs),
@@ -242,22 +206,15 @@ struct GroupedConvolutionBackwardWeightTwoStageInvoker
                     ck_tile::make_tuple(shape[1], 1), // Output Stride
                     input_tensors,
                     static_cast<WeiDataType*>(c_ptr)));
-
-            return ave_time;
         };
 
-        const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {
-            if(args.k_batch == 1)
-            {
-                Run(has_hot_loop_, tail_number_, MemoryOpSet{});
-            }
-            else
-            {
-                Run(has_hot_loop_, tail_number_, MemoryOpAtomicAdd{});
-            }
-        };
-
-        BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
-        return ave_time;
+        if(args.k_batch == 1)
+        {
+            return Run(MemoryOpSet{});
+        }
+        else
+        {
+            return Run(MemoryOpAtomicAdd{});
+        }
     }
 };

example/ck_tile/20_grouped_convolution/grouped_convolution_forward_invoker.hpp

@@ -65,148 +65,96 @@ struct GroupedConvolutionForwardInvoker
             GroupedConvTraitsType::FixedGemmParams::UseStructuredSparsity,
             GroupedConvTraitsType::FixedGemmParams::Persistent,
             ConvConfig::NumWaveGroups>;
-
-        using GemmPipelineProblem = ck_tile::GemmPipelineProblem<
-            InDataType,
-            WeiDataType,
-            AccDataType,
-            GemmShape,
-            typename GroupedConvTraitsType::template GroupedConvImplicitGemmTraitsFwd<
-                ConvConfig::NumWaveGroups>,
-            ck_tile::element_wise::PassThrough,
-            ck_tile::element_wise::PassThrough,
-            OutDataType,
-            GroupedConvTraitsType::FixedGemmParams::FixedVectorSize,
-            GroupedConvTraitsType::VectorSizeA,
-            GroupedConvTraitsType::VectorSizeB>;
-
-        using BaseGemmPipeline = typename PipelineTypeTraits<
-            ConvConfig::Pipeline>::template UniversalGemmPipeline<GemmPipelineProblem>;
-
-        const ck_tile::index_t gemm_k =
-            args.C_ * std::accumulate(args.filter_spatial_lengths_.begin(),
-                                      args.filter_spatial_lengths_.end(),
-                                      1,
-                                      std::multiplies<ck_tile::index_t>());
-
-        // Split-K parameters
-        const ck_tile::index_t k_grain     = args.k_batch * ConvConfig::K_Tile;
-        const ck_tile::index_t K_split     = (gemm_k + k_grain - 1) / k_grain * ConvConfig::K_Tile;
-        const ck_tile::index_t num_loop    = TilePartitioner::GetLoopNum(K_split);
-        const bool has_hot_loop            = BaseGemmPipeline::BlockHasHotloop(num_loop);
-        const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
-        float ave_time{0};
+        constexpr auto scheduler = ConvConfig::Scheduler;
 
         // =====================================================================
         // Regular Convolution: Simple, no split-image
         // =====================================================================
-        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        = ConvConfig::Scheduler;
-                constexpr auto memory_operation = memory_operation_.value;
+        const auto Run = [&](const auto memory_operation_) {
+            constexpr auto memory_operation = memory_operation_.value;
 
-                using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem<
-                    InDataType,
-                    WeiDataType,
-                    AccDataType,
-                    GemmShape,
-                    GemmUniversalTraits,
-                    scheduler,
-                    has_hot_loop_v,
-                    tail_number_v,
-                    ck_tile::element_wise::PassThrough,
-                    ck_tile::element_wise::PassThrough,
-                    OutDataType,
-                    GroupedConvTraitsType::FixedGemmParams::FixedVectorSize,
-                    GroupedConvTraitsType::VectorSizeA,
-                    GroupedConvTraitsType::VectorSizeB>;
+            using UniversalGemmProblem = ck_tile::UniversalGemmPipelineProblem<
+                InDataType,
+                WeiDataType,
+                AccDataType,
+                GemmShape,
+                GemmUniversalTraits,
+                scheduler,
+                ck_tile::element_wise::PassThrough,
+                ck_tile::element_wise::PassThrough,
+                OutDataType,
+                GroupedConvTraitsType::FixedGemmParams::FixedVectorSize,
+                GroupedConvTraitsType::VectorSizeA,
+                GroupedConvTraitsType::VectorSizeB>;
 
-                using GemmPipeline = typename PipelineTypeTraits<
-                    ConvConfig::Pipeline>::template GemmPipeline<UniversalGemmProblem>;
+            using GemmPipeline = typename PipelineTypeTraits<
+                ConvConfig::Pipeline>::template GemmPipeline<UniversalGemmProblem>;
 
-                using ConvEpilogue = ck_tile::CShuffleEpilogue<ck_tile::CShuffleEpilogueProblem<
-                    InDataType,
-                    WeiDataType,
-                    DsDataType,
-                    AccDataType,
-                    OutDataType,
-                    typename GroupedConvTraitsType::ImplicitGemmDsLayout,
-                    typename GroupedConvTraitsType::FixedGemmParams::ELayout,
-                    CDElementWise,
-                    TilePartitioner::MPerBlock,
-                    TilePartitioner::NPerBlock,
-                    ConvConfig::M_Warp,
-                    ConvConfig::N_Warp,
-                    ConvConfig::M_Warp_Tile,
-                    ConvConfig::N_Warp_Tile,
-                    ConvConfig::K_Warp_Tile,
-                    GroupedConvTraitsType::FixedGemmParams::TransposeC,
-                    memory_operation,
-                    ConvConfig::NumWaveGroups,
-                    GroupedConvTraitsType::FixedGemmParams::FixedVectorSize,
-                    GroupedConvTraitsType::VectorSizeC>>;
+            using ConvEpilogue = ck_tile::CShuffleEpilogue<ck_tile::CShuffleEpilogueProblem<
+                InDataType,
+                WeiDataType,
+                DsDataType,
+                AccDataType,
+                OutDataType,
+                typename GroupedConvTraitsType::ImplicitGemmDsLayout,
+                typename GroupedConvTraitsType::FixedGemmParams::ELayout,
+                CDElementWise,
+                TilePartitioner::MPerBlock,
+                TilePartitioner::NPerBlock,
+                ConvConfig::M_Warp,
+                ConvConfig::N_Warp,
+                ConvConfig::M_Warp_Tile,
+                ConvConfig::N_Warp_Tile,
+                ConvConfig::K_Warp_Tile,
+                GroupedConvTraitsType::FixedGemmParams::TransposeC,
+                memory_operation,
+                ConvConfig::NumWaveGroups,
+                GroupedConvTraitsType::FixedGemmParams::FixedVectorSize,
+                GroupedConvTraitsType::VectorSizeC>>;
 
-                using Kernel = ck_tile::GroupedConvolutionForwardKernel<GroupedConvTraitsType,
-                                                                        TilePartitioner,
-                                                                        GemmPipeline,
-                                                                        ConvEpilogue>;
-                auto kargs   = Kernel::MakeKernelArgs(args);
+            using Kernel = ck_tile::GroupedConvolutionForwardKernel<GroupedConvTraitsType,
+                                                                    TilePartitioner,
+                                                                    GemmPipeline,
+                                                                    ConvEpilogue>;
+            auto kargs   = Kernel::MakeKernelArgs(args);
 
-                const dim3 grids  = Kernel::GridSize(kargs);
-                const dim3 blocks = Kernel::BlockSize();
+            const dim3 grids  = Kernel::GridSize(kargs);
+            const dim3 blocks = Kernel::BlockSize();
 
-                if(!Kernel::IsSupportedArgument(kargs))
-                {
-                    throw std::runtime_error("Wrong! Arguments not supported! Skipping conv!\n");
-                }
-
-                if(s.log_level_ > 0)
-                {
-                    std::cout << "Launching kernel with args: " << Kernel::GetName() << '\n'
-                              << "shape: " << GemmShape::GetName() << '\n'
-                              << "problem: " << UniversalGemmProblem::GetName() << '\n'
-                              << "pipeline: " << GemmPipeline::GetName() << '\n'
-                              << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}"
-                              << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z
-                              << "}" << '\n'
-                              << "Vector size A: " << GemmPipeline::GetVectorSizeA()
-                              << ", Vector size B: " << GemmPipeline::GetVectorSizeB()
-                              << ", Vector size C: " << ConvEpilogue::GetVectorSizeC() << std::endl;
-                }
-
-                ave_time = ck_tile::launch_kernel(s,
-                                                  ck_tile::make_kernel<ConvConfig::kBlockPerCu>(
-                                                      Kernel{}, grids, blocks, 0, kargs));
-
-                return ave_time;
-            };
-
-        // =====================================================================
-        // Split-K lambda
-        // =====================================================================
-        const auto RunSplitk = [&](const auto has_hot_loop_, const auto tail_number_) {
-            if(args.k_batch == 1)
+            if(!Kernel::IsSupportedArgument(kargs))
             {
-                Run.template operator()(has_hot_loop_, tail_number_, MemoryOpSet{});
+                throw std::runtime_error("Wrong! Arguments not supported! Skipping conv!\n");
             }
-            else
+
+            if(s.log_level_ > 0)
             {
-                Run.template operator()(has_hot_loop_, tail_number_, MemoryOpAtomicAdd{});
+                std::cout << "Launching kernel with args: " << Kernel::GetName() << '\n'
+                          << "shape: " << GemmShape::GetName() << '\n'
+                          << "problem: " << UniversalGemmProblem::GetName() << '\n'
+                          << "pipeline: " << GemmPipeline::GetName() << '\n'
+                          << "grid: {" << grids.x << ", " << grids.y << ", " << grids.z << "}"
+                          << ", blocks: {" << blocks.x << ", " << blocks.y << ", " << blocks.z
+                          << "}" << '\n'
+                          << "Vector size A: " << GemmPipeline::GetVectorSizeA()
+                          << ", Vector size B: " << GemmPipeline::GetVectorSizeB()
+                          << ", Vector size C: " << ConvEpilogue::GetVectorSizeC() << std::endl;
             }
+
+            return ck_tile::launch_kernel(
+                s,
+                ck_tile::make_kernel<ConvConfig::kBlockPerCu>(Kernel{}, grids, blocks, 0, kargs));
         };
 
         // =====================================================================
-        // Regular Convolution Example: ALWAYS uses regular path (Kernel<false>)
+        // Split-K dispatch
         // =====================================================================
-        // This example demonstrates regular convolution without split-image.
-        // For large images that don't fit in memory, use
-        // grouped_convolution_forward_split_image.cpp
-
-        // Launch kernel using regular path (no split-image)
-        BaseGemmPipeline::TailHandler(RunSplitk, has_hot_loop, tail_num);
-
-        return ave_time;
+        if(args.k_batch == 1)
+        {
+            return Run(MemoryOpSet{});
+        }
+        else
+        {
+            return Run(MemoryOpAtomicAdd{});
+        }
     }
 };

example/ck_tile/20_grouped_convolution/grouped_convolution_forward_large_tensor_invoker.hpp

@@ -72,36 +72,6 @@ struct GroupedConvolutionForwardInvoker
             GroupedConvTraitsTypeDefault::FixedGemmParams::Persistent,
             ConvConfig::NumWaveGroups>;
 
-        using GemmPipelineProblem = ck_tile::GemmPipelineProblem<
-            InDataType,
-            WeiDataType,
-            AccDataType,
-            GemmShape,
-            typename GroupedConvTraitsTypeDefault::template GroupedConvImplicitGemmTraitsFwd<
-                ConvConfig::NumWaveGroups>,
-            ck_tile::element_wise::PassThrough,
-            ck_tile::element_wise::PassThrough,
-            OutDataType,
-            GroupedConvTraitsTypeDefault::FixedGemmParams::FixedVectorSize,
-            GroupedConvTraitsTypeDefault::VectorSizeA,
-            GroupedConvTraitsTypeDefault::VectorSizeB>;
-
-        using BaseGemmPipeline = typename PipelineTypeTraits<
-            ConvConfig::Pipeline>::template UniversalGemmPipeline<GemmPipelineProblem>;
-
-        const ck_tile::index_t gemm_k =
-            args.C_ * std::accumulate(args.filter_spatial_lengths_.begin(),
-                                      args.filter_spatial_lengths_.end(),
-                                      1,
-                                      std::multiplies<ck_tile::index_t>());
-
-        // Split-K parameters
-        const ck_tile::index_t k_grain     = args.k_batch * ConvConfig::K_Tile;
-        const ck_tile::index_t K_split     = (gemm_k + k_grain - 1) / k_grain * ConvConfig::K_Tile;
-        const ck_tile::index_t num_loop    = TilePartitioner::GetLoopNum(K_split);
-        const bool has_hot_loop            = BaseGemmPipeline::BlockHasHotloop(num_loop);
-        const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
-
         using TransformType =
             ck_tile::TransformConvFwdToGemm<NDimSpatial,
                                             ck_tile::ConvolutionSpecialization::Default,
@@ -238,16 +208,12 @@ struct GroupedConvolutionForwardInvoker
             }
         }
 
+        constexpr auto scheduler = ConvConfig::Scheduler;
+
         // =====================================================================
         // Kernel launch lambda: Uses EnableSplitImage based on layout support
         // =====================================================================
-        const auto Run = [&](const auto has_hot_loop_,
-                             const auto tail_number_,
-                             const auto memory_operation_,
-                             const auto enable_split_image_) {
-            constexpr bool has_hot_loop_v   = has_hot_loop_.value;
-            constexpr auto tail_number_v    = tail_number_.value;
-            constexpr auto scheduler        = ConvConfig::Scheduler;
+        const auto Run = [&](const auto memory_operation_, const auto enable_split_image_) {
             constexpr auto memory_operation = memory_operation_.value;
             constexpr bool EnableSplitImage = enable_split_image_.value;
 
@@ -262,8 +228,6 @@ struct GroupedConvolutionForwardInvoker
                 GemmShape,
                 GemmUniversalTraits,
                 scheduler,
-                has_hot_loop_v,
-                tail_number_v,
                 ck_tile::element_wise::PassThrough,
                 ck_tile::element_wise::PassThrough,
                 OutDataType,
@@ -368,33 +332,17 @@ struct GroupedConvolutionForwardInvoker
         // =====================================================================
         if(use_split_image)
         {
-            const auto RunSplitImage = [&](const auto has_hot_loop_, const auto tail_number_) {
-                if(args.k_batch == 1)
-                    return Run(
-                        has_hot_loop_, tail_number_, MemoryOpSet{}, ck_tile::bool_constant<true>{});
-                else
-                    return Run(has_hot_loop_,
-                               tail_number_,
-                               MemoryOpAtomicAdd{},
-                               ck_tile::bool_constant<true>{});
-            };
-            return BaseGemmPipeline::TailHandler(RunSplitImage, has_hot_loop, tail_num);
+            if(args.k_batch == 1)
+                return Run(MemoryOpSet{}, ck_tile::bool_constant<true>{});
+            else
+                return Run(MemoryOpAtomicAdd{}, ck_tile::bool_constant<true>{});
         }
         else
         {
-            const auto RunRegular = [&](const auto has_hot_loop_, const auto tail_number_) {
-                if(args.k_batch == 1)
-                    return Run(has_hot_loop_,
-                               tail_number_,
-                               MemoryOpSet{},
-                               ck_tile::bool_constant<false>{});
-                else
-                    return Run(has_hot_loop_,
-                               tail_number_,
-                               MemoryOpAtomicAdd{},
-                               ck_tile::bool_constant<false>{});
-            };
-            return BaseGemmPipeline::TailHandler(RunRegular, has_hot_loop, tail_num);
+            if(args.k_batch == 1)
+                return Run(MemoryOpSet{}, ck_tile::bool_constant<false>{});
+            else
+                return Run(MemoryOpAtomicAdd{}, ck_tile::bool_constant<false>{});
         }
     }
 };