From ce51d2d0997a32e12d5be5acacfb759022415422 Mon Sep 17 00:00:00 2001
From: Your Name <you@example.com>
Date: Mon, 12 May 2025 16:24:30 +0800
Subject: [PATCH] Revert "use mem_op::set when topk=1"

This reverts commit def952a178bbb73e0940cf6a3cf69802e38b4dd7.
---
 .../gpu/device/impl/device_moe_gemm.hpp       | 171 +++++-------------
 1 file changed, 48 insertions(+), 123 deletions(-)

diff --git a/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm.hpp b/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm.hpp
index a6110d2bfc..08d177035e 100644
--- a/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm.hpp
+++ b/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm.hpp
@@ -264,152 +264,77 @@ struct DeviceMoeGemm : public DeviceGemmMultipleDSplitKBPreShuffle<ALayout,
 
             constexpr index_t minimum_occupancy = (estimated_reg_total >= 256) ? 1 : 2;
 
-            if(IsInputGemm || arg.TopK == 1)
+            constexpr auto MemoryDataOp =
+                IsInputGemm ? InMemoryDataOperationEnum::Set : InMemoryDataOperationEnum::AtomicAdd;
+            if(has_main_k_block_loop)
             {
-                constexpr auto MemoryDataOp = InMemoryDataOperationEnum::Set;
-
-                if(has_main_k_block_loop)
+                // Tail number always full
+                if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1)
                 {
-                    // Tail number always full
-                    if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1)
-                    {
-                        {
-                            if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd)
-                            {
-                                const auto kernel = kernel_moe_gemm<GridwiseGemm,
-                                                                    true,
-                                                                    MemoryDataOp,
-                                                                    minimum_occupancy,
-                                                                    TailNumber::Odd>;
-                                RunKernel(kernel);
-                            }
-                            else
-                            {
-                                const auto kernel = kernel_moe_gemm<GridwiseGemm,
-                                                                    true,
-                                                                    MemoryDataOp,
-                                                                    minimum_occupancy,
-                                                                    TailNumber::Even>;
-                                RunKernel(kernel);
-                            }
-                        }
-                    }
-                    else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v2 ||
-                                      BlkGemmPipelineVer == BlockGemmPipelineVersion::v3)
                     {
                         if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd)
                         {
-                            const auto kernel = kernel_moe_gemm_2lds<GridwiseGemm,
-                                                                     true,
-                                                                     MemoryDataOp,
-                                                                     minimum_occupancy,
-                                                                     TailNumber::Odd>;
+                            const auto kernel = kernel_moe_gemm<GridwiseGemm,
+                                                                true,
+                                                                MemoryDataOp,
+                                                                minimum_occupancy,
+                                                                TailNumber::Odd>;
                             RunKernel(kernel);
                         }
                         else
                         {
-                            const auto kernel = kernel_moe_gemm_2lds<GridwiseGemm,
-                                                                     true,
-                                                                     MemoryDataOp,
-                                                                     minimum_occupancy,
-                                                                     TailNumber::Even>;
+                            const auto kernel = kernel_moe_gemm<GridwiseGemm,
+                                                                true,
+                                                                MemoryDataOp,
+                                                                minimum_occupancy,
+                                                                TailNumber::Even>;
                             RunKernel(kernel);
                         }
                     }
+                }
+                else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v2 ||
+                                  BlkGemmPipelineVer == BlockGemmPipelineVersion::v3)
+                {
+                    if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd)
+                    {
+                        const auto kernel = kernel_moe_gemm_2lds<GridwiseGemm,
+                                                                 true,
+                                                                 MemoryDataOp,
+                                                                 minimum_occupancy,
+                                                                 TailNumber::Odd>;
+                        RunKernel(kernel);
+                    }
                     else
                     {
-                        throw std::runtime_error("todo: only v1 & v2 support now");
-                    }
-                }
-#if 1
-                else
-                {
-                    // Tail number always 1
-                    if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1)
-                    {
-                        const auto kernel = kernel_moe_gemm<GridwiseGemm,
-                                                            true,
-                                                            InMemoryDataOperationEnum::Set,
-                                                            minimum_occupancy,
-                                                            TailNumber::Odd>;
+                        const auto kernel = kernel_moe_gemm_2lds<GridwiseGemm,
+                                                                 true,
+                                                                 MemoryDataOp,
+                                                                 minimum_occupancy,
+                                                                 TailNumber::Even>;
                         RunKernel(kernel);
                     }
                 }
-#endif
+                else
+                {
+                    throw std::runtime_error("todo: only v1 & v2 support now");
+                }
             }
+#if 1
             else
             {
-                constexpr auto MemoryDataOp = InMemoryDataOperationEnum::AtomicAdd;
-
-                if(has_main_k_block_loop)
+                // Tail number always 1
+                if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1)
                 {
-                    // Tail number always full
-                    if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1)
-                    {
-                        {
-                            if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd)
-                            {
-                                const auto kernel = kernel_moe_gemm<GridwiseGemm,
-                                                                    true,
-                                                                    MemoryDataOp,
-                                                                    minimum_occupancy,
-                                                                    TailNumber::Odd>;
-                                RunKernel(kernel);
-                            }
-                            else
-                            {
-                                const auto kernel = kernel_moe_gemm<GridwiseGemm,
-                                                                    true,
-                                                                    MemoryDataOp,
-                                                                    minimum_occupancy,
-                                                                    TailNumber::Even>;
-                                RunKernel(kernel);
-                            }
-                        }
-                    }
-                    else if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v2 ||
-                                      BlkGemmPipelineVer == BlockGemmPipelineVersion::v3)
-                    {
-                        if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd)
-                        {
-                            const auto kernel = kernel_moe_gemm_2lds<GridwiseGemm,
-                                                                     true,
-                                                                     MemoryDataOp,
-                                                                     minimum_occupancy,
-                                                                     TailNumber::Odd>;
-                            RunKernel(kernel);
-                        }
-                        else
-                        {
-                            const auto kernel = kernel_moe_gemm_2lds<GridwiseGemm,
-                                                                     true,
-                                                                     MemoryDataOp,
-                                                                     minimum_occupancy,
-                                                                     TailNumber::Even>;
-                            RunKernel(kernel);
-                        }
-                    }
-                    else
-                    {
-                        throw std::runtime_error("todo: only v1 & v2 support now");
-                    }
+                    const auto kernel = kernel_moe_gemm<GridwiseGemm,
+                                                        true,
+                                                        InMemoryDataOperationEnum::Set,
+                                                        minimum_occupancy,
+                                                        TailNumber::Odd>;
+                    RunKernel(kernel);
                 }
-#if 1
-                else
-                {
-                    // Tail number always 1
-                    if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1)
-                    {
-                        const auto kernel = kernel_moe_gemm<GridwiseGemm,
-                                                            true,
-                                                            InMemoryDataOperationEnum::Set,
-                                                            minimum_occupancy,
-                                                            TailNumber::Odd>;
-                        RunKernel(kernel);
-                    }
-                }
-#endif
             }
+#endif
+
             return ave_time;
         }