Fix the CK Tile related operators (#2356)

* fix the flatmm

* Fix the pipeline

* address the comment

include/ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5.hpp

@@ -1,5 +1,6 @@
 #include "ck_tile/core.hpp"
 #include "ck_tile/ops/gemm/pipeline/gemm_universal_pipeline_ag_bg_cr_policy.hpp"
+#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_comp_v5_default_policy.hpp"
 #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_scheduler.hpp"
 #include "ck_tile/ops/gemm/pipeline/gemm_pipeline_ag_bg_cr_base.hpp"
 #include "ck_tile/host/concat.hpp"

include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp

@@ -47,6 +47,8 @@ struct GemmPipelineAGmemBGmemCRegV1
     static constexpr bool kPadN = Problem::kPadN;
     static constexpr bool kPadK = Problem::kPadK;
 
+    static constexpr index_t NumWaveGroups = Problem::NumWaveGroups;
+
     static constexpr index_t kLdsAlignmentInBytes = 16;
 
     [[nodiscard]] CK_TILE_HOST static const std::string GetName()

include/ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp

@@ -32,6 +32,8 @@ struct GemmPipelineProblemBase
 
     static constexpr bool TransposeC = Traits::TransposeC;
 
+    static constexpr index_t NumWaveGroups = Traits::NumWaveGroups;
+
     static constexpr bool UseStructuredSparsity = Traits::UseStructuredSparsity;
 
     static constexpr index_t kBlockSize = BlockGemmShape::NumWarps * get_warp_size();

include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp

@@ -28,6 +28,7 @@ struct TileGemmTraits
 
     static constexpr bool TransposeC            = false;
     static constexpr bool UseStructuredSparsity = false;
+    static constexpr index_t NumWaveGroups      = 1;
 };
 
 template <bool kPadM_,