Use un-rolled gemm for Gemm-0

include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp

@@ -330,23 +330,14 @@ struct BlockFmhaPipelineQRKSVSAsync
             // ensure k is completely updated on LDS
             block_sync_lds();
 
-            // for kQKHeaddim == 96 (kSubQKHeaddim == 128), we need to use k0_loops
-            if constexpr(kQKHeaddim == kSubQKHeaddim)
-            {
-                gemm_0(s_acc, q, k_lds_window);
-            }
-            else
-            {
-
             static_for<0, k0_loops, 1>{}([&](auto i_k0) {
-                    gemm_0(s_acc,
-                           get_slice_tile(
-                               q, sequence<0, i_k0 * kK0>{}, sequence<kM0, (i_k0 + 1) * kK0>{}),
+                gemm_0(
+                    s_acc,
+                    get_slice_tile(q, sequence<0, i_k0 * kK0>{}, sequence<kM0, (i_k0 + 1) * kK0>{}),
                     get_slice_tile(k_lds_window,
                                    sequence<0, i_k0 * kK0>{},
                                    sequence<kN0, (i_k0 + 1) * kK0>{}));
             });
-            }
 
             __builtin_amdgcn_sched_barrier(0); // prevent from messing up the order of global loads
 

include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async_default_policy.hpp

@@ -57,6 +57,7 @@ struct BlockFmhaPipelineQRKSVSAsyncDefaultPolicy
                                        sequence<0, 1>>{});
     }
 
+    /*
         template <typename Problem>
         CK_TILE_HOST_DEVICE static constexpr auto GetQKBlockGemm()
         {
@@ -71,13 +72,14 @@ struct BlockFmhaPipelineQRKSVSAsyncDefaultPolicy
                 typename Problem::SaccDataType,
                 Problem::kNumGemm0Warps * get_warp_size(),
                 TileGemmShape<
-                sequence<Problem::BlockFmhaShape::kM0, Problem::BlockFmhaShape::kN0, BlockGemmK>,
-                typename Problem::BlockFmhaShape::Gemm0BlockWarps,
-                typename Problem::BlockFmhaShape::Gemm0WarpTile>>;
+                    sequence<Problem::BlockFmhaShape::kM0, Problem::BlockFmhaShape::kN0,
+       BlockGemmK>, typename Problem::BlockFmhaShape::Gemm0BlockWarps, typename
+       Problem::BlockFmhaShape::Gemm0WarpTile>>;
 
             constexpr auto warp_gemm = []() {
-            constexpr index_t WarpGemmM = Problem::BlockFmhaShape::Gemm0WarpTile::at(number<0>{});
-            static_assert(WarpGemmM == 4 || WarpGemmM == 16 || WarpGemmM == 32);
+                constexpr index_t WarpGemmM =
+       Problem::BlockFmhaShape::Gemm0WarpTile::at(number<0>{}); static_assert(WarpGemmM == 4 ||
+       WarpGemmM == 16 || WarpGemmM == 32);
 
                 if constexpr(std::is_same_v<typename Problem::QDataType, half_t> &&
                              std::is_same_v<typename Problem::KDataType, half_t> &&
@@ -118,14 +120,15 @@ struct BlockFmhaPipelineQRKSVSAsyncDefaultPolicy
                 BlockGemmARegBSmemCRegV2CustomPolicy<typename Problem::QDataType,
                                                      typename Problem::KDataType,
                                                      typename Problem::SaccDataType,
-                                                 typename Problem::BlockFmhaShape::Gemm0BlockWarps,
-                                                 decltype(warp_gemm)>;
+                                                     typename
+       Problem::BlockFmhaShape::Gemm0BlockWarps, decltype(warp_gemm)>;
 
             if constexpr(1 < Problem::kNumGemm0Warps)
                 return BlockGemmARegBSmemCRegV2<GemmProblem, BlockGemmPolicy>{};
             else
                 return BlockGemmARegBSmemCRegOneWarpV1<GemmProblem, BlockGemmPolicy>{};
         }
+    */
 
     template <typename Problem>
     CK_TILE_HOST_DEVICE static constexpr auto GetSmemKPackQ()