From 7024136ca1b0f2bd4bef6b054988fe7a5fcc764b Mon Sep 17 00:00:00 2001
From: OscarXu <huaiguxu@amd.com>
Date: Tue, 29 Apr 2025 16:09:49 +0800
Subject: [PATCH] Fix no-mainloop scenario

---
 .../impl/device_moe_gemm_blockscale.hpp       | 32 ++++++++++++-------
 1 file changed, 21 insertions(+), 11 deletions(-)

diff --git a/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm_blockscale.hpp b/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm_blockscale.hpp
index 2343185589..2a2d4116ca 100644
--- a/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm_blockscale.hpp
+++ b/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm_blockscale.hpp
@@ -317,17 +317,27 @@ struct DeviceMoeGemmBlockScale
 #if 1
             else
             {
-                // Tail number always 1
-                if constexpr(BlkGemmPipelineVer == BlockGemmPipelineVersion::v1)
-                {
-                    const auto kernel = kernel_moe_gemm<GridwiseGemm,
-                                                        true,
-                                                        InMemoryDataOperationEnum::Set,
-                                                        minimum_occupancy,
-                                                        IsInputGemm,
-                                                        TailNumber::Odd>;
-                    RunKernel(kernel);
-                }
+                    if(GridwiseGemm::CalculateKBlockLoopTailNum(K_split) == TailNumber::Odd)
+                    {
+                        printf("here to do odd.\n");
+                        const auto kernel = kernel_moe_gemm<GridwiseGemm,
+                                                            false,
+                                                            MemoryDataOp,
+                                                            minimum_occupancy,
+                                                            IsInputGemm,
+                                                            TailNumber::Odd>;
+                        RunKernel(kernel);
+                    }
+                    else
+                    {
+                        const auto kernel = kernel_moe_gemm<GridwiseGemm,
+                                                            false,
+                                                            MemoryDataOp,
+                                                            minimum_occupancy,
+                                                            IsInputGemm,
+                                                            TailNumber::Even>;
+                        RunKernel(kernel);
+                    }
             }
 #endif