Allow SM90 pingpong kernel to use custom tile schedulers (NVIDIA#1194)

Co-authored-by: Sergey Klevtsov <sklevtsov@nvidia.com>
WeiMa01 · Nov 15, 2023 · b5d8a5d · b5d8a5d
1 parent 6e60b9b
commit b5d8a5d
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp b/include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp
@@ -98,8 +98,7 @@ class GemmUniversal<
   using EpilogueArguments = typename CollectiveEpilogue::Arguments;
   using EpilogueParams = typename CollectiveEpilogue::Params;
 
-  static_assert(cute::is_void_v<TileScheduler_> or cute::is_same_v<TileScheduler_, PersistentScheduler>,
-    "Ping-pong kernel only supports the default scheduler.");
+  static_assert(!cute::is_same_v<TileScheduler_, StreamKScheduler>, "Ping-pong kernel does not currently support stream-K scheduler.");
   using TileSchedulerTag = TileScheduler_;
   using TileScheduler = typename detail::TileSchedulerSelector<
     TileScheduler_, ArchTag, TileShape, ClusterShape>::Scheduler;