From b5d8a5d9ccb435268e2215144dca162b0fc6a826 Mon Sep 17 00:00:00 2001
From: Sergey Klevtsov <s.klevzoff@gmail.com>
Date: Wed, 15 Nov 2023 10:45:17 -0800
Subject: [PATCH] Allow SM90 pingpong kernel to use custom tile schedulers
 (#1194)

Co-authored-by: Sergey Klevtsov <sklevtsov@nvidia.com>
---
 .../gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp     | 3 +--
 1 file changed, 1 insertion(+), 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
index ba483c2b..b34e56f2 100644
--- 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 @@ public:
   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;