From e4391834addced6c041b64141aae9d5eec50ef10 Mon Sep 17 00:00:00 2001
From: valarLip <340077269@qq.com>
Date: Fri, 26 Sep 2025 08:11:59 +0000
Subject: [PATCH] fix wp gemm bug when permuteN is false

---
 example/ck_tile/03_gemm/gemm_utils.hpp                     | 3 ++-
 example/ck_tile/03_gemm/gemm_weight_preshuffle_invoker.hpp | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/example/ck_tile/03_gemm/gemm_utils.hpp b/example/ck_tile/03_gemm/gemm_utils.hpp
index 588b66ca43..2315c39d29 100644
--- a/example/ck_tile/03_gemm/gemm_utils.hpp
+++ b/example/ck_tile/03_gemm/gemm_utils.hpp
@@ -72,6 +72,7 @@ struct GemmConfigBase
     static constexpr ck_tile::index_t Pipeline      = CK_TILE_PIPELINE_COMPUTE_V3;
     static constexpr ck_tile::index_t NumWaveGroups = 1;
     static constexpr bool Preshuffle                = false;
+    static constexpr bool TiledMMAPermuteN          = false;
 };
 
 template <typename PrecType>
@@ -300,7 +301,7 @@ struct GemmConfigPreshufflePrefill : public GemmConfigBase
     static constexpr bool Preshuffle           = true;
     static constexpr bool DoubleSmemBuffer     = true;
     static constexpr int N_Repeat              = N_Tile / N_Warp_Tile / N_Warp;
-    static constexpr bool TiledMMAPermuteN     = N_Repeat % 2 == 0;
+    static constexpr bool TiledMMAPermuteN     = N_Repeat % 4 == 0;
 };
 
 template <typename PrecType>
diff --git a/example/ck_tile/03_gemm/gemm_weight_preshuffle_invoker.hpp b/example/ck_tile/03_gemm/gemm_weight_preshuffle_invoker.hpp
index b47dd8d8a7..d737a0f864 100644
--- a/example/ck_tile/03_gemm/gemm_weight_preshuffle_invoker.hpp
+++ b/example/ck_tile/03_gemm/gemm_weight_preshuffle_invoker.hpp
@@ -109,7 +109,7 @@ struct WeightPreshuffleInvoker
                                                  GemmConfig::NumWaveGroups,
                                                  false,
                                                  1,
-                                                 true>>;
+                                                 GemmConfig::TiledMMAPermuteN>>;
             using Kernel = ck_tile::GemmKernel<TilePartitioner, GemmPipeline, GemmEpilogue>;
             auto kargs   = Kernel::MakeKernelArgs(args);