CUTLASS 3.4.0 (#1286)

* CUTLASS 3.4.0

* Update CHANGELOG.md

---------

Co-authored-by: Pradeep Ramani <prramani@nvidia.com>

include/cutlass/gemm/kernel/gemm_streamk_with_fused_epilogue.h

@@ -851,7 +851,9 @@ protected:
       }
       ptr_D += tile_work.tiled_coord.k() * params.batch_stride_D;
       if (ptr_Tensor) {
-        ptr_Tensor += tile_work.tiled_coord.k() * params.batch_stride_Tensor;
+        ptr_Tensor = ReferenceFactory<typename Epilogue::ElementTensor>::add_pointer_offset(
+          ptr_Tensor,
+          tile_work.tiled_coord.k() * params.batch_stride_Tensor);
       }
       if (ptr_Vector) {
         ptr_Vector += tile_work.tiled_coord.k() * params.batch_stride_Vector;
@@ -2024,7 +2026,9 @@ protected:
       ptr_C += tile_work.tiled_coord.k() * params.batch_stride_C;
       ptr_D += tile_work.tiled_coord.k() * params.batch_stride_D;
       if (ptr_Tensor) {
-        ptr_Tensor += tile_work.tiled_coord.k() * params.batch_stride_Tensor;
+        ptr_Tensor = ReferenceFactory<typename Epilogue::ElementTensor>::add_pointer_offset(
+          ptr_Tensor,
+          tile_work.tiled_coord.k() * params.batch_stride_Tensor);
       }
       if (ptr_Vector) {
         ptr_Vector += tile_work.tiled_coord.k() * params.batch_stride_Vector;

include/cutlass/gemm/kernel/gemm_universal.h

@@ -67,9 +67,9 @@ class GemmUniversal<
   Epilogue_,
   ThreadblockSwizzle_,
   void,
-  // 3.x kernels use the first template argument to define the ProblemShape tuple
+  // 3.x kernels use the first template argument to define the ProblemShape
   // We use this invariant to SFINAE dispatch against either the 2.x API or the 3.x API
-  cute::enable_if_t<not cute::is_tuple<Mma_>::value>
+  cute::enable_if_t<not (cute::is_tuple<Mma_>::value || IsCutlass3ArrayKernel<Mma_>::value)>
 > {
 public:
 

include/cutlass/gemm/kernel/gemm_universal.hpp

@@ -61,6 +61,19 @@ template <
 >
 class GemmUniversal;
 
+
+////////////////////////////////////////////////////////////////////////////////
+
+// In cases where ProblemShape is not a tuple, this is used to check if the
+// underlying problem shape type is aliased within or not.
+// Used for dispatching GemmUniversal to 2.x API or 3.x API
+template <class ProblemShape, class = void>
+struct IsCutlass3ArrayKernel : cute::false_type { };
+
+template <typename ProblemShape>
+struct IsCutlass3ArrayKernel<ProblemShape, cute::void_t<typename ProblemShape::UnderlyingProblemShape>>
+    : cute::true_type { };
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace cutlass::gemm::kernel
@@ -75,4 +88,5 @@ class GemmUniversal;
 #include "cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp"
 #include "cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp"
 #include "cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp"
+#include "cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp"
 ////////////////////////////////////////////////////////////////////////////////

include/cutlass/gemm/kernel/gemm_with_fused_epilogue.h

@@ -42,7 +42,7 @@
 #include "cutlass/complex.h"
 #include "cutlass/semaphore.h"
 #include "cutlass/gemm/kernel/params_universal_base.h"
-
+#include "cutlass/subbyte_reference.h"
 #include "cutlass/trace.h"
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
@@ -676,7 +676,9 @@ public:
       }
       ptr_D += threadblock_tile_offset.k() * params.batch_stride_D;
       if (ptr_Tensor) {
-        ptr_Tensor += threadblock_tile_offset.k() * params.batch_stride_Tensor;
+        ptr_Tensor = ReferenceFactory<typename Epilogue::ElementTensor>::add_pointer_offset(
+          ptr_Tensor,
+          threadblock_tile_offset.k() * params.batch_stride_Tensor);
       }
       if (ptr_Vector) {
         ptr_Vector += threadblock_tile_offset.k() * params.batch_stride_Vector;
@@ -1387,7 +1389,9 @@ public:
       ptr_C += threadblock_tile_offset.k() * params.batch_stride_C;
       ptr_D += threadblock_tile_offset.k() * params.batch_stride_D;
       if (ptr_Tensor) {
-        ptr_Tensor += threadblock_tile_offset.k() * params.batch_stride_Tensor;
+        ptr_Tensor = ReferenceFactory<typename Epilogue::ElementTensor>::add_pointer_offset(
+          ptr_Tensor,
+          threadblock_tile_offset.k() * params.batch_stride_Tensor);
       }
       if (ptr_Vector) {
         ptr_Vector += threadblock_tile_offset.k() * params.batch_stride_Vector;

include/cutlass/gemm/kernel/rank_2k_grouped.h

@@ -217,6 +217,8 @@ public:
     // Only used by device-level operator
     GemmCoord *host_problem_sizes;
 
+    bool allow_early_exit;
+
     //
     // Methods
     //
@@ -235,7 +237,8 @@ public:
       ldb(nullptr),
       ldc(nullptr),
       ldd(nullptr),
-      host_problem_sizes(nullptr)
+      host_problem_sizes(nullptr),
+      allow_early_exit(false)
     {
 
     }
@@ -256,7 +259,8 @@ public:
       typename LayoutB::Stride::LongIndex *ldb,
       typename LayoutC::Stride::LongIndex *ldc,
       typename LayoutC::Stride::LongIndex *ldd,
-      GemmCoord *host_problem_sizes=nullptr
+      GemmCoord *host_problem_sizes=nullptr,
+      bool allow_early_exit=false
     ):
       mode(mode),
       problem_sizes(problem_sizes),
@@ -271,7 +275,8 @@ public:
       ldb(ldb),
       ldc(ldc),
       ldd(ldd),
-      host_problem_sizes(host_problem_sizes)
+      host_problem_sizes(host_problem_sizes),
+      allow_early_exit(allow_early_exit)
     {
 
     }
@@ -303,6 +308,7 @@ public:
     typename LayoutC::Stride::LongIndex *ldc;
     typename LayoutC::Stride::LongIndex *ldd;
 
+    bool allow_early_exit;
 
     //
     // Methods
@@ -318,7 +324,8 @@ public:
       lda(nullptr),
       ldb(nullptr),
       ldc(nullptr),
-      ldd(nullptr)
+      ldd(nullptr),
+      allow_early_exit(false)
     { }
 
     CUTLASS_HOST_DEVICE
@@ -333,7 +340,8 @@ public:
       lda(args.lda),
       ldb(args.ldb),
       ldc(args.ldc),
-      ldd(args.ldd)
+      ldd(args.ldd),
+      allow_early_exit(args.allow_early_exit)
     {
 
     }
@@ -388,6 +396,12 @@ public:
   CUTLASS_DEVICE
   void operator()(Params const &params, SharedStorage &shared_storage) {
 
+    // Early exit following LAPACK's definition
+    if (params.allow_early_exit &&
+        (params.output_op.alpha == ElementC(0)) && (params.output_op.beta == ElementC(1))) {
+      return;
+    }
+
     //
     // Problem visitor.
     //

include/cutlass/gemm/kernel/rank_2k_universal.h

@@ -140,6 +140,8 @@ public:
     typename LayoutC::Stride::Index ldc;
     typename LayoutC::Stride::Index ldd;
 
+    bool allow_early_exit;
+
     //
     // Methods
     //
@@ -147,7 +149,8 @@ public:
     Arguments(): 
       mode(GemmUniversalMode::kGemm), 
       batch_count(1), 
-      ptr_A(nullptr), ptr_B(nullptr), ptr_C(nullptr), ptr_D(nullptr) { }
+      ptr_A(nullptr), ptr_B(nullptr), ptr_C(nullptr), ptr_D(nullptr),
+      allow_early_exit(false) { }
 
     /// constructs an arguments structure
     Arguments(
@@ -166,7 +169,8 @@ public:
       typename LayoutA::Stride::Index lda,
       typename LayoutB::Stride::Index ldb,
       typename LayoutC::Stride::Index ldc,
-      typename LayoutC::Stride::Index ldd
+      typename LayoutC::Stride::Index ldd,
+      bool allow_early_exit = false
     ):
       mode(mode), 
       problem_size(problem_size), 
@@ -174,7 +178,8 @@ public:
       epilogue(epilogue), 
       ptr_A(ptr_A), ptr_B(ptr_B), ptr_C(ptr_C), ptr_D(ptr_D), 
       batch_stride_A(batch_stride_A), batch_stride_C(batch_stride_C), batch_stride_D(batch_stride_D), 
-      lda(lda), ldb(ldb), ldc(ldc), ldd(ldd) {
+      lda(lda), ldb(ldb), ldc(ldc), ldd(ldd),
+      allow_early_exit(allow_early_exit) {
 
       }
 
@@ -231,6 +236,8 @@ public:
 
     int *semaphore;
 
+    bool allow_early_exit;
+
     //
     // Methods
     //
@@ -255,7 +262,8 @@ public:
       batch_stride_B(0),
       batch_stride_C(0),
       batch_stride_D(0),
-      semaphore(nullptr) { }
+      semaphore(nullptr),
+      allow_early_exit(false) { }
 
     CUTLASS_HOST_DEVICE
     Params(
@@ -285,7 +293,8 @@ public:
       batch_stride_B(args.batch_stride_B),
       batch_stride_C(args.batch_stride_C),
       batch_stride_D(args.batch_stride_D),
-      semaphore(static_cast<int *>(workspace)) {
+      semaphore(static_cast<int *>(workspace)),
+      allow_early_exit(args.allow_early_exit) {
     }
 
     CUTLASS_HOST_DEVICE
@@ -347,6 +356,12 @@ public:
   CUTLASS_DEVICE
   void operator()(Params const &params, SharedStorage &shared_storage) {
 
+    // Early exit following LAPACK's definition
+    if (params.allow_early_exit &&
+        (params.output_op.alpha == ElementC(0)) && (params.output_op.beta == ElementC(1))) {
+      return;
+    }
+
     // Compute threadblock location
     ThreadblockSwizzle threadblock_swizzle;
 

include/cutlass/gemm/kernel/rank_k_universal.h

@@ -125,6 +125,8 @@ public:
     typename LayoutC::Stride::Index ldc;
     typename LayoutC::Stride::Index ldd;
 
+    bool allow_early_exit;
+
     //
     // Methods
     //
@@ -132,7 +134,8 @@ public:
     Arguments(): 
       mode(GemmUniversalMode::kGemm), 
       batch_count(1), 
-      ptr_A(nullptr), ptr_C(nullptr), ptr_D(nullptr) { }
+      ptr_A(nullptr), ptr_C(nullptr), ptr_D(nullptr),
+      allow_early_exit(false) { }
 
     /// constructs an arguments structure
     Arguments(
@@ -148,7 +151,8 @@ public:
       int64_t batch_stride_D,
       typename LayoutA::Stride::Index lda,
       typename LayoutC::Stride::Index ldc,
-      typename LayoutC::Stride::Index ldd
+      typename LayoutC::Stride::Index ldd,
+      bool allow_early_exit = false
     ):
       mode(mode), 
       problem_size(problem_size), 
@@ -156,7 +160,8 @@ public:
       epilogue(epilogue), 
       ptr_A(ptr_A), ptr_C(ptr_C), ptr_D(ptr_D), 
       batch_stride_A(batch_stride_A), batch_stride_C(batch_stride_C), batch_stride_D(batch_stride_D), 
-      lda(lda), ldb(ldb), ldc(ldc), ldd(ldd) {
+      lda(lda), ldb(ldb), ldc(ldc), ldd(ldd),
+      allow_early_exit(allow_early_exit) {
 
       }
 
@@ -196,6 +201,8 @@ public:
 
     int *semaphore;
 
+    bool allow_early_exit;
+
     //
     // Methods
     //
@@ -218,7 +225,8 @@ public:
       batch_stride_B(0),
       batch_stride_C(0),
       batch_stride_D(0),
-      semaphore(nullptr) { }
+      semaphore(nullptr),
+      allow_early_exit(false) { }
 
     CUTLASS_HOST_DEVICE
     Params(
@@ -246,7 +254,8 @@ public:
       batch_stride_B(args.batch_stride_A),
       batch_stride_C(args.batch_stride_C),
       batch_stride_D(args.batch_stride_D),
-      semaphore(static_cast<int *>(workspace)) {
+      semaphore(static_cast<int *>(workspace)),
+      allow_early_exit(args.allow_early_exit) {
     }
 
     CUTLASS_HOST_DEVICE
@@ -313,6 +322,12 @@ public:
     cutlass::gemm::GemmCoord threadblock_tile_offset =
         threadblock_swizzle.get_tile_offset(params.swizzle_log_tile);
 
+    // Early exit following LAPACK's definition
+    if (params.allow_early_exit &&
+        (params.output_op.alpha == ElementC(0)) && (params.output_op.beta == ElementC(1))) {
+      return;
+    }
+
     // Early exit if CTA is out of range
     if (params.grid_tiled_shape.m() <= threadblock_tile_offset.m() ||
       params.grid_tiled_shape.n() <= threadblock_tile_offset.n()) {

include/cutlass/gemm/kernel/sm70_gemm.hpp

@@ -60,7 +60,7 @@ public:
   //
   using ProblemShape = ProblemShape_;
 
-  static_assert(cute::rank(ProblemShape{}) == 3 or cute::rank(ProblemShape{}) == 4,
+  static_assert(rank(ProblemShape{}) == 3 or rank(ProblemShape{}) == 4,
     "ProblemShape{} should be <M,N,K> or <M,N,K,L>");
 
   // Mainloop derived types
@@ -101,7 +101,7 @@ public:
       sizeof(typename CollectiveMainloop::SharedStorage),
       sizeof(typename CollectiveEpilogue::SharedStorage)));
 
-  static constexpr uint32_t MaxThreadsPerBlock = cute::size(TiledMma{});
+  static constexpr uint32_t MaxThreadsPerBlock = CUTE_STATIC_V(cute::size(TiledMma{}));
   static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
 
   // Device side arguments
@@ -141,8 +141,9 @@ public:
 
   static bool
   can_implement(Arguments const& args) {
-    return args.mode == GemmUniversalMode::kGemm or
-          (args.mode == GemmUniversalMode::kBatched && cute::rank(ProblemShape{}) == 4);
+    bool mode_implementable = args.mode == GemmUniversalMode::kGemm or
+          (args.mode == GemmUniversalMode::kBatched && rank(ProblemShape{}) == 4);
+    return mode_implementable && TileScheduler::can_implement(args.scheduler);
   }
 
   static int

include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp

@@ -0,0 +1,749 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/workspace.h"
+#include "cutlass/fast_math.h"
+#include "cutlass/kernel_hardware_info.hpp"
+#include "cute/arch/cluster_sm90.hpp"
+#include "cutlass/arch/reg_reconfig.h"
+#include "cutlass/arch/mma_sm90.h"
+#include "cutlass/epilogue/collective/detail.hpp"
+#include "cutlass/gemm/gemm.h"
+#include "cutlass/gemm/dispatch_policy.hpp"
+#include "cutlass/gemm/group_array_problem_shape.hpp"
+#include "cutlass/gemm/kernel/tile_scheduler.hpp"
+#include "cutlass/pipeline/pipeline.hpp"
+#include "cute/tensor.hpp"
+#include "cutlass/trace.h"
+
+///////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass::gemm::kernel {
+
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+  class ProblemShape_,
+  class CollectiveMainloop_,
+  class CollectiveEpilogue_,
+  class TileScheduler_
+>
+class GemmUniversal<
+  ProblemShape_,
+  CollectiveMainloop_,
+  CollectiveEpilogue_,
+  TileScheduler_,
+  cute::enable_if_t<cute::is_base_of_v<KernelArrayTmaWarpSpecializedCooperative, typename CollectiveMainloop_::DispatchPolicy::Schedule> ||
+                    cute::is_base_of_v<KernelGroupTmaWarpSpecializedCooperative, typename CollectiveMainloop_::DispatchPolicy::Schedule>>
+>
+{
+public:
+  //
+  // Type Aliases
+  //
+  using ProblemShape = ProblemShape_;
+  static_assert(rank(typename ProblemShape::UnderlyingProblemShape{}) == 3 or rank(typename ProblemShape::UnderlyingProblemShape{}) == 4,
+    "ProblemShape{} should be <M,N,K> or <M,N,K,L>");
+
+  // Mainloop derived types
+  using CollectiveMainloop = CollectiveMainloop_;
+  using TileShape = typename CollectiveMainloop::TileShape;
+  using TiledMma  = typename CollectiveMainloop::TiledMma;
+  using ArchTag   = typename CollectiveMainloop::ArchTag;
+  using ElementA  = typename CollectiveMainloop::ElementA;
+  using StrideA   = typename CollectiveMainloop::StrideA;
+  using ElementB  = typename CollectiveMainloop::ElementB;
+  using StrideB   = typename CollectiveMainloop::StrideB;
+  using DispatchPolicy = typename CollectiveMainloop::DispatchPolicy;
+  using Schedule = typename DispatchPolicy::Schedule;
+  using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator;
+  using ClusterShape = typename DispatchPolicy::ClusterShape;
+  using MainloopArguments = typename CollectiveMainloop::Arguments;
+  using MainloopParams = typename CollectiveMainloop::Params;
+
+  // Epilogue derived types
+  using CollectiveEpilogue = CollectiveEpilogue_;
+  using ElementC = typename CollectiveEpilogue::ElementC;
+  using StrideC  = typename CollectiveEpilogue::StrideC;
+  using ElementD = typename CollectiveEpilogue::ElementD;
+  using StrideD  = typename CollectiveEpilogue::StrideD;
+  using EpilogueArguments = typename CollectiveEpilogue::Arguments;
+  using EpilogueParams = typename CollectiveEpilogue::Params;
+
+  static_assert(ArchTag::kMinComputeCapability >= 90);
+  static_assert(cute::is_void_v<TileScheduler_>,
+    "Ptr-Array Cooperative and Grouped Gemm Cooperative kernel only supports the default scheduler.");
+  
+  static constexpr bool IsGroupedGemmKernel = cute::is_base_of_v<KernelGroupTmaWarpSpecializedCooperative, Schedule>;
+
+  using TileScheduler = cute::conditional_t<IsGroupedGemmKernel,
+    typename detail::TileSchedulerSelector<
+      GroupScheduler, ArchTag,
+      TileShape, ClusterShape,
+      ProblemShape>::Scheduler,
+    typename detail::TileSchedulerSelector<
+    void, ArchTag, TileShape, ClusterShape>::Scheduler>;
+  using TileSchedulerArguments = typename TileScheduler::Arguments;
+  using TileSchedulerParams = typename TileScheduler::Params;
+
+  static constexpr uint32_t NumLoadWarpGroups = 1;
+  static constexpr uint32_t NumMmaWarpGroups = CUTE_STATIC_V(size(TiledMma{})) / NumThreadsPerWarpGroup;
+  static constexpr uint32_t MaxThreadsPerBlock = CUTE_STATIC_V(size(TiledMma{})) + (NumLoadWarpGroups * NumThreadsPerWarpGroup);
+  static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
+
+  /// Register requirement for Load and Math WGs
+  static constexpr uint32_t LoadRegisterRequirement = 40;
+  static constexpr uint32_t MmaRegisterRequirement = 232;
+
+  // 1 stage ordered sequence between mainloop and epilogue producer load threads
+  using LoadWarpOrderBarrier = cutlass::OrderedSequenceBarrier<1,2>;
+
+  // Kernel level shared memory storage
+  struct SharedStorage {
+    struct TensorStorage : cute::aligned_struct<128> {
+      using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage;
+      using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage;
+
+      MainloopTensorStorage mainloop;
+      EpilogueTensorStorage epilogue;
+    } tensors;
+
+    struct PipelineStorage : cute::aligned_struct<16> {
+      using MainloopPipelineStorage = typename CollectiveMainloop::PipelineStorage;
+      using EpiLoadPipelineStorage = typename CollectiveEpilogue::PipelineStorage;
+
+      alignas(16) MainloopPipelineStorage mainloop;
+      alignas(16) EpiLoadPipelineStorage epi_load;
+      alignas(16) typename LoadWarpOrderBarrier::SharedStorage load_order;
+    } pipelines;
+
+    struct TensorMapStorage : cute::aligned_struct<128> {
+      using MainloopTensorMapStorage = typename CollectiveMainloop::TensorMapStorage;
+      alignas(128) MainloopTensorMapStorage mainloop;
+    } tensormaps;
+  };
+
+  static constexpr int SharedStorageSize = sizeof(SharedStorage);
+
+  // Device side arguments
+  struct Arguments {
+    GemmUniversalMode mode{};
+    ProblemShape problem_shape{};
+    MainloopArguments mainloop{};
+    EpilogueArguments epilogue{};
+    KernelHardwareInfo hw_info{};
+    TileSchedulerArguments scheduler{};
+  };
+
+  // Kernel entry point API
+  struct Params {
+    GemmUniversalMode mode;
+    ProblemShape problem_shape;
+    MainloopParams mainloop;
+    EpilogueParams epilogue;
+    KernelHardwareInfo hw_info;
+    TileSchedulerParams scheduler;
+    void* workspace;
+  };
+
+  //
+  // Methods
+  //
+
+  // Convert to underlying arguments. In this case, a simple copy for the aliased type.
+  static
+  Params
+  to_underlying_arguments(Arguments const& args, void* workspace) {
+    CUTLASS_TRACE_HOST("to_underlying_arguments():");
+
+    ProblemShape problem_shapes = args.problem_shape;
+
+    // Get SM count if needed, otherwise use user supplied SM count
+    int sm_count = args.hw_info.sm_count;
+    if (sm_count <= 0) {
+      CUTLASS_TRACE_HOST("  WARNING: Arguments do not include a valid SM count.\n"
+          "  For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count.");
+      sm_count = KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id);
+    }
+
+    CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count);
+
+    KernelHardwareInfo hw_info{args.hw_info.device_id, sm_count};
+
+    // Calculate workspace pointers
+    uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
+    size_t workspace_offset = 0;
+
+    void* scheduler_workspace = workspace_ptr;
+    workspace_offset += TileScheduler::template get_workspace_size<typename ProblemShape::UnderlyingProblemShape, ElementAccumulator>(
+      args.scheduler, problem_shapes.get_host_problem_shape(0), args.hw_info, NumMmaWarpGroups);
+    workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
+
+    void* epilogue_workspace = workspace_ptr + workspace_offset;
+    workspace_offset += CollectiveEpilogue::get_workspace_size(problem_shapes, args.epilogue);
+    workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
+
+    void* mainloop_workspace = workspace_ptr + workspace_offset;
+    workspace_offset += CollectiveMainloop::get_workspace_size(problem_shapes, args.mainloop, args.hw_info.sm_count);
+    workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
+
+    // Precompute the sub tiles numbers in epilogue, pass into tile scheduler.  Therefore it will be used
+    // in separate reduction scheme for streamk case, NumEpilogueSubTiles default value is 1, which means
+    // subtile will not be used, therefore separate reduction will not be enabled.
+    constexpr uint32_t NumEpilogueSubTiles = CollectiveEpilogue::get_store_pipe_increment(TileShape{});
+    TileSchedulerParams scheduler;
+    if constexpr (IsGroupedGemmKernel) {
+      scheduler = TileScheduler::to_underlying_arguments(
+      problem_shapes, TileShape{}, ClusterShape{}, hw_info, args.scheduler, scheduler_workspace, NumEpilogueSubTiles);
+    }
+    else {
+      scheduler = TileScheduler::to_underlying_arguments(
+      problem_shapes.get_host_problem_shape(), TileShape{}, ClusterShape{}, hw_info, args.scheduler, scheduler_workspace, NumEpilogueSubTiles);
+    } 
+
+    return {
+      args.mode,
+      problem_shapes,
+      CollectiveMainloop::to_underlying_arguments(problem_shapes, args.mainloop, mainloop_workspace),
+      CollectiveEpilogue::to_underlying_arguments(problem_shapes, args.epilogue, epilogue_workspace),
+      hw_info,
+      scheduler,
+      workspace
+    };
+  }
+
+  CUTLASS_HOST_DEVICE static
+  bool
+  can_implement(Arguments const& args) {
+    bool implementable = true;
+    if constexpr (cute::is_base_of_v<KernelArrayTmaWarpSpecializedCooperative, Schedule>) {
+      implementable &= (args.mode == GemmUniversalMode::kArray && rank(typename ProblemShape::UnderlyingProblemShape{}) == 4);
+    } else if constexpr (IsGroupedGemmKernel) {
+      // Group GEMM currently only supports rank-3 problem shapes
+      implementable &= (args.mode == GemmUniversalMode::kGrouped && rank(typename ProblemShape::UnderlyingProblemShape{}) == 3);
+    }
+    else {
+      implementable = false;
+    }
+    if (!implementable) {
+      CUTLASS_TRACE_HOST("  CAN IMPLEMENT: Arguments or Problem Shape don't meet the requirements for Ptr Array Gemm or Grouped Gemm.\n");
+      return implementable;
+    }
+    implementable &= CollectiveMainloop::can_implement(args.problem_shape, args.mainloop);
+    implementable &= CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue);
+    implementable &= TileScheduler::can_implement(args.scheduler);
+    return implementable;
+  }
+
+  static size_t
+  get_workspace_size(Arguments const& args) {
+    size_t workspace_size = 0;
+    constexpr uint32_t NumEpilogueSubTiles = CollectiveEpilogue::get_store_pipe_increment(TileShape{});
+
+    workspace_size += TileScheduler::template get_workspace_size<typename ProblemShape::UnderlyingProblemShape, ElementAccumulator>(
+      args.scheduler, args.problem_shape.get_host_problem_shape(0), args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles);
+    workspace_size = round_nearest(workspace_size,  MinWorkspaceAlignment);
+
+    workspace_size += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
+    workspace_size = round_nearest(workspace_size,  MinWorkspaceAlignment);
+
+    // Get SM count if needed, otherwise use user supplied SM count
+    int sm_count = args.hw_info.sm_count;
+    if (sm_count <= 0) {
+      CUTLASS_TRACE_HOST("  WARNING: Arguments do not include a valid SM count.\n"
+          "  For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count.");
+      sm_count = KernelHardwareInfo::query_device_multiprocessor_count(args.hw_info.device_id);
+    }
+
+    workspace_size += CollectiveMainloop::get_workspace_size(args.problem_shape, args.mainloop, sm_count);
+    workspace_size = round_nearest(workspace_size,  MinWorkspaceAlignment);
+
+    return workspace_size;
+  }
+
+  static cutlass::Status
+  initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
+    Status status = Status::kSuccess;
+    uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
+    size_t workspace_offset = 0;
+    constexpr uint32_t NumEpilogueSubTiles = CollectiveEpilogue::get_store_pipe_increment(TileShape{});
+
+    status = TileScheduler::template initialize_workspace<typename ProblemShape::UnderlyingProblemShape, ElementAccumulator>(
+      args.scheduler, workspace_ptr + workspace_offset, stream, args.problem_shape.get_host_problem_shape(0), args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles);
+    workspace_offset += TileScheduler::template get_workspace_size<typename ProblemShape::UnderlyingProblemShape, ElementAccumulator>(
+      args.scheduler, args.problem_shape.get_host_problem_shape(0), args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles);
+    workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    status = CollectiveEpilogue::initialize_workspace(args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream);
+    workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
+    workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
+
+    status = CollectiveMainloop::initialize_workspace(args.problem_shape, args.mainloop, workspace_ptr + workspace_offset, stream);
+    workspace_offset += CollectiveMainloop::get_workspace_size(args.problem_shape, args.mainloop, args.hw_info.sm_count);
+    workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
+
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    return status;
+  }
+
+  // Computes the kernel launch grid shape based on runtime parameters
+  static dim3
+  get_grid_shape(Params const& params) {
+    // Given device SM count, set grid size s.t. we do not launch more thread blocks than we can run concurrently
+    TileSchedulerArguments args{};
+    if constexpr (!std::is_const_v<decltype(args.max_swizzle_size)>) {
+      args.max_swizzle_size = 1 << params.scheduler.log_swizzle_size_;
+    }
+    args.raster_order = params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN ? TileScheduler::RasterOrderOptions::AlongN : TileScheduler::RasterOrderOptions::AlongM;
+    dim3 grid_shape;
+    if constexpr (IsGroupedGemmKernel) {
+      grid_shape = TileScheduler::get_grid_shape(params.problem_shape, TileShape{}, ClusterShape{}, params.hw_info, args);
+    }
+    else {
+      grid_shape = TileScheduler::get_grid_shape(params.problem_shape.get_host_problem_shape(), TileShape{}, ClusterShape{}, params.hw_info, args);
+    } 
+    return grid_shape;
+  }
+
+  static dim3
+  get_block_shape() {
+    return dim3(MaxThreadsPerBlock, 1, 1);
+  }
+
+  CUTLASS_DEVICE
+  void
+  operator()(Params const& params, char* smem_buf) {
+    using namespace cute;
+    using X = Underscore;
+
+    // Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a.
+    #if ! defined(__CUDA_ARCH_FEAT_SM90_ALL)
+      if constexpr(size<0>(typename TiledMma::AtomShape_MNK{}) == 64) {
+        printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n");
+        return;
+      }
+    #endif
+
+    // Preconditions
+    static_assert(size(TiledMma{}) == 256, "Cooperative kernel must have TiledMMA operating using 256 threads.");
+    static_assert(size<0>(TileShape{}) >= 128,
+        "Cooperative kernel requires Tile Size to be greater than or equal to 128 along the M-dimension.");
+
+    static_assert(cute::rank(StrideA{}) == 3, "StrideA must be rank-3: [M, K, L]. If batch mode is not needed, set L stride to Int<0>.");
+    static_assert(cute::rank(StrideB{}) == 3, "StrideB must be rank-3: [N, K, L]. If batch mode is not needed, set L stride to Int<0>.");
+    static_assert(cute::rank(StrideC{}) == 3, "StrideC must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
+    static_assert(cute::rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
+
+    /* In the Cooperative kernel, Consumer0 and Consumer1 collaborate on the same tile */
+    enum class WarpGroupRole {
+      Producer = 0,
+      Consumer0 = 1,
+      Consumer1 = 2
+    };
+    enum class ProducerWarpRole {
+      Mainloop = 0,
+      Warp1 = 1,
+      Epilogue = 2,
+      Warp3 = 3
+    };
+
+    // Kernel level shared memory storage
+    SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
+
+    int thread_idx = int(threadIdx.x);
+    int lane_idx = canonical_lane_idx();
+    int warp_idx = canonical_warp_idx_sync();
+    int warp_idx_in_warp_group = warp_idx % NumWarpsPerWarpGroup;
+    int warp_group_thread_idx = thread_idx % NumThreadsPerWarpGroup;
+    int mma_thread_idx = thread_idx % size(TiledMma{});
+    auto warp_group_role = WarpGroupRole(canonical_warp_group_idx());
+    auto producer_warp_role = ProducerWarpRole(warp_idx_in_warp_group);
+    int lane_predicate = cute::elect_one_sync();
+    uint32_t block_rank_in_cluster = cute::block_rank_in_cluster();
+
+    // Note: Tma Descriptor Prefetch (from either const or param) is not applicable here
+
+    // Mainloop Load pipeline
+    using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
+    typename MainloopPipeline::Params mainloop_pipeline_params;
+    if (warp_group_role == WarpGroupRole::Producer && producer_warp_role == ProducerWarpRole::Mainloop) {
+      mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer;
+    }
+    if (warp_group_role == WarpGroupRole::Consumer0 || warp_group_role == WarpGroupRole::Consumer1) {
+      mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer;
+    }
+    mainloop_pipeline_params.is_leader = warp_group_thread_idx == 0;
+    mainloop_pipeline_params.num_consumers = size(TiledMma{});
+    mainloop_pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytes;
+    MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params, ClusterShape{});
+
+    // Epilogue Load pipeline
+    using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline;
+    typename EpiLoadPipeline::Params epi_load_pipeline_params;
+    if (warp_group_role == WarpGroupRole::Producer && producer_warp_role == ProducerWarpRole::Epilogue) {
+      epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Producer;
+    }
+    if (warp_group_role == WarpGroupRole::Consumer0 || warp_group_role == WarpGroupRole::Consumer1) {
+      epi_load_pipeline_params.role = EpiLoadPipeline::ThreadCategory::Consumer;
+    }
+    epi_load_pipeline_params.dst_blockid = cute::block_rank_in_cluster();
+    epi_load_pipeline_params.producer_arv_count = NumThreadsPerWarp;
+    epi_load_pipeline_params.consumer_arv_count = size(TiledMma{});
+    epi_load_pipeline_params.transaction_bytes = CollectiveEpilogue::TmaTransactionBytes;
+    EpiLoadPipeline epi_load_pipeline(shared_storage.pipelines.epi_load, epi_load_pipeline_params);
+
+    // Epilogue Store pipeline
+    using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline;
+    typename EpiStorePipeline::Params epi_store_pipeline_params;
+    epi_store_pipeline_params.always_wait = true;
+    EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params);
+
+    typename LoadWarpOrderBarrier::Params params_load_order_barrier;
+    params_load_order_barrier.group_id = producer_warp_role == ProducerWarpRole::Mainloop ? 0 : 1;
+    params_load_order_barrier.group_size = NumThreadsPerWarp;
+    LoadWarpOrderBarrier load_order_barrier(shared_storage.pipelines.load_order, params_load_order_barrier);
+
+    // Initialize starting pipeline states for the collectives
+    // Epilogue store pipe is producer-only (consumer is TMA unit, waits via scoreboarding)
+    typename CollectiveMainloop::PipelineState mainloop_pipe_consumer_state;
+    typename CollectiveEpilogue::LoadPipelineState epi_load_pipe_consumer_state;
+    // Purpose of maintaining this pipeline state is to make sure TMA loads have finished before doing descriptor updates
+    typename CollectiveMainloop::PipelineState mainloop_pipe_tma_consumer_state;
+
+    // For the DMA Load (producer) we start with an opposite phase
+    // i.e., we skip all waits since we know that the buffer is indeed empty
+    PipelineState mainloop_pipe_producer_state = cutlass::make_producer_start_state<MainloopPipeline>();
+    PipelineState epi_load_pipe_producer_state = cutlass::make_producer_start_state<EpiLoadPipeline>();
+    PipelineState epi_store_pipe_producer_state = cutlass::make_producer_start_state<EpiStorePipeline>();
+
+    auto cluster_wait_fn = [] () {
+      // We need this to guarantee that the Pipeline init is visible
+      // To all producers and consumer thread blocks in the Cluster
+      if constexpr (size(ClusterShape{}) > 1) {
+        cute::cluster_arrive_relaxed();
+        return [] () { cute::cluster_wait(); };
+      }
+      else {
+        __syncthreads();
+        return [] () {}; // do nothing
+      }
+    } ();
+
+    // Get the appropriate blocks for this thread block -- potential for thread block locality
+    TiledMma tiled_mma;
+    auto blk_shape = TileShape{};                                                                // (BLK_M,BLK_N,BLK_K)
+
+    TileScheduler scheduler{params.scheduler};
+    auto work_tile_info = scheduler.get_current_work();
+
+    // Optionally append 1s until problem shape is rank-4 in case it is only rank-3 (MNK)
+    auto problem_shape_MNKL = append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), Int<1>{});
+
+    // In a warp specialized kernel, collectives expose data movement and compute operations separately
+    CollectiveMainloop collective_mainloop;
+    CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
+
+    // Prepare and partition the input tensors. Expects a tuple of tensors where:
+    // get<0>(load_inputs) is the tma tensor A after local tiling so that it has shape (BLK_M,BLK_K,m,k,l)
+    // get<1>(load_inputs) is the tma tensor B after local tiling so that it has shape (BLK_N,BLK_K,n,k,l)
+    auto load_inputs = collective_mainloop.load_init(problem_shape_MNKL, params.mainloop);
+    static_assert(cute::tuple_size_v<decltype(load_inputs)> >= 2, "Output of load_init must have at least two elements (A, B)");
+
+    // Extract out partitioned A and B.
+    Tensor gA_mkl = get<0>(load_inputs);
+    Tensor gB_nkl = get<1>(load_inputs);
+
+    // Get pipeline stage increments from tensor shapes
+    auto k_tile_count = size<3>(gA_mkl);
+
+    // Wait for all thread blocks in the Cluster
+    cluster_wait_fn();
+
+    if (warp_group_role == WarpGroupRole::Producer) {
+      cutlass::arch::warpgroup_reg_dealloc<LoadRegisterRequirement>();
+
+      // Mainloop Producer Warp
+      if (producer_warp_role == ProducerWarpRole::Mainloop) {
+        int32_t curr_batch = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl)); // Usually just returns work_tile_info.L_idx;
+        int32_t next_batch = curr_batch;
+        int32_t const mock_l_coord = 0;
+        int32_t const sm_idx = blockIdx.x + (blockIdx.y * gridDim.x);
+        int32_t const sm_count = params.hw_info.sm_count;
+
+        // Fetch a copy of tensormaps for the CTA
+        auto input_tensormaps = collective_mainloop.tensormaps_init(params.mainloop, sm_count, sm_idx);
+
+        // Update tensormap for the initial batch for the CTA
+        if (work_tile_info.is_valid()) {
+          collective_mainloop.tensormaps_perform_update(
+            shared_storage.tensormaps.mainloop,
+            params.mainloop,
+            input_tensormaps,
+            problem_shape_MNKL,
+            next_batch
+          );
+          // Ensure warp is converged before issuing tensor replace
+          __syncwarp();
+          // Entire warp must do this (ie its aligned)
+          collective_mainloop.tensormaps_cp_fence_release(shared_storage.tensormaps.mainloop, input_tensormaps);
+        }
+
+        bool do_load_order_arrive = true;
+        while (work_tile_info.is_valid()) {
+          if (!TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) {
+            work_tile_info = fetch_next_work(work_tile_info, scheduler);
+            continue;
+          }
+
+          // Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape
+          auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl));
+          auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl));
+          auto blk_coord = make_coord(m_coord, n_coord, _, mock_l_coord);
+
+          // Get the number of K tiles to compute for this work as well as the starting K tile offset of the work.
+          auto work_k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, blk_shape);
+          auto work_k_tile_start = TileScheduler::get_work_k_tile_start(work_tile_info);
+          auto k_tile_iter = cute::make_coord_iterator(idx2crd(work_k_tile_start, shape<3>(gA_mkl)), shape<3>(gA_mkl));
+
+          collective_mainloop.tensormaps_fence_acquire(input_tensormaps);
+
+          collective_mainloop.load(
+            params.mainloop,
+            mainloop_pipeline,
+            mainloop_pipe_producer_state,
+            load_inputs,
+            input_tensormaps,
+            blk_coord,
+            k_tile_iter, work_k_tile_count,
+            lane_idx,
+            block_rank_in_cluster,
+            shared_storage.tensors.mainloop
+          );
+          // Update starting pipeline state for the next tile
+          mainloop_pipe_producer_state.advance(work_k_tile_count);
+
+          // Signal for the epilogue load warp to begin
+          if (do_load_order_arrive) {
+            load_order_barrier.arrive();
+            do_load_order_arrive = false;
+          }
+
+          // Get next work tile
+          work_tile_info = fetch_next_work(work_tile_info, scheduler);
+          next_batch = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl)); // Usually just returns work_tile_info.L_idx
+
+          if (work_tile_info.is_valid() && next_batch != curr_batch ) {
+            if constexpr (IsGroupedGemmKernel) {
+              problem_shape_MNKL = append<4>(params.problem_shape.get_problem_shape(next_batch), Int<1>{});
+            }
+            // Wait for the last TMA stage to complete loading, before issuing tensormap updates
+            mainloop_pipe_tma_consumer_state.advance(work_k_tile_count-1);
+            mainloop_pipeline.consumer_wait(mainloop_pipe_tma_consumer_state);
+            collective_mainloop.tensormaps_perform_update(
+              shared_storage.tensormaps.mainloop,
+              params.mainloop,
+              input_tensormaps,
+              problem_shape_MNKL,
+              next_batch
+            );
+            // Ensure warp is converged before issuing tensor replace
+            __syncwarp();
+            // Entire warp must do this (ie its aligned)
+            collective_mainloop.tensormaps_cp_fence_release(shared_storage.tensormaps.mainloop, input_tensormaps);
+            curr_batch = next_batch;
+            // Advance the TMA consumer state for the last remaining stage that was being waited for above
+            mainloop_pipe_tma_consumer_state.advance(1);
+          } 
+          else if (work_tile_info.is_valid()) { // case where batch/group didn't change between tiles
+            // Advance the TMA consumer state for all the stages to be in sync
+            mainloop_pipe_tma_consumer_state.advance(work_k_tile_count);
+          }
+        } // Scheduler work fetch loop
+
+        // Make sure all Consumer Warp Groups have been waited upon
+        collective_mainloop.load_tail(mainloop_pipeline, mainloop_pipe_producer_state);
+      } // Mainloop Producer Warp End
+
+      // Epilogue Producer Warp
+      else if (producer_warp_role == ProducerWarpRole::Epilogue && collective_epilogue.is_producer_load_needed()) {
+        while (work_tile_info.is_valid()) {
+          if (!TileScheduler::requires_separate_reduction(params.scheduler)) {
+            load_order_barrier.wait();
+          }
+          if (TileScheduler::compute_epilogue(work_tile_info, params.scheduler)) {
+            // Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape
+            auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl));
+            auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl));
+            auto l_coord = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl));
+            auto blk_coord = make_coord(m_coord, n_coord, _, l_coord);
+
+            epi_load_pipe_producer_state =
+            collective_epilogue.load(
+              epi_load_pipeline,
+              epi_load_pipe_producer_state,
+              problem_shape_MNKL,
+              blk_shape,
+              blk_coord,
+              tiled_mma,
+              lane_idx,
+              shared_storage.tensors.epilogue,
+              work_tile_info.reduction_subtile_idx()
+            );
+          }
+
+          // Get next work tile
+          work_tile_info = fetch_next_work(work_tile_info, scheduler);
+          if constexpr (IsGroupedGemmKernel) {
+            problem_shape_MNKL = append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), Int<1>{});
+          }
+        } // Scheduler work fetch loop
+
+        // Make sure all Consumer Warp Groups have been waited upon
+        collective_epilogue.load_tail(epi_load_pipeline, epi_load_pipe_producer_state);
+      } // Epilogue Producer Warp End
+    } // Producer Warp Group End
+
+    else if (warp_group_role == WarpGroupRole::Consumer0 || warp_group_role == WarpGroupRole::Consumer1) {
+      cutlass::arch::warpgroup_reg_alloc<MmaRegisterRequirement>();
+
+      // Do we potentially issue tail arrives for TMA stores, if epilogue load is waiting for it
+      bool do_store_tail = false;
+      while (work_tile_info.is_valid()) {
+        // Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape
+        auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl));
+        auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl));
+        auto l_coord = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl));
+        auto blk_coord = make_coord(m_coord, n_coord, _, l_coord);
+        auto work_k_tile_count = TileScheduler::get_work_k_tile_count(work_tile_info, problem_shape_MNKL, blk_shape);
+
+        // Allocate the accumulators for the (M,N) blk_shape
+        //
+        // MSVC CTAD breaks if we say "Tensor" here, so we use "auto" instead.
+        auto accumulators = partition_fragment_C(tiled_mma, take<0,2>(blk_shape));               // (MMA,MMA_M,MMA_N)
+        if(TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) {
+          collective_mainloop.mma(
+            mainloop_pipeline,
+            mainloop_pipe_consumer_state,
+            accumulators,
+            work_k_tile_count,
+            mma_thread_idx,
+            shared_storage.tensors.mainloop,
+            params.mainloop
+          );
+
+          // Make sure the math instructions are done and free buffers before entering the epilogue
+          collective_mainloop.mma_tail(
+            mainloop_pipeline,
+            mainloop_pipe_consumer_state,
+            work_k_tile_count
+          );
+
+          // Update starting mainloop pipeline state for the next tile
+          mainloop_pipe_consumer_state.advance(work_k_tile_count);
+        }
+        // Index of warp group within consumer warp groups
+        int consumer_warp_group_idx = canonical_warp_group_idx() - NumLoadWarpGroups;
+
+        // Perform reduction across splits, if needed
+        TileScheduler::fixup(
+          params.scheduler, work_tile_info, accumulators, NumMmaWarpGroups, consumer_warp_group_idx);
+
+        if (TileScheduler::compute_epilogue(work_tile_info, params.scheduler)) {
+          // Epilogue and write to gD
+          auto [epi_load_pipe_consumer_state_next, epi_store_pipe_producer_state_next] =
+          collective_epilogue.store(
+            epi_load_pipeline,
+            epi_load_pipe_consumer_state,
+            epi_store_pipeline,
+            epi_store_pipe_producer_state,
+            problem_shape_MNKL,
+            blk_shape,
+            blk_coord,
+            accumulators,
+            tiled_mma,
+            mma_thread_idx,
+            shared_storage.tensors.epilogue,
+            work_tile_info.reduction_subtile_idx()
+          );
+          epi_load_pipe_consumer_state = epi_load_pipe_consumer_state_next;
+          epi_store_pipe_producer_state = epi_store_pipe_producer_state_next;
+          do_store_tail = true;
+        }
+
+        // Get next work tile
+        work_tile_info = fetch_next_work(work_tile_info, scheduler);
+        if constexpr (IsGroupedGemmKernel) {
+          problem_shape_MNKL = append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), Int<1>{});
+        }
+      } // Scheduler work fetch loop
+
+      if (do_store_tail) {
+        collective_epilogue.store_tail(
+          epi_load_pipeline,
+          epi_load_pipe_consumer_state,
+          epi_store_pipeline,
+          epi_store_pipe_producer_state
+        );
+      }
+    } // Consumer Warp Groups End
+  }
+
+private:
+  // Kernel helper function to get next work unit
+  CUTLASS_DEVICE
+  typename TileScheduler::WorkTileInfo
+  fetch_next_work(
+    typename TileScheduler::WorkTileInfo& work_tile_info,
+    TileScheduler& scheduler) const {
+    // Check whether we should continue on with the current work unit. If this is the case,
+    // the work unit will have been updated in continue_current_work to reflect the new
+    // tile to be computed.
+    if (scheduler.continue_current_work(work_tile_info)) {
+      return work_tile_info;
+    }
+
+    // Get next work tile
+    scheduler.advance_to_next_work();
+    return scheduler.get_current_work();
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+
+} // namespace cutlass::gemm::kernel

include/cutlass/gemm/kernel/sm90_gemm_tma.hpp

@@ -121,7 +121,7 @@ public:
       sizeof(typename CollectiveMainloop::SharedStorage),
       sizeof(typename CollectiveEpilogue::SharedStorage)));
 
-  static constexpr uint32_t MaxThreadsPerBlock = size(TiledMma{});
+  static constexpr uint32_t MaxThreadsPerBlock = CUTE_STATIC_V(size(TiledMma{}));
   static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
 
   // Device side arguments
@@ -176,6 +176,8 @@ public:
     }
     implementable &= CollectiveMainloop::can_implement(args.problem_shape, args.mainloop);
     implementable &= CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue);
+    implementable &= TileScheduler::can_implement(args.scheduler);
+
     return implementable;
   }
 

include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized.hpp

@@ -128,7 +128,7 @@ public:
 
   static constexpr uint32_t NumLoadWarpGroups = 1;
   static constexpr uint32_t NumMmaWarpGroups = 1;
-  static constexpr uint32_t MaxThreadsPerBlock = size(TiledMma{}) + (NumLoadWarpGroups * NumThreadsPerWarpGroup);
+  static constexpr uint32_t MaxThreadsPerBlock = CUTE_STATIC_V(size(TiledMma{})) + (NumLoadWarpGroups * NumThreadsPerWarpGroup);
   static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
 
   // Device side arguments
@@ -183,6 +183,8 @@ public:
     }
     implementable &= CollectiveMainloop::can_implement(args.problem_shape, args.mainloop);
     implementable &= CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue);
+    implementable &= TileScheduler::can_implement(args.scheduler);
+
     return implementable;
   }
 
@@ -270,7 +272,7 @@ public:
     mainloop_pipeline_params.is_leader = warp_group_thread_idx == 0;
     mainloop_pipeline_params.num_consumers = NumThreadsPerWarpGroup;
     mainloop_pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytes;
-    MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params);
+    MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params, ClusterShape{});
 
     // Epilogue Load pipeline
     using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline;
@@ -335,14 +337,14 @@ public:
     CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
 
     // Prepare and partition the input tensors. Expects a tuple of tensors where:
-    // get<0>(tiled_tensors) is the tma tensor A after local tiling so that it has shape (BLK_M,BLK_K,m,k,l)
-    // get<1>(tiled_tensors) is the tma tensor B after local tiling so that it has shape (BLK_N,BLK_K,n,k,l)
-    auto tiled_tensors = collective_mainloop.tile_input_tensors(problem_shape_MNKL, params.mainloop, blk_shape);
-    static_assert(cute::tuple_size_v<decltype(tiled_tensors)> >= 2, "Output of tile_input_tensors must have at least two elements (A, B)");
+    // get<0>(load_inputs) is the tma tensor A after local tiling so that it has shape (BLK_M,BLK_K,m,k,l)
+    // get<1>(load_inputs) is the tma tensor B after local tiling so that it has shape (BLK_N,BLK_K,n,k,l)
+    auto load_inputs = collective_mainloop.load_init(problem_shape_MNKL, params.mainloop);
+    static_assert(cute::tuple_size_v<decltype(load_inputs)> >= 2, "Output of load_init must have at least two elements (A, B)");
 
     // Extract out partitioned A and B.
-    Tensor gA_mkl = get<0>(tiled_tensors);
-    Tensor gB_nkl = get<1>(tiled_tensors);
+    Tensor gA_mkl = get<0>(load_inputs);
+    Tensor gB_nkl = get<1>(load_inputs);
 
     // Compute m_coord, n_coord, and l_coord with their post-tiled shapes
     auto m_coord = idx2crd(int(blockIdx.x), shape<2>(gA_mkl));
@@ -363,7 +365,7 @@ public:
           params.mainloop,
           mainloop_pipeline,
           mainloop_pipe_producer_state,
-          tiled_tensors,
+          load_inputs,
           blk_coord,
           k_tile_iter, k_tile_count,
           lane_idx,
@@ -378,8 +380,7 @@ public:
         if (collective_epilogue.is_producer_load_needed()) {
           // Ensure warp is converged before issuing epilogue loads
           __syncwarp();
-          epi_load_pipe_producer_state =
-          collective_epilogue.load(
+          epi_load_pipe_producer_state = collective_epilogue.load(
             epi_load_pipeline,
             epi_load_pipe_producer_state,
             problem_shape_MNKL,

include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp

@@ -105,8 +105,8 @@ public:
   using TileSchedulerParams = typename TileScheduler::Params;
 
   static constexpr uint32_t NumLoadWarpGroups = 1;
-  static constexpr uint32_t NumMmaWarpGroups = size(TiledMma{}) / NumThreadsPerWarpGroup;
-  static constexpr uint32_t MaxThreadsPerBlock = size(TiledMma{}) + (NumLoadWarpGroups * NumThreadsPerWarpGroup);
+  static constexpr uint32_t NumMmaWarpGroups = CUTE_STATIC_V(size(TiledMma{})) / NumThreadsPerWarpGroup;
+  static constexpr uint32_t MaxThreadsPerBlock = CUTE_STATIC_V(size(TiledMma{})) + (NumLoadWarpGroups * NumThreadsPerWarpGroup);
   static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
 
   /// Register requirement for Load and Math WGs
@@ -203,6 +203,12 @@ public:
     workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
 
     void* mainloop_workspace = nullptr;
+    // Precompute the sub tiles numbers in epilogue, pass into tile scheduler.  Therefore it will be used
+    // in separate reduction scheme for streamk case, NumEpilogueSubTiles default value is 1, which means
+    // subtile will not be used, therefore separate reduction will not be enabled.
+    constexpr uint32_t NumEpilogueSubTiles = CollectiveEpilogue::get_store_pipe_increment(TileShape{});
+    TileSchedulerParams scheduler = TileScheduler::to_underlying_arguments(
+      problem_shape_MNKL, TileShape{}, ClusterShape{}, hw_info, args.scheduler, scheduler_workspace, NumEpilogueSubTiles);
 
     return {
       args.mode,
@@ -210,7 +216,7 @@ public:
       CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, mainloop_workspace),
       CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, epilogue_workspace),
       hw_info,
-      TileScheduler::to_underlying_arguments(problem_shape_MNKL, TileShape{}, ClusterShape{}, hw_info, args.scheduler, scheduler_workspace),
+      scheduler,
       workspace
     };
   }
@@ -226,14 +232,17 @@ public:
     }
     implementable &= CollectiveMainloop::can_implement(args.problem_shape, args.mainloop);
     implementable &= CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue);
+    implementable &= TileScheduler::can_implement(args.scheduler);
     return implementable;
   }
 
   static size_t
   get_workspace_size(Arguments const& args) {
     size_t workspace_size = 0;
+    constexpr uint32_t NumEpilogueSubTiles = CollectiveEpilogue::get_store_pipe_increment(TileShape{});
+
     workspace_size += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
-      args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups);
+      args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles);
     workspace_size = round_nearest(workspace_size,  MinWorkspaceAlignment);
 
     workspace_size += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
@@ -247,11 +256,12 @@ public:
     Status status = Status::kSuccess;
     uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
     size_t workspace_offset = 0;
+    constexpr uint32_t NumEpilogueSubTiles = CollectiveEpilogue::get_store_pipe_increment(TileShape{});
 
     status = TileScheduler::template initialize_workspace<ProblemShape, ElementAccumulator>(
-      args.scheduler, workspace_ptr + workspace_offset, stream, args.problem_shape, args.hw_info, NumMmaWarpGroups);
+      args.scheduler, workspace_ptr + workspace_offset, stream, args.problem_shape, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles);
     workspace_offset += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
-      args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups);
+      args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles);
     workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
     if (status != Status::kSuccess) {
       return status;
@@ -353,7 +363,7 @@ public:
     mainloop_pipeline_params.is_leader = warp_group_thread_idx == 0;
     mainloop_pipeline_params.num_consumers = size(TiledMma{});
     mainloop_pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytes;
-    MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params);
+    MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params, ClusterShape{});
 
     // Epilogue Load pipeline
     using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline;
@@ -392,7 +402,7 @@ public:
     PipelineState epi_load_pipe_producer_state = cutlass::make_producer_start_state<EpiLoadPipeline>();
     PipelineState epi_store_pipe_producer_state = cutlass::make_producer_start_state<EpiStorePipeline>();
 
-    auto cluster_wait_fn = [&] () {
+    auto cluster_wait_fn = [] () {
       // We need this to guarantee that the Pipeline init is visible
       // To all producers and consumer thread blocks in the Cluster
       if constexpr (size(ClusterShape{}) > 1) {
@@ -420,14 +430,14 @@ public:
     CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
 
     // Prepare and partition the input tensors. Expects a tuple of tensors where:
-    // get<0>(tiled_tensors) is the tma tensor A after local tiling so that it has shape (BLK_M,BLK_K,m,k,l)
-    // get<1>(tiled_tensors) is the tma tensor B after local tiling so that it has shape (BLK_N,BLK_K,n,k,l)
-    auto tiled_tensors = collective_mainloop.tile_input_tensors(problem_shape_MNKL, params.mainloop, blk_shape);
-    static_assert(cute::tuple_size_v<decltype(tiled_tensors)> >= 2, "Output of tile_input_tensors must have at least two elements (A, B)");
+    // get<0>(load_inputs) is the tma tensor A after local tiling so that it has shape (BLK_M,BLK_K,m,k,l)
+    // get<1>(load_inputs) is the tma tensor B after local tiling so that it has shape (BLK_N,BLK_K,n,k,l)
+    auto load_inputs = collective_mainloop.load_init(problem_shape_MNKL, params.mainloop);
+    static_assert(cute::tuple_size_v<decltype(load_inputs)> >= 2, "Output of load_init must have at least two elements (A, B)");
 
     // Extract out partitioned A and B.
-    Tensor gA_mkl = get<0>(tiled_tensors);
-    Tensor gB_nkl = get<1>(tiled_tensors);
+    Tensor gA_mkl = get<0>(load_inputs);
+    Tensor gB_nkl = get<1>(load_inputs);
 
     // Get pipeline stage increments from tensor shapes
     auto k_tile_count = size<3>(gA_mkl);
@@ -442,6 +452,11 @@ public:
       if (producer_warp_role == ProducerWarpRole::Mainloop) {
         bool do_load_order_arrive = true;
         while (work_tile_info.is_valid()) {
+          if (!TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) {
+            work_tile_info = fetch_next_work(work_tile_info, scheduler);
+            continue;
+          }
+
           // Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape
           auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl));
           auto n_coord = idx2crd(work_tile_info.N_idx, shape<2>(gB_nkl));
@@ -457,7 +472,7 @@ public:
             params.mainloop,
             mainloop_pipeline,
             mainloop_pipe_producer_state,
-            tiled_tensors,
+            load_inputs,
             blk_coord,
             k_tile_iter, work_k_tile_count,
             lane_idx,
@@ -483,8 +498,10 @@ public:
 
       // Epilogue Producer Warp
       else if (producer_warp_role == ProducerWarpRole::Epilogue && collective_epilogue.is_producer_load_needed()) {
-        load_order_barrier.wait();
         while (work_tile_info.is_valid()) {
+          if (!TileScheduler::requires_separate_reduction(params.scheduler)) {
+            load_order_barrier.wait();
+          }
           if (TileScheduler::compute_epilogue(work_tile_info, params.scheduler)) {
             // Compute m_coord, n_coord, l_coord with the post-tiled m-shape and n-shape
             auto m_coord = idx2crd(work_tile_info.M_idx, shape<2>(gA_mkl));
@@ -501,7 +518,8 @@ public:
               blk_coord,
               tiled_mma,
               lane_idx,
-              shared_storage.tensors.epilogue
+              shared_storage.tensors.epilogue,
+              work_tile_info.reduction_subtile_idx()
             );
           }
 
@@ -531,27 +549,27 @@ public:
         //
         // MSVC CTAD breaks if we say "Tensor" here, so we use "auto" instead.
         auto accumulators = partition_fragment_C(tiled_mma, take<0,2>(blk_shape));               // (MMA,MMA_M,MMA_N)
+        if(TileScheduler::valid_warpgroup_in_work_tile(work_tile_info)) {
+          collective_mainloop.mma(
+            mainloop_pipeline,
+            mainloop_pipe_consumer_state,
+            accumulators,
+            work_k_tile_count,
+            mma_thread_idx,
+            shared_storage.tensors.mainloop,
+            params.mainloop
+          );
 
-        collective_mainloop.mma(
-          mainloop_pipeline,
-          mainloop_pipe_consumer_state,
-          accumulators,
-          work_k_tile_count,
-          mma_thread_idx,
-          shared_storage.tensors.mainloop,
-          params.mainloop
-        );
-
-        // Make sure the math instructions are done and free buffers before entering the epilogue
-        collective_mainloop.mma_tail(
-          mainloop_pipeline,
-          mainloop_pipe_consumer_state,
-          work_k_tile_count
-        );
-
-        // Update starting mainloop pipeline state for the next tile
-        mainloop_pipe_consumer_state.advance(work_k_tile_count);
+          // Make sure the math instructions are done and free buffers before entering the epilogue
+          collective_mainloop.mma_tail(
+            mainloop_pipeline,
+            mainloop_pipe_consumer_state,
+            work_k_tile_count
+          );
 
+          // Update starting mainloop pipeline state for the next tile
+          mainloop_pipe_consumer_state.advance(work_k_tile_count);
+        }
         // Index of warp group within consumer warp groups
         int consumer_warp_group_idx = canonical_warp_group_idx() - NumLoadWarpGroups;
 
@@ -573,7 +591,8 @@ public:
             accumulators,
             tiled_mma,
             mma_thread_idx,
-            shared_storage.tensors.epilogue
+            shared_storage.tensors.epilogue,
+            work_tile_info.reduction_subtile_idx()
           );
           epi_load_pipe_consumer_state = epi_load_pipe_consumer_state_next;
           epi_store_pipe_producer_state = epi_store_pipe_producer_state_next;

include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp

@@ -107,7 +107,7 @@ public:
 
   static constexpr uint32_t NumLoadWarpGroups = 1;
   static constexpr uint32_t NumMmaWarpGroups = 2;
-  static constexpr uint32_t MaxThreadsPerBlock = size(TiledMma{}) + (NumMmaWarpGroups * NumThreadsPerWarpGroup);
+  static constexpr uint32_t MaxThreadsPerBlock = CUTE_STATIC_V(size(TiledMma{})) + (NumMmaWarpGroups * NumThreadsPerWarpGroup);
   static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
 
   /// Register requirement for Load and Math WGs
@@ -232,6 +232,8 @@ public:
     }
     implementable &= CollectiveMainloop::can_implement(args.problem_shape, args.mainloop);
     implementable &= CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue);
+    implementable &= TileScheduler::can_implement(args.scheduler);
+
     return implementable;
   }
 
@@ -324,7 +326,7 @@ public:
 
     // Kernel level shared memory storage
     SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
-    
+
     int thread_idx = int(threadIdx.x);
     int lane_idx = canonical_lane_idx();
     int warp_idx = canonical_warp_idx_sync();
@@ -353,7 +355,7 @@ public:
     mainloop_pipeline_params.is_leader = warp_group_thread_idx == 0;
     mainloop_pipeline_params.num_consumers = NumThreadsPerWarpGroup;
     mainloop_pipeline_params.transaction_bytes = CollectiveMainloop::TmaTransactionBytes;
-    MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params);
+    MainloopPipeline mainloop_pipeline(shared_storage.pipelines.mainloop, mainloop_pipeline_params, ClusterShape{});
 
     // Epilogue Load pipeline
     using EpiLoadPipeline = typename CollectiveEpilogue::LoadPipeline;
@@ -424,14 +426,14 @@ public:
     CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
 
     // Prepare and partition the input tensors. Expects a tuple of tensors where:
-    // get<0>(tiled_tensors) is the tma tensor A after local tiling so that it has shape (BLK_M,BLK_K,m,k,l)
-    // get<1>(tiled_tensors) is the tma tensor B after local tiling so that it has shape (BLK_N,BLK_K,n,k,l)
-    auto tiled_tensors = collective_mainloop.tile_input_tensors(problem_shape_MNKL, params.mainloop, blk_shape);
-    static_assert(cute::tuple_size_v<decltype(tiled_tensors)> >= 2, "Output of tile_input_tensors must have at least two elements (A, B)");
+    // get<0>(load_inputs) is the tma tensor A after local tiling so that it has shape (BLK_M,BLK_K,m,k,l)
+    // get<1>(load_inputs) is the tma tensor B after local tiling so that it has shape (BLK_N,BLK_K,n,k,l)
+    auto load_inputs = collective_mainloop.load_init(problem_shape_MNKL, params.mainloop);
+    static_assert(cute::tuple_size_v<decltype(load_inputs)> >= 2, "Output of load_init must have at least two elements (A, B)");
 
     // Extract out partitioned A and B.
-    Tensor gA_mkl = get<0>(tiled_tensors);
-    Tensor gB_nkl = get<1>(tiled_tensors);
+    Tensor gA_mkl = get<0>(load_inputs);
+    Tensor gB_nkl = get<1>(load_inputs);
 
     // Get pipeline stage increments from tensor shapes
     auto k_tile_count = size<3>(gA_mkl);
@@ -472,7 +474,7 @@ public:
             params.mainloop,
             mainloop_pipeline,
             mainloop_pipe_producer_state,
-            tiled_tensors,
+            load_inputs,
             blk_coord,
             k_tile_iter, k_tile_count,
             lane_idx,

include/cutlass/gemm/kernel/sm90_gemm_warpspecialized.hpp

@@ -187,6 +187,8 @@ public:
     }
     implementable &= CollectiveMainloop::can_implement(args.problem_shape, args.mainloop);
     implementable &= CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue);
+    implementable &= TileScheduler::can_implement(args.scheduler);
+
     return implementable;
   }
 

include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_cooperative.hpp

@@ -207,6 +207,8 @@ public:
     }
     implementable &= CollectiveMainloop::can_implement(args.problem_shape, args.mainloop);
     implementable &= CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue);
+    implementable &= TileScheduler::can_implement(args.scheduler);
+
     return implementable;
   }
 

include/cutlass/gemm/kernel/sm90_gemm_warpspecialized_pingpong.hpp

@@ -219,6 +219,8 @@ public:
     }
     implementable &= CollectiveMainloop::can_implement(args.problem_shape, args.mainloop);
     implementable &= CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue);
+    implementable &= TileScheduler::can_implement(args.scheduler);
+
     return implementable;
   }
 

include/cutlass/gemm/kernel/sm90_tile_scheduler.hpp

@@ -76,6 +76,12 @@ public:
     is_final_split(uint32_t k_tiles_per_output_tile) const {
       return true;
     }
+
+    CUTLASS_HOST_DEVICE
+    int32_t
+    reduction_subtile_idx() const {
+      return -1;
+    }
   };
 
   using Params = PersistentTileSchedulerSm90Params;
@@ -101,7 +107,8 @@ public:
     ClusterShape cluster_shape,
     [[maybe_unused]] KernelHardwareInfo const& hw_info,
     Arguments const& arguments,
-    [[maybe_unused]] void* workspace=nullptr) {
+    [[maybe_unused]] void* workspace=nullptr,
+    [[maybe_unused]] const uint32_t epilogue_subtile = 1) {
 
     // We only need the tile and cluster shape during scheduler setup, so let FTAD do the magic
     static_assert(cute::is_static<TileShape>::value);
@@ -114,13 +121,19 @@ public:
       problem_blocks,
       to_gemm_coord(cluster_shape),
       hw_info,
-      arguments.max_swizzle_size, 
+      arguments.max_swizzle_size,
       arguments.raster_order
     );
 
     return params;
   }
 
+  CUTLASS_HOST_DEVICE
+  static bool
+  can_implement(Arguments const& args) {
+    return true;
+  }
+
   CUTLASS_HOST_DEVICE
   PersistentTileSchedulerSm90() { };
 
@@ -164,7 +177,7 @@ public:
                                                          scheduler_params.divmod_cluster_shape_major_,
                                                          scheduler_params.divmod_cluster_shape_minor_,
                                                          scheduler_params.divmod_cluster_blk_major_,
-                                                         scheduler_params.log_swizzle_size_, 
+                                                         scheduler_params.log_swizzle_size_,
                                                          scheduler_params.raster_order_);
 
     return {work_idx_m, work_idx_n, static_cast<int32_t>(work_idx_l), true};
@@ -180,11 +193,11 @@ public:
   static CUTLASS_DEVICE
   cute::tuple<int32_t, int32_t>
   get_work_idx_m_and_n(
-      uint64_t blk_per_grid_dim, 
+      uint64_t blk_per_grid_dim,
       FastDivmodU64Pow2 const& divmod_cluster_shape_major,
       FastDivmodU64Pow2 const& divmod_cluster_shape_minor,
       FastDivmodU64 const& divmod_cluster_blk_major,
-      int32_t log_swizzle_size, 
+      int32_t log_swizzle_size,
       RasterOrder raster_order) {
 
     uint64_t cluster_id, cluster_major_offset = 0, cluster_minor_offset = 0;
@@ -199,26 +212,26 @@ public:
     }
 
     uint64_t cluster_idx_minor, cluster_idx_major;
-    
+
     uint64_t cluster_idx_minor_div_swizzle, extra, offset;
 
     offset = cluster_id & ((1 << log_swizzle_size) - 1);
     extra = cluster_id >> log_swizzle_size;
-    
+
     divmod_cluster_blk_major(cluster_idx_minor_div_swizzle, cluster_idx_major, extra);
 
     cluster_idx_minor = cluster_idx_minor_div_swizzle * (1 << log_swizzle_size) + offset;
 
-    auto minor_work_idx = static_cast<int32_t>(cluster_idx_minor * divmod_cluster_shape_minor.divisor + 
+    auto minor_work_idx = static_cast<int32_t>(cluster_idx_minor * divmod_cluster_shape_minor.divisor +
                                                cluster_minor_offset);
-    auto major_work_idx = static_cast<int32_t>(cluster_idx_major * divmod_cluster_shape_major.divisor + 
+    auto major_work_idx = static_cast<int32_t>(cluster_idx_major * divmod_cluster_shape_major.divisor +
                                                cluster_major_offset);
 
     if (raster_order == RasterOrder::AlongN) {
       return {minor_work_idx, major_work_idx};
     }
     else {
-      return {major_work_idx, minor_work_idx}; 
+      return {major_work_idx, minor_work_idx};
     }
 
   }
@@ -331,13 +344,14 @@ public:
   // The basic tile scheduler does not require any additional workspace
   template <class ProblemShape, class ElementAccumulator>
   static int
-  get_workspace_size(Arguments const&, ProblemShape, KernelHardwareInfo const&, uint32_t) {
+  get_workspace_size(Arguments const&, ProblemShape, KernelHardwareInfo const&, uint32_t, const uint32_t = 1) {
     return 0;
   }
 
   template <class ProblemShape, class ElementAccumulator>
   static cutlass::Status
-  initialize_workspace(Arguments const&, void*, cudaStream_t, ProblemShape, KernelHardwareInfo const&, uint32_t) {
+  initialize_workspace(Arguments const&, void*, cudaStream_t, ProblemShape, KernelHardwareInfo const&,
+    uint32_t, const uint32_t = 1) {
     return Status::kSuccess;
   }
 
@@ -353,8 +367,61 @@ public:
   CUTLASS_HOST_DEVICE
   static uint32_t
   get_work_k_tile_start(WorkTileInfo const&) {
-    // All work units returned by this scheduler start from K tile 0
-    return 0u;
+  // All work units returned by this scheduler start from K tile 0
+  return 0u;
+  }
+
+  CUTLASS_DEVICE
+  static bool
+  need_separate_reduction(Params const& params) {
+    return false;
+  }
+
+  CUTLASS_DEVICE
+  bool
+  is_work_tile_for_reduction(WorkTileInfo const& work_tile_info, Params const& params) {
+    return false;
+  }
+
+  CUTLASS_DEVICE
+  uint32_t
+  epilgoue_subtile_idx(WorkTileInfo const& work_tile_info, Params const& params) const {
+    return 0;
+  }
+
+  template <class FrgTensorC>
+  CUTLASS_DEVICE
+  void
+  separate_reduction(
+    Params const& params,
+    WorkTileInfo const& work_tile_info,
+    FrgTensorC& accumulators,
+    uint32_t num_barriers,
+    uint32_t barrier_idx) {
+  }
+
+  // Shares the accumulator set with peers in the global workspace
+  template <class FrgTensorC>
+  CUTLASS_DEVICE
+  static void
+  share(
+    Params const& params,
+    WorkTileInfo const& work_tile_info,
+    FrgTensorC& accumulators,
+    uint32_t num_barriers,
+    uint32_t barrier_idx) {
+  }
+
+  CUTLASS_DEVICE
+  static bool
+  valid_warpgroup_in_work_tile(WorkTileInfo const& work_tile_info) {
+    return true;
+  }
+
+  CUTLASS_DEVICE
+  static bool
+  requires_separate_reduction(Params const& params) {
+    return false;
   }
 };
 

include/cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp

@@ -0,0 +1,431 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the documentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "cutlass/fast_math.h"
+#include "cutlass/gemm_coord.hpp"
+#include "cutlass/kernel_hardware_info.hpp"
+#include "cutlass/gemm/kernel/tile_scheduler_params.h"
+#include "cute/layout.hpp"
+#include "cute/tensor.hpp"
+#include "cute/arch/cluster_sm90.hpp"
+
+namespace cutlass::gemm::kernel::detail {
+
+///////////////////////////////////////////////////////////////////////////////
+
+// Persistent Thread Block (TB) scheduler
+template <class GroupProblemShape>
+class PersistentTileSchedulerSm90Group {
+  //
+  // Data members
+  //
+
+private:
+  uint64_t current_work_linear_idx_ = 0;
+  uint64_t total_grid_size_ = 0;
+
+  // Tracking current group, its starting linear idx and total tiles
+  struct GroupInfo {
+    uint64_t group = 0;
+    uint64_t start_linear_idx = 0;
+    uint64_t total_tiles = 0;
+  } current_group_info_;
+
+public:
+  struct WorkTileInfo {
+    int32_t M_idx = 0;
+    int32_t N_idx = 0;
+    int32_t L_idx = 0;
+    bool is_valid_tile = false;
+
+    CUTLASS_HOST_DEVICE
+    bool
+    is_valid() const {
+      return is_valid_tile;
+    }
+
+    CUTLASS_HOST_DEVICE
+    static WorkTileInfo
+    invalid_work_tile() {
+      return {-1, -1, -1, false};
+    }
+
+    CUTLASS_HOST_DEVICE
+    bool
+    is_final_split(uint32_t k_tiles_per_output_tile) const {
+      return true;
+    }
+
+    CUTLASS_HOST_DEVICE
+    int32_t
+    reduction_subtile_idx() const {
+      return -1;
+    }
+  };
+
+  using ProblemShape = typename GroupProblemShape::UnderlyingProblemShape;
+  using Params = PersistentTileSchedulerSm90GroupParams<ProblemShape>;
+  using RasterOrder = typename Params::RasterOrder;
+  using RasterOrderOptions = typename Params::RasterOrderOptions;
+  struct Arguments {
+    int max_swizzle_size = 1;
+    // Not applying Heuristics for Grouped problems, since largest dimension can change per group
+    RasterOrderOptions raster_order = RasterOrderOptions::AlongM;
+  };
+
+  // Sink scheduler params as a member
+  Params scheduler_params;
+
+  //
+  // Methods
+  //
+
+  template <class TileShape, class ClusterShape>
+  static Params
+  to_underlying_arguments(
+    GroupProblemShape problem_shapes,
+    TileShape tile_shape,
+    ClusterShape cluster_shape,
+    [[maybe_unused]] KernelHardwareInfo const& hw_info,
+    Arguments const& arguments,
+    [[maybe_unused]] void* workspace=nullptr,
+    [[maybe_unused]] const uint32_t epilogue_subtile = 1) {
+
+    // We only need the tile and cluster shape during scheduler setup, so let FTAD do the magic
+    static_assert(cute::is_static<TileShape>::value);
+    static_assert(cute::is_static<ClusterShape>::value);
+
+    dim3 problem_blocks = get_tiled_cta_shape_mnl(
+      problem_shapes.groups(),
+      reinterpret_cast<ProblemShape const*>(problem_shapes.host_problem_shapes),
+      tile_shape, cluster_shape);
+
+    Params params;
+    params.initialize(
+      problem_blocks,
+      problem_shapes.groups(),
+      reinterpret_cast<ProblemShape*>(problem_shapes.problem_shapes),
+      to_gemm_coord(tile_shape),
+      to_gemm_coord(cluster_shape),
+      hw_info,
+      arguments.max_swizzle_size, 
+      arguments.raster_order
+    );
+
+    return params;
+  }
+
+  CUTLASS_HOST_DEVICE
+  static bool
+  can_implement(Arguments const& args) {
+    return true;
+  }
+
+  PersistentTileSchedulerSm90Group() = default;
+
+  CUTLASS_DEVICE explicit PersistentTileSchedulerSm90Group(Params const& params_) : scheduler_params(params_) {
+    // MSVC requires protecting use of CUDA-specific nonstandard syntax,
+    // like blockIdx and gridDim, with __CUDA_ARCH__.
+#if defined(__CUDA_ARCH__)
+    if (params_.raster_order_ == RasterOrder::AlongN) {
+      current_work_linear_idx_ = uint64_t(blockIdx.x) + uint64_t(blockIdx.y) * uint64_t(gridDim.x);
+    }
+    else {
+      current_work_linear_idx_ = uint64_t(blockIdx.x) * uint64_t(gridDim.y) + uint64_t(blockIdx.y);
+    }
+
+    total_grid_size_ = uint64_t(gridDim.x) * uint64_t(gridDim.y) * uint64_t(gridDim.z);
+
+    auto cta_m = cute::size(cute::ceil_div(cute::shape<0>(params_.problem_shapes_[0]), params_.cta_shape_.m()));
+    auto cta_n = cute::size(cute::ceil_div(cute::shape<1>(params_.problem_shapes_[0]), params_.cta_shape_.n()));
+    current_group_info_.total_tiles = cta_m * cta_n;
+#else
+    CUTLASS_ASSERT(false && "This line should never be reached");
+#endif
+  }
+
+  CUTLASS_DEVICE
+  WorkTileInfo
+  get_current_work() {
+    return get_current_work_for_linear_idx(current_work_linear_idx_);
+  }
+
+  CUTLASS_DEVICE
+  WorkTileInfo
+  get_current_work_for_linear_idx(uint64_t linear_idx) {
+    if (linear_idx >= scheduler_params.blocks_per_problem_) {
+      return WorkTileInfo::invalid_work_tile();
+    }
+
+    uint64_t blk_per_grid_dim = scheduler_params.divmod_cluster_shape_minor_.divide(linear_idx);
+
+    auto [work_idx_m, work_idx_n, new_group_info, valid_tile] = get_work_idx_m_and_n(blk_per_grid_dim,
+                                                         current_group_info_,
+                                                         scheduler_params.groups_,
+                                                         scheduler_params.problem_shapes_,
+                                                         scheduler_params.cta_shape_,
+                                                         scheduler_params.divmod_cluster_shape_major_,
+                                                         scheduler_params.divmod_cluster_shape_minor_,
+                                                         scheduler_params.log_swizzle_size_, 
+                                                         scheduler_params.raster_order_);
+
+    current_group_info_ = new_group_info;
+    return {work_idx_m, work_idx_n, static_cast<int>(current_group_info_.group), valid_tile};
+  }
+
+  CUTLASS_DEVICE
+  void
+  advance_to_next_work(uint32_t advance_count = 1) {
+    current_work_linear_idx_ += total_grid_size_ * uint64_t(advance_count);
+  }
+
+  // get work_idx_m, work_idx_n from blk_per_grid_dim while applying swizzle
+  static CUTLASS_DEVICE
+  cute::tuple<int32_t, int32_t, struct GroupInfo, bool>
+  get_work_idx_m_and_n(
+      uint64_t blk_per_grid_dim,
+      struct GroupInfo group_info,
+      int32_t total_problem_groups,
+      ProblemShape* problem_shapes,
+      GemmCoord cta_shape,
+      FastDivmodU64Pow2 const& divmod_cluster_shape_major,
+      FastDivmodU64Pow2 const& divmod_cluster_shape_minor,
+      int32_t log_swizzle_size, 
+      RasterOrder raster_order) {
+
+    bool valid_tile = true;
+    int cta_m = cute::size(cute::ceil_div(cute::shape<0>(problem_shapes[group_info.group]), cta_shape.m()));
+    int cta_n = cute::size(cute::ceil_div(cute::shape<1>(problem_shapes[group_info.group]), cta_shape.n()));
+
+    while (group_info.start_linear_idx + group_info.total_tiles <= blk_per_grid_dim) {
+      group_info.group++;
+      group_info.start_linear_idx += group_info.total_tiles;
+      cta_m = cute::size(cute::ceil_div(cute::shape<0>(problem_shapes[group_info.group]), cta_shape.m()));
+      cta_n = cute::size(cute::ceil_div(cute::shape<1>(problem_shapes[group_info.group]), cta_shape.n()));
+      group_info.total_tiles = cta_m * cta_n;
+    }
+
+    uint64_t cluster_id, cluster_major_offset = 0, cluster_minor_offset = 0;
+    divmod_cluster_shape_major(cluster_id, cluster_major_offset, blk_per_grid_dim - group_info.start_linear_idx);
+
+    auto [cta_m_in_cluster, cta_n_in_cluster, _] = cute::block_id_in_cluster();
+    if (raster_order == RasterOrder::AlongN) {
+      cluster_minor_offset = cta_m_in_cluster;
+    }
+    else {
+      cluster_minor_offset = cta_n_in_cluster;
+    }
+
+    uint64_t cluster_idx_minor, cluster_idx_major;
+    
+    uint64_t cluster_idx_minor_div_swizzle, extra, offset;
+
+    offset = cluster_id & ((1 << log_swizzle_size) - 1);
+    extra = cluster_id >> log_swizzle_size;
+
+    uint64_t curr_group_cluster_blk_major, remainder;
+    divmod_cluster_shape_major(curr_group_cluster_blk_major, remainder, cta_m);
+    cluster_idx_minor_div_swizzle = extra / curr_group_cluster_blk_major;
+    cluster_idx_major = extra % curr_group_cluster_blk_major;
+
+    cluster_idx_minor = cluster_idx_minor_div_swizzle * (1 << log_swizzle_size) + offset;
+
+    auto minor_work_idx = static_cast<int32_t>(cluster_idx_minor * divmod_cluster_shape_minor.divisor + 
+                                               cluster_minor_offset);
+    auto major_work_idx = static_cast<int32_t>(cluster_idx_major * divmod_cluster_shape_major.divisor + 
+                                               cluster_major_offset);
+
+    if (raster_order == RasterOrder::AlongN) {
+      return {minor_work_idx, major_work_idx, group_info, valid_tile};
+    }
+    else {
+      return {major_work_idx, minor_work_idx, group_info, valid_tile}; 
+    }
+
+  }
+
+  // Given the inputs, computes the total number of output blocks this problem will compute over
+  // Note that this is only the logical size of our grid, not the physical grid we will actually launch.
+  template<class BlockShape, class ClusterShape>
+  CUTLASS_HOST_DEVICE static
+  dim3
+  get_tiled_cta_shape_mnl(int groups, ProblemShape const* problem_shapes, BlockShape cta_shape, ClusterShape cluster_shape) {
+    uint32_t total_ctas = 0;
+    uint32_t cta_in_N_dim = 1; // We linearize the blocks across all the problems here
+    for (int group = 0; group < groups; group++) {
+      auto cta_m = cute::size(cute::ceil_div(cute::shape<0>(problem_shapes[group]), cute::shape<0>(cta_shape)));
+      auto cta_n = cute::size(cute::ceil_div(cute::shape<1>(problem_shapes[group]), cute::shape<1>(cta_shape)));
+      total_ctas += cta_m * cta_n;
+    }
+
+    return Params::get_tiled_cta_shape_mnl(
+      to_gemm_coord(cluster_shape),
+      total_ctas, cta_in_N_dim
+    );
+  }
+
+  // Given the inputs, computes the physical grid we should launch.
+  template<class BlockShape, class ClusterShape>
+  CUTLASS_HOST_DEVICE static
+  dim3
+  get_grid_shape(
+    GroupProblemShape problem_shapes,
+    BlockShape cta_shape,
+    ClusterShape cluster_shape,
+    KernelHardwareInfo hw_info,
+    Arguments arguments,
+    bool truncate_by_problem_size=true) {
+
+    dim3 problem_blocks = get_tiled_cta_shape_mnl(
+      problem_shapes.groups(),
+      reinterpret_cast<ProblemShape const*>(problem_shapes.host_problem_shapes),
+      cta_shape, cluster_shape);
+
+    return Params::get_grid_shape(
+      problem_blocks,
+      to_gemm_coord(cluster_shape),
+      hw_info,
+      arguments.max_swizzle_size,
+      arguments.raster_order,
+      /* truncate_by_problem_size = */true
+    );
+  }
+
+  // Returns whether the block assigned this work should compute the epilogue for the corresponding
+  // output tile. For the basic tile scheduler, this is always true.
+  CUTLASS_HOST_DEVICE
+  static bool
+  compute_epilogue(WorkTileInfo const&, Params const&) {
+    return true;
+  }
+
+  // Performs the reduction across splits for a given output tile. Since this scheduler does
+  // not split output tiles, no reduction is needed.
+  template <class FrgTensorC>
+  CUTLASS_DEVICE
+  static void
+  fixup(Params const&, WorkTileInfo const&, FrgTensorC&, uint32_t, uint32_t) {}
+
+  // Returns whether the current WorkTileInfo passed in should continue to be used. Since
+  // this scheduler only schedules work in units of single, full output tiles, the WorkTileInfo
+  // passed in should not be used after having been processed.
+  CUTLASS_DEVICE
+  static bool
+  continue_current_work(WorkTileInfo&) {
+    return false;
+  }
+
+  // The basic tile scheduler does not require any additional workspace
+  template <class ProblemShape, class ElementAccumulator>
+  static int
+  get_workspace_size(Arguments const&, ProblemShape, KernelHardwareInfo const&, uint32_t, const uint32_t = 1) {
+    return 0;
+  }
+
+  template <class ProblemShape, class ElementAccumulator>
+  static cutlass::Status
+  initialize_workspace(Arguments const&, void*, cudaStream_t, ProblemShape, KernelHardwareInfo const&,
+    uint32_t, const uint32_t = 1) {
+    return Status::kSuccess;
+  }
+
+  template <class ProblemShape_MNKL, class TileShape>
+  CUTLASS_HOST_DEVICE
+  static int
+  get_work_k_tile_count(WorkTileInfo const& work_tile_info, ProblemShape_MNKL problem_shape, TileShape tile_shape) {
+    // All work units returned by this scheduler cover the entire K iteration
+    // space of the output tile assigned to the work unit.
+    return cute::size(cute::ceil_div(cute::get<2>(problem_shape), cute::get<2>(tile_shape)));
+  }
+
+  CUTLASS_HOST_DEVICE
+  static uint32_t
+  get_work_k_tile_start(WorkTileInfo const&) {
+    // All work units returned by this scheduler start from K tile 0
+    return 0u;
+  }
+
+  CUTLASS_DEVICE
+  static bool
+  need_separate_reduction(Params const& params) {
+    return false;
+  }
+
+  CUTLASS_DEVICE
+  bool
+  is_work_tile_for_reduction(WorkTileInfo const& work_tile_info, Params const& params) {
+    return false;
+  }
+
+  CUTLASS_DEVICE
+  uint32_t
+  epilgoue_subtile_idx(WorkTileInfo const& work_tile_info, Params const& params) const {
+    return 0;
+  }
+
+  template <class FrgTensorC>
+  CUTLASS_DEVICE
+  void
+  separate_reduction(
+    Params const& params,
+    WorkTileInfo const& work_tile_info,
+    FrgTensorC& accumulators,
+    uint32_t num_barriers,
+    uint32_t barrier_idx) {
+  }
+
+  // Shares the accumulator set with peers in the global workspace
+  template <class FrgTensorC>
+  CUTLASS_DEVICE
+  static void
+  share(
+    Params const& params,
+    WorkTileInfo const& work_tile_info,
+    FrgTensorC& accumulators,
+    uint32_t num_barriers,
+    uint32_t barrier_idx) {
+  }
+
+  CUTLASS_DEVICE
+  static bool
+  valid_warpgroup_in_work_tile(WorkTileInfo const& work_tile_info) {
+    return true;
+  }
+
+  CUTLASS_DEVICE
+  static bool
+  requires_separate_reduction(Params const& params) {
+    return false;
+  }
+};
+
+} // namespace cutlass::gemm::kernel::detail

include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp

@@ -69,6 +69,7 @@ public:
 
   using Params = PersistentTileSchedulerSm90StreamKParams;
   using ReductionMode = Params::ReductionMode;
+  using DecompositionMode = Params::DecompositionMode;
 
   struct WorkTileInfo {
     int32_t M_idx = 0;
@@ -83,11 +84,43 @@ public:
     // Number of k tiles remaining for the work unit as a whole
     uint32_t k_tile_remaining = 0;
 
+    // Whether this unit of work is the final split for the given tile
+    bool is_separate_reduction = false;
+
     CUTLASS_HOST_DEVICE
     bool
     is_valid() const {
-      // Use negative indices to denote invalid work
-      return M_idx >= 0;
+      // A work tile that computes no K tiles is invalid unless it is a separate-reduction work tile
+      // (which only performs reduction and epilogue)
+      return k_tile_count > 0 || is_separate_reduction;
+    }
+
+    CUTLASS_HOST_DEVICE
+    bool
+    is_reduction_unit() const {
+      return is_separate_reduction;
+    }
+
+    CUTLASS_HOST_DEVICE
+    int32_t
+    reduction_subtile_idx() const {
+      // For separate reduction units, the K_idx of the work tile is unused.
+      // Therefore, we override it to contain the subtile of that the reduction
+      // unit operates on.
+      return is_reduction_unit() ? K_idx : -1;
+    }
+
+    CUTLASS_HOST_DEVICE
+    void
+    setup_separate_reduction(int32_t epilogue_subtile_idx) {
+      // Set the epilogue subtile in the K_idx, since this is otherwise unused
+      // by separate reduction units.
+      K_idx = epilogue_subtile_idx;
+
+      is_separate_reduction = true;
+      k_tile_count = 0;
+      // Clean up remaining k tiles
+      k_tile_remaining = 0;
     }
 
     CUTLASS_HOST_DEVICE
@@ -113,7 +146,10 @@ public:
     Arguments&
     operator=(Arguments const& args) {
       splits = args.splits;
+      max_swizzle_size = args.max_swizzle_size;
       raster_order = args.raster_order;
+      reduction_mode = args.reduction_mode;
+      decomposition_mode = args.decomposition_mode;
       return *this;
     }
 
@@ -121,7 +157,10 @@ public:
     Arguments&
     operator=(Arguments&& args) noexcept {
       splits = args.splits;
+      max_swizzle_size = args.max_swizzle_size;
       raster_order = args.raster_order;
+      reduction_mode = args.reduction_mode;
+      decomposition_mode = args.decomposition_mode;
       return *this;
     }
 
@@ -129,18 +168,20 @@ public:
     Arguments(int splits_) : splits(splits_) {}
 
     CUTLASS_HOST_DEVICE
-    Arguments(int splits_, int max_swizzle_size_, RasterOrderOptions raster_order_) :
+    Arguments(int splits_, int max_swizzle_size_, RasterOrderOptions raster_order_, DecompositionMode decomposition_mode_) :
       splits(splits_),
       max_swizzle_size(max_swizzle_size_),
-      raster_order(raster_order_) {}
+      raster_order(raster_order_),
+      decomposition_mode(decomposition_mode_) {}
 
     // The splitting factor to be used in a split-K decomposition of the problem.
     // If this is set to a value greater than 1, stream-K decomposition logic
     // is bypassed in favor of a split-K decomposition.
     int splits = 1;
-    const int max_swizzle_size = 1;
+    int max_swizzle_size = 1;
     RasterOrderOptions raster_order = RasterOrderOptions::Heuristic;
     ReductionMode reduction_mode = ReductionMode::Deterministic;
+    DecompositionMode decomposition_mode = DecompositionMode::Heuristic;
   };
 
   // Sink scheduler params as a member
@@ -158,7 +199,8 @@ public:
     ClusterShape cluster_shape,
     KernelHardwareInfo const& hw_info,
     Arguments const& args,
-    void* workspace) {
+    void* workspace,
+    const uint32_t epilogue_subtile = 1) {
 
     static_assert(cute::is_static<TileShape>::value);
     static_assert(cute::is_static<ClusterShape>::value);
@@ -177,11 +219,22 @@ public:
       args.max_swizzle_size,
       args.raster_order,
       args.reduction_mode,
-      workspace
+      args.decomposition_mode,
+      workspace,
+      epilogue_subtile
     );
     return params;
   }
 
+  CUTLASS_HOST_DEVICE
+  static bool
+  can_implement(Arguments const& args) {
+    // Split count > 1 is only valid for heuristic and split-K decomposition modes
+    return (args.splits == 1 ||
+            args.decomposition_mode == DecompositionMode::Heuristic ||
+            args.decomposition_mode == DecompositionMode::SplitK);
+  }
+
   CUTLASS_HOST_DEVICE
   PersistentTileSchedulerSm90StreamK() { };
 
@@ -210,7 +263,7 @@ public:
     // for the fact that we have splits_ peers per output tile, we multiply this
     // value by splits_. For stream-K, this multiplication ends up being a no-op
     // because splits_ is set to 1 for stream-K.
-    if (linear_idx >= params.units_per_problem_ * params.splits_) {
+    if(linear_idx >= (params.units_per_problem_ * params.splits_ + params.separate_reduction_units_)) {
       // Invalid work. Return an empty result.
       return WorkTileInfo::invalid_work_tile();
     }
@@ -231,8 +284,8 @@ public:
       current_work_linear_idx_, work_tile_info, scheduler_params);
   }
 
-  CUTLASS_DEVICE static
-  bool
+  CUTLASS_DEVICE
+  static bool
   continue_current_work_for_linear_idx(
     uint64_t linear_idx,
     WorkTileInfo& work_tile_info,
@@ -243,9 +296,8 @@ public:
     if (work_tile_info.k_tile_remaining == 0) {
       return false;
     }
-
     assign_work(params, linear_idx, work_tile_info);
-    return true;
+    return work_tile_info.is_valid();
   }
 
   CUTLASS_DEVICE
@@ -281,7 +333,7 @@ public:
       problem_blocks,
       to_gemm_coord(cluster_shape),
       hw_info,
-      arguments.max_swizzle_size, 
+      arguments.max_swizzle_size,
       arguments.raster_order
     );
   }
@@ -290,8 +342,22 @@ public:
   CUTLASS_HOST_DEVICE
   static bool
   requires_fixup(Params const& params, WorkTileInfo const& work_tile_info) {
-    // Fixup is not needed for data-parallel tiles
-    return work_tile_info.k_tile_count != params.divmod_tiles_per_output_tile_.divisor;
+    // Fixup is not needed for invalid or data-parallel tiles
+    return work_tile_info.is_valid() && work_tile_info.k_tile_count != params.divmod_tiles_per_output_tile_.divisor;
+  }
+
+  CUTLASS_HOST_DEVICE
+  static bool
+  requires_separate_reduction(Params const& params) {
+    return params.requires_separate_reduction();
+  }
+
+  // When the work tile is not special for reduction, it's valid. Otherwise need to skip
+  // global loading that producer warpgroup do, also math computation that consumer warpgroup do.
+  CUTLASS_DEVICE
+  static bool
+  valid_warpgroup_in_work_tile(WorkTileInfo const& work_tile_info) {
+    return !work_tile_info.is_reduction_unit();
   }
 
   // Performs the reduction across splits for a given output tile.
@@ -304,7 +370,7 @@ public:
     FrgTensorC& accumulators,
     uint32_t num_barriers,
     uint32_t barrier_idx) {
-    static constexpr uint32_t Offset = 2;
+    static constexpr uint32_t Offset = static_cast<int>(cutlass::arch::ReservedNamedBarriers::StreamkBarrier0);
     static constexpr uint32_t MaxNumNamedBarriers = 2;
     using BarrierManager = NamedBarrierManager<NumThreadsPerWarpGroup, Offset, MaxNumNamedBarriers>;
     return fixup_helper<FrgTensorC, BarrierManager>(
@@ -327,16 +393,27 @@ public:
     if (!requires_fixup(params, work_tile_info)) {
       return;
     }
-
     auto tile_idx = output_tile_index(params, work_tile_info);
 
     // Index of the lock on which to wait
     auto lock_idx = (tile_idx * num_barriers) + barrier_idx;
 
+    auto reduction_tile_idx = tile_idx;
+    auto [first_peer_id, my_peer_id, last_peer_id] = tile_peer_range(params, tile_idx, static_cast<uint32_t>(work_tile_info.K_idx));
+    auto reduction_peer_offset = 0;
+    if (params.requires_separate_reduction()) {
+      // If separate reduction is to be performed, each stream-K unit writes its partials
+      // to a separate portion of the workspace. There are as many of these portions as there
+      // are peers for a given output tile, so we multiply the tile index by the maximum peer count.
+      reduction_tile_idx *= Params::max_peers_per_tile(params.sk_units_, params.sk_tiles_);
+      reduction_peer_offset = my_peer_id * cute::size<0>(TileShape{}) * cute::size<1>(TileShape{});
+    }
+
     // Reductions use BlockStripedReduce with a width of BarrierManager::ThreadCount under the hood.
     // Thus, the start of the reduction space is the same across all threads in a warp group.
     int reduction_offset =
-      (cute::size<0>(TileShape{}) * cute::size<1>(TileShape{}) * tile_idx) +
+      (cute::size<0>(TileShape{}) * cute::size<1>(TileShape{}) * reduction_tile_idx) +
+      reduction_peer_offset +
       (size(accumulators) * barrier_idx * BarrierManager::ThreadCount);
 
     ElementAccumulator* group_reduction_workspace = reinterpret_cast<ElementAccumulator*>(params.reduction_workspace_) + reduction_offset;
@@ -344,56 +421,86 @@ public:
     using AccumulatorArrayT = Array<typename FrgTensorC::value_type, size(FrgTensorC{})>;
     using BlockStripedReduceT = BlockStripedReduce<BarrierManager::ThreadCount, AccumulatorArrayT>;
 
+    AccumulatorArrayT* reduction_workspace_array = reinterpret_cast<AccumulatorArrayT*>(group_reduction_workspace);
+    AccumulatorArrayT* accumulator_array = reinterpret_cast<AccumulatorArrayT*>(&accumulators);
+
+    int barrier_group_thread_idx = threadIdx.x % BarrierManager::ThreadCount;
+
     // The number of tiles for which reduction is required is either:
     //   (a) the total number of output tiles (in the case of split-K)
-    //   (b) the number of stream-K tiles
+    //   (b) the number of stream-K tiles (potentially multiplied by peer count if using separate reduction)
     // To calculate the total number of output tiles in the split-K case, we
     // note that, in the split-K case, the units_per_problem_ member of Params will be
     // the total number of output tiles.
-    auto reduction_tiles = params.splits_ > 1 ? params.units_per_problem_ : params.sk_tiles_;
+    uint32_t reduction_tiles = 0;
+    if (params.splits_ > 1) {
+      reduction_tiles = params.units_per_problem_;
+    }
+    else if (params.requires_separate_reduction()) {
+      reduction_tiles = params.sk_tiles_ * Params::max_peers_per_tile(params.sk_units_, params.sk_tiles_);
+    }
+    else {
+      reduction_tiles = params.sk_tiles_;
+    }
+
     auto reduction_workspace_size = Params::get_reduction_workspace_size(
       reduction_tiles, to_gemm_coord(TileShape{}), sizeof_bits<ElementAccumulator>::value);
     BarrierType* lock_workspace = reinterpret_cast<BarrierType*>(
       reinterpret_cast<uint8_t*>(params.reduction_workspace_) + reduction_workspace_size);
 
-    AccumulatorArrayT* reduction_workspace_array = reinterpret_cast<AccumulatorArrayT*>(group_reduction_workspace);
-    AccumulatorArrayT* accumulator_array = reinterpret_cast<AccumulatorArrayT*>(&accumulators);
-    int barrier_group_thread_idx = threadIdx.x % BarrierManager::ThreadCount;
+    if (work_tile_info.is_reduction_unit()) {
+      plus<AccumulatorArrayT> add_fragments;
+      auto peer_offset = size(accumulators) * num_barriers * BarrierManager::ThreadCount;
 
-    if (!work_tile_info.is_final_split(params.divmod_tiles_per_output_tile_.divisor)) {
-      if (work_tile_info.K_idx == 0) {
-        // First peer initializes the workspace partials
+      // Wait until the peers collaborating on this output tile have all written
+      // their accumulators to workspace.
+      uint32_t num_peers = last_peer_id - first_peer_id + 1;
+      BarrierManager::wait_eq(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, num_peers);
+
+      // Load the first peer's data
+      BlockStripedReduceT::load(*accumulator_array, reduction_workspace_array, barrier_group_thread_idx);
+
+      for (int i = 1; i < num_peers; ++i) {
+        // Load peer fragment
+        AccumulatorArrayT addend_fragment;
+        auto peer_reduction_workspace = reinterpret_cast<AccumulatorArrayT*>(group_reduction_workspace + (i * peer_offset));
+
+        BlockStripedReduceT::load(addend_fragment, peer_reduction_workspace, barrier_group_thread_idx);
+
+        // Add peer fragment
+        *accumulator_array = add_fragments(*accumulator_array, addend_fragment);
+      }
+    }
+    else if (!compute_epilogue(work_tile_info, params)) {
+      if (params.requires_separate_reduction() || work_tile_info.K_idx == 0) {
+        // The first peer initializes the workspace partials in the non-separate-reduction case,
+        // and all peers write to their own location in workspace when using separate reduction
         BlockStripedReduceT::store(reduction_workspace_array, *accumulator_array, barrier_group_thread_idx);
       }
       else {
-        if (params.reduction_mode_ == ReductionMode::Deterministic) {
-          // Wait until the preceding split added its accumulators
-          BarrierManager::wait_eq(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, work_tile_info.K_idx);
-        }
-        else {
-          // Wait until the first split has stored its accumulators. Note that the first split will have
-          // accumulated a value into the lock potentially greater than one (since the locked value is
-          // incremented by work_tile_info.k_tile_count below for both the deterministic and non-deterministic)
-          // cases. For non-deterministic reductions, all that non-first or last splits care about is whether
-          // the first split has been written, so we only wait while the locked value is less than 1. This
-          // avoids having to add logic to determine the work_tile_info.k_tile_count for the first split.
-          BarrierManager::wait_lt(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, 1);
-        }
+        // Wait until the preceding split added its accumulators
+        BarrierManager::wait_eq(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, work_tile_info.K_idx);
 
         // Perform reduction in workspace
         BlockStripedReduceT::reduce(reduction_workspace_array, *accumulator_array, barrier_group_thread_idx);
       }
 
+      // If separate reduction is being performed, each participating stream-K unit increments the barrier
+      // by only 1. Otherwise, increment by the K tile count that this unit has processed.
+      int32_t increment = params.requires_separate_reduction() ? 1 : work_tile_info.k_tile_count;
+
       // Signal our arrival
-      BarrierManager::arrive_inc(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, work_tile_info.k_tile_count);
+      BarrierManager::arrive_inc(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, increment);
     }
     else {
-      // Wait until the preceding split added its accumulators.
-      // For both the deterministic and non-deterministic case, each preceding split will have incremented
-      // the locked value by work_tile_info.k_tile_count. Thus, the final split konws that it can begin
-      // loading the partially-reduced value when the locked value reaches its starting K tile index (i.e.,
-      // work_tile_info.K_idx).
-      BarrierManager::wait_eq(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, work_tile_info.K_idx);
+      if (params.reduction_mode_ == ReductionMode::Deterministic) {
+        // Wait until the preceding split added its accumulators
+        BarrierManager::wait_eq(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, work_tile_info.K_idx);
+      }
+      else {
+        // Wait unitl the first split has stored its accumulators
+        BarrierManager::wait_lt(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, 1);
+      }
 
       // The block computing the final split for the tile adds previously-reduced partials
       // to its accumulators and computes the epilogue.
@@ -406,7 +513,13 @@ public:
   CUTLASS_HOST_DEVICE
   static bool
   compute_epilogue(WorkTileInfo const& work_tile_info, Params const& params) {
-    return work_tile_info.is_final_split(params.divmod_tiles_per_output_tile_.divisor);
+    // `is_final_split` will be set to `true` for the following scenarios, all of which must compute the epilogue:
+    //  1. The tile is computed in data-parallel mode
+    //  2. The tile is computed in split-/stream-K mode and this work unit represents the final split of the tile
+    //  3. The tile is computed in split-/stream-K mode and separate reduction is used, and this is a separate reduction unit
+    return work_tile_info.is_valid() &&
+            (work_tile_info.is_final_split(params.divmod_tiles_per_output_tile_.divisor) &&
+             !params.requires_separate_reduction()) || work_tile_info.is_separate_reduction;
   }
 
   // Returns the linearized index of the output tile corresponding to the tile with offset [L, M, K]
@@ -432,7 +545,8 @@ public:
     Arguments const& args,
     ProblemShape problem_shape,
     KernelHardwareInfo const& hw_info,
-    uint32_t mma_warp_groups) {
+    uint32_t mma_warp_groups,
+    const uint32_t epilogue_subtile = 1) {
 
     auto problem_shape_mnkl = cute::append<4>(problem_shape, 1);
 
@@ -451,9 +565,11 @@ public:
       args.splits,
       args.max_swizzle_size,
       args.raster_order,
+      args.decomposition_mode,
       mma_warp_groups,
       sizeof_bits<BarrierType>::value,
-      sizeof_bits<ElementAccumulator>::value
+      sizeof_bits<ElementAccumulator>::value,
+      epilogue_subtile
     );
   }
 
@@ -464,8 +580,9 @@ public:
     void* workspace,
     cudaStream_t stream,
     ProblemShape const& problem_shape,
-     KernelHardwareInfo const& hw_info,
-    uint32_t mma_warp_groups) {
+    KernelHardwareInfo const& hw_info,
+    uint32_t mma_warp_groups,
+    const uint32_t epilogue_subtile = 1) {
 
     auto problem_shape_mnkl = cute::append<4>(problem_shape, 1);
 
@@ -486,9 +603,11 @@ public:
       args.splits,
       args.max_swizzle_size,
       args.raster_order,
+      args.decomposition_mode,
       mma_warp_groups,
       sizeof_bits<BarrierType>::value,
-      sizeof_bits<ElementAccumulator>::value
+      sizeof_bits<ElementAccumulator>::value,
+      epilogue_subtile
     );
   }
 
@@ -505,6 +624,7 @@ public:
     return work_tile_info.K_idx;
   }
 
+private:
   // Sets the current stream-K work to compute within work_tile_info. If new_unit is true, work_tile_info
   // is populated as a new unit of work. Otherwise, state existing in work_tile_info (e.g., remaining
   // iterations) is used to find the next tile in the current work unit.
@@ -515,10 +635,22 @@ public:
     uint64_t linear_idx,
     WorkTileInfo& work_tile_info) {
 
-    uint64_t true_tile_id = linear_idx;
-    if (linear_idx >= params.sk_units_ && params.splits_ == 1) {
+    uint64_t output_tile_id = linear_idx;
+    if (linear_idx >= params.units_per_problem_ * params.splits_) {
+      // Separate-reduction work
+      auto cluster_size = params.get_cluster_size();
+      // Divide up the linearized separate reduction units into clusters
+      auto cluster_linear_reduction_unit_idx = params.div_cluster_size((linear_idx - params.units_per_problem_));
+      uint64_t cluster_tile_idx, epi_subtile_idx;
+      params.divmod_epilogue_subtile_(cluster_tile_idx, epi_subtile_idx, cluster_linear_reduction_unit_idx);
+      // Bring the linearized tile ID back into the space of tiles, rather than clusters
+      output_tile_id = cluster_tile_idx * cluster_size;
+
+      work_tile_info.setup_separate_reduction(epi_subtile_idx);
+    }
+    else if (linear_idx >= params.sk_units_ && params.splits_ == 1) {
       // Data-parallel work
-      true_tile_id = linear_idx - params.sk_units_ + params.sk_tiles_;
+      output_tile_id = linear_idx - params.sk_units_ + params.sk_tiles_;
       work_tile_info.K_idx = 0;
       work_tile_info.k_tile_count = params.divmod_tiles_per_output_tile_.divisor;
       work_tile_info.k_tile_remaining = params.divmod_tiles_per_output_tile_.divisor;
@@ -540,48 +672,114 @@ public:
       // To do so, we divide up the linearized stream-K units into clusters and share the same K
       // offsets for work within clusters.
 
-      // Equivalent to linear_idx / cluster_size
-      auto cluster_linear_work_idx = params.divmod_cluster_shape_minor_.divide(
-        params.divmod_cluster_shape_major_.divide(linear_idx)
-      );
+      auto cluster_linear_work_idx = params.div_cluster_size(linear_idx);
+
+      uint64_t group_idx;
+      params.divmod_sk_groups_(cluster_linear_work_idx, group_idx, cluster_linear_work_idx);
+
+      // Determine whether we are in a "big group" that will process an additional
+      // stream-K cluster tile.
+      auto sk_cluster_tiles = params.div_cluster_size(params.sk_tiles_);
+      auto sk_cluster_tiles_in_group = params.divmod_sk_groups_.divide(sk_cluster_tiles);
+      if (group_idx < params.big_groups_) {
+        ++sk_cluster_tiles_in_group;
+      }
+
+      // Determine whether we are in a "big unit" within the group, that will process
+      // an additional K chunk in the group.
+      auto sk_tiles_in_group = sk_cluster_tiles_in_group * params.get_cluster_size();
+      auto k_tiles_in_group = sk_tiles_in_group * params.divmod_tiles_per_output_tile_.divisor;
+      auto k_tiles_per_unit_in_group = params.divmod_sk_units_per_group_.divide(k_tiles_in_group);
+      auto big_units_in_group = params.div_cluster_size(
+        k_tiles_in_group - (k_tiles_per_unit_in_group * params.divmod_sk_units_per_group_.divisor));
 
       uint64_t split;
       params.divmod_clusters_mnl_(split, cluster_linear_work_idx, cluster_linear_work_idx);
-      auto big_unit_cmp = params.splits_ > 1 ? split : cluster_linear_work_idx;
-      auto linear_idx_mult = params.splits_ > 1 ? params.divmod_tiles_per_output_tile_.divisor : params.k_tiles_per_sk_unit_;
+
+      bool is_split_k = params.splits_ > 1;
+      auto big_unit_cmp_lhs = is_split_k ? split : cluster_linear_work_idx;
+      auto big_unit_cmp_rhs = is_split_k ? params.big_units_ : big_units_in_group;
+      auto linear_idx_mult = is_split_k ? params.divmod_tiles_per_output_tile_.divisor : k_tiles_per_unit_in_group;
+      auto k_tiles_per_split = is_split_k ? params.k_tiles_per_sk_unit_ : k_tiles_per_unit_in_group;
 
       // Determine the starting k iteration computed by this stream-K work unit
-      uint32_t unit_iter_start = (linear_idx_mult * cluster_linear_work_idx) + (params.k_tiles_per_sk_unit_ * split);
+      uint32_t unit_iter_start = (linear_idx_mult * cluster_linear_work_idx) +
+                                 (k_tiles_per_split * split);
 
       // Adjust the starting position and number of k iterations for "big units," which
-      // compute one extra iteration. These are the first big_units_ units in the
-      // linearized ID space.
-      bool is_big_unit = big_unit_cmp < params.big_units_;
-      if (is_big_unit) {
+      // compute one extra iteration. If there are any big units, they will be the first
+      // in the linearized ID space.
+      auto k_tiles_in_my_split = k_tiles_per_split;
+      if (big_unit_cmp_lhs < big_unit_cmp_rhs) {
         // Since the "big units" are the first units in the linearized ID space, each
         // of the units preceding this big unit computed one extra iteration. Thus,
         // we must offset our start iteration by the number of units that precede
         // the current unit in the linearized ID space.
-        unit_iter_start += big_unit_cmp;
+        unit_iter_start += big_unit_cmp_lhs;
+        ++k_tiles_in_my_split;
       }
       else {
         // Increment by one for each of the big clusters (since all big units precede this unit)
-        unit_iter_start += params.big_units_;
+        unit_iter_start += big_unit_cmp_rhs;
+      }
+
+      if (!is_split_k) {
+        // Adjust the unit starting position and number of tiles to avoid
+        // computing splits of size less than min_iters_per_sk_unit_
+        int unused, start_tile_k_tile;
+        params.divmod_tiles_per_output_tile_(unused, start_tile_k_tile, unit_iter_start);
+        if (start_tile_k_tile < Params::min_iters_per_sk_unit_) {
+          // Starting K tile is in range [0, Params::min_iters_per_sk_unit_), which means that another
+          // stream-K unit will be computing a split with fewer than Params::min_iters_per_sk_unit_ K tiles.
+          // Adjust our work to take over these K tiles.
+          unit_iter_start -= start_tile_k_tile;
+          k_tiles_in_my_split += start_tile_k_tile;
+        }
+        else if (start_tile_k_tile > (params.divmod_tiles_per_output_tile_.divisor - Params::min_iters_per_sk_unit_)) {
+          // Starting K tile is within the final Params::min_iters_per_sk_unit_ K tiles of some output tile,
+          // which means that this unit will compute a split with fewer than Params::min_iters_per_sk_unit_ K tiles.
+          // Adjust our work to shed these K tiles to a neighboring stream-K unit that will compute more consecutive K tiles.
+          auto adjustment_tiles = (params.divmod_tiles_per_output_tile_.divisor - start_tile_k_tile);
+          unit_iter_start += adjustment_tiles;
+          k_tiles_in_my_split -= adjustment_tiles;
+        }
       }
 
       if (work_tile_info.k_tile_count == 0) {
         // This is a new unit
-        work_tile_info.k_tile_remaining = params.k_tiles_per_sk_unit_;
 
-        // Only adjust iteration count for big unit if we are initializing this
-        // work unit. For existing work units, the extra iteration for big units
-        // has already been accounted for in k_tiles_reamaining
-        if (is_big_unit) {
-          ++work_tile_info.k_tile_remaining;
+        if (!is_split_k) {
+          //
+          // Adjust the unit ending position and number of tiles to avoid
+          // computing splits of size less than min_iters_per_sk_unit_
+          //
+
+          // Begin by assuming that no adjustment is needed
+          auto initial_unit_iter_end = unit_iter_start + k_tiles_in_my_split;
+
+          int unused, end_tile_k_tile;
+          params.divmod_tiles_per_output_tile_(unused, end_tile_k_tile, initial_unit_iter_end);
+
+          if (end_tile_k_tile < Params::min_iters_per_sk_unit_) {
+            // Ending K tile is within the first Params::min_iters_per_sk_unit_ K tiles of some output tile,
+            // which means that this unit will compute a split with fewer than Params::min_iters_per_sk_unit_ K tiles.
+            // Adjust our work to shed these K tiles to a neighboring stream-K unit that will compute more consecutive K tiles.
+            k_tiles_in_my_split -= end_tile_k_tile;
+          }
+          else if (end_tile_k_tile > (params.divmod_tiles_per_output_tile_.divisor - Params::min_iters_per_sk_unit_)) {
+            // Ending K tile is within the final Params::min_iters_per_sk_unit_ K tiles of some output tile,
+            // which means that some other unit will compute a split with fewer than Params::min_iters_per_sk_unit_ K tiles.
+            // Adjust our work to take on these K tiles.
+            k_tiles_in_my_split += (params.divmod_tiles_per_output_tile_.divisor - end_tile_k_tile);
+          }
         }
+
+        work_tile_info.k_tile_remaining = k_tiles_in_my_split;
       }
 
-      // Find the output tile corresponding to the final k iteration covered by this
+      uint32_t unit_iter_end = unit_iter_start + work_tile_info.k_tile_remaining - 1;
+
+      // Find the output tile corresponding to the final k tile covered by this
       // work unit. Stream-K work units will work backwards in terms of the tiles they
       // are responsible computing. This is beneficial because the final (partial)
       // tile computed by a stream-K block is typically the beginning of the output
@@ -590,43 +788,45 @@ public:
       // other work units computing portions of that output tile, it is preferable
       // for them to be computed later, so as to reduce the likelihood of blocking
       // on other work.
-      uint32_t unit_iter_end = unit_iter_start + work_tile_info.k_tile_remaining - 1;
 
-      true_tile_id = params.divmod_tiles_per_output_tile_.divide(unit_iter_end);
-      uint32_t true_tile_iter_start = true_tile_id * params.divmod_tiles_per_output_tile_.divisor;
-      uint32_t true_tile_iter_end = true_tile_iter_start + params.divmod_tiles_per_output_tile_.divisor;
+      auto output_tile_id_in_group = params.divmod_tiles_per_output_tile_.divide(unit_iter_end);
+      uint32_t output_tile_iter_start = output_tile_id_in_group * params.divmod_tiles_per_output_tile_.divisor;
+      uint32_t output_tile_iter_end = output_tile_iter_start + params.divmod_tiles_per_output_tile_.divisor;
+
+      // Convert the output tile from the linearized space within each group to the
+      // overall linearized space.
+      output_tile_id = (output_tile_id_in_group * params.divmod_sk_groups_.divisor) + group_idx;
 
       // Bring the linearized tile ID back into the space of tiles, rather than clusters
-      true_tile_id *= params.divmod_cluster_shape_major_.divisor * params.divmod_cluster_shape_minor_.divisor;
+      output_tile_id *= params.get_cluster_size();
 
       auto [cta_m_in_cluster, cta_n_in_cluster, _] = cute::block_id_in_cluster();
 
       // The final linearized tile ID is in units of the cluster dimension over which we rasterize.
       if (params.raster_order_ == RasterOrder::AlongN) {
-        true_tile_id += cta_n_in_cluster * params.divmod_cluster_shape_minor_.divisor;
+        output_tile_id += cta_n_in_cluster * params.divmod_cluster_shape_minor_.divisor;
       }
       else {
-        true_tile_id += cta_m_in_cluster * params.divmod_cluster_shape_minor_.divisor;
+        output_tile_id += cta_m_in_cluster * params.divmod_cluster_shape_minor_.divisor;
       }
 
       // The unit's starting k iteration in the current tile is either the starting
       // iteration for the tile as a whole, or the starting k iteration for the unit
       // as a whole (if the latter is greater than the former).
-      uint32_t tile_iter_start = max(true_tile_iter_start, unit_iter_start);
+      uint32_t tile_iter_start = max(output_tile_iter_start, unit_iter_start);
 
       // Similarly, the unit's ending k iteration (exclusive) is either the end of
       // the current tile it is assigned, or the ending iteration of the unit as a whole
       // (if the latter is less than the former).
-      uint32_t tile_iter_end = min(true_tile_iter_end, unit_iter_end + 1);
+      uint32_t tile_iter_end = min(output_tile_iter_end, unit_iter_end + 1);
 
       // Set the k offset to be the starting k tile for this output tile
-      work_tile_info.K_idx = static_cast<int32_t>(tile_iter_start - true_tile_iter_start);
-
+      work_tile_info.K_idx = static_cast<int32_t>(tile_iter_start - output_tile_iter_start);
       work_tile_info.k_tile_count = tile_iter_end - tile_iter_start;
     }
 
     uint64_t work_idx_l, remainder;
-    params.divmod_batch_(work_idx_l, remainder, true_tile_id);
+    params.divmod_batch_(work_idx_l, remainder, output_tile_id);
 
     uint64_t cta_per_grid_dim = params.divmod_cluster_shape_minor_.divide(remainder);
 
@@ -642,7 +842,57 @@ public:
     work_tile_info.M_idx = work_idx_m;
     work_tile_info.N_idx = work_idx_n;
     work_tile_info.L_idx = static_cast<int32_t>(work_idx_l);
+  }
 
+  // Returns the starting and ending peer ID of this tile
+  CUTLASS_HOST_DEVICE
+  static auto
+  tile_peer_range(Params const& params, uint32_t tile_idx, uint32_t cur_k_tile) {
+    auto tile_idx_in_cluster_path = params.div_cluster_size(tile_idx);
+    auto start_k_tile = params.divmod_tiles_per_output_tile_.divisor * tile_idx_in_cluster_path;
+    auto end_k_tile = start_k_tile + params.divmod_tiles_per_output_tile_.divisor - 1;
+    auto big_unit_k_tiles = params.big_units_ * (params.k_tiles_per_sk_unit_ + 1);
+
+    auto adjust_unit = [&](uint32_t k_tile, uint32_t unit_idx, uint32_t k_tiles_per_unit) {
+      auto unit_k_start = unit_idx * k_tiles_per_unit;
+      auto unit_k_end = unit_k_start + k_tiles_per_unit;
+      if (k_tile - start_k_tile < Params::min_iters_per_sk_unit_ &&
+          unit_k_end - start_k_tile < Params::min_iters_per_sk_unit_) {
+        // k_tile is within the first min_iters_per_sk_unit_ K tiles of this output tile,
+        // and the stream-K unit computes fewer than min_iters_per_sk_unit_ K tiles for this
+        // output tile. This work will thus be subsumed by the next stream-K unit.
+        ++unit_idx;
+      }
+
+      if (end_k_tile + 1 - k_tile < Params::min_iters_per_sk_unit_ &&
+          end_k_tile + 1 - unit_k_start < Params::min_iters_per_sk_unit_) {
+        // k_tile is within the last min_iters_per_sk_unit_ K tiles of this output tile,
+        // and the stream-K unit computes fewer than min_iters_per_sk_unit_ K tiles for this
+        // output tile. This work will thus be subsumed by the previous stream-K unit.
+        --unit_idx;
+      }
+
+      return unit_idx;
+    };
+
+    // Lambda to find the ID of the stream-K unit that computes this K tile
+    auto find_unit = [&](uint32_t k_tile) {
+      if (k_tile < big_unit_k_tiles) {
+        // The tile is within the "big unit range"
+        auto k_tiles_per_unit = params.k_tiles_per_sk_unit_ + 1;
+        auto unit_idx = k_tile / k_tiles_per_unit;
+        return static_cast<uint64_t>(adjust_unit(k_tile, unit_idx, k_tiles_per_unit));
+      }
+      else {
+        // The tile is after the "big unit range." Account for this by finding the "normal unit"
+        // that it belongs to, and then offsetting by the number of big units
+        auto k_tiles_per_unit = params.k_tiles_per_sk_unit_;
+        auto unit_idx = ((k_tile - big_unit_k_tiles) / params.k_tiles_per_sk_unit_) + (params.big_units_);
+        return static_cast<uint64_t>(adjust_unit(k_tile, unit_idx, k_tiles_per_unit));
+      }
+    };
+
+    return cute::make_tuple(find_unit(start_k_tile), find_unit(cur_k_tile), find_unit(end_k_tile));
   }
 };
 

include/cutlass/gemm/kernel/tile_scheduler.hpp

@@ -38,6 +38,7 @@
 #include "cutlass/detail/dependent_false.hpp"
 #include "cutlass/gemm/kernel/sm90_tile_scheduler.hpp"
 #include "cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp"
+#include "cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp"
 ////////////////////////////////////////////////////////////////////////////////
 
 namespace cutlass::gemm {
@@ -52,6 +53,8 @@ struct PersistentScheduler { };
 
 struct StreamKScheduler { };
 
+struct GroupScheduler { }; // Only used for Grouped GEMMs
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace cutlass::gemm
@@ -69,6 +72,7 @@ template <
   class ArchTag,
   class TileShape,
   class ClusterShape
+  , class ProblemShapeType = void
 >
 struct TileSchedulerSelector {
   static_assert(cutlass::detail::dependent_false<ArchTag>,
@@ -122,6 +126,21 @@ struct TileSchedulerSelector<
   using Scheduler = PersistentTileSchedulerSm90StreamK<TileShape, ClusterShape>;
 };
 
+template <
+  class TileShape,
+  class ClusterShape
+  , class GroupProblemShape
+>
+struct TileSchedulerSelector<
+  GroupScheduler,
+  arch::Sm90,
+  TileShape,
+  ClusterShape
+  , GroupProblemShape
+  > {
+  using Scheduler = PersistentTileSchedulerSm90Group<GroupProblemShape>;
+};
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace cutlass::gemm::kernel::detail

include/cutlass/gemm/kernel/tile_scheduler_params.h

@@ -218,7 +218,7 @@ struct PersistentTileSchedulerSm90Params {
 
     auto possibly_truncate = [&](int x, int y) {
       if (truncate_by_problem_size) {
-        return cutlass::platform::min(x, y);
+        return platform::min(x, y);
       }
       else {
         return x;
@@ -272,7 +272,7 @@ struct PersistentTileSchedulerSm90Params {
   CUTLASS_HOST_DEVICE
   static int32_t
   get_log_swizzle_size(int problem_ctas_m, int problem_ctas_n, int max_swizzle_size) {
-    int min_cta_dim = cutlass::platform::min(problem_ctas_m, problem_ctas_n);
+    int min_cta_dim = platform::min(problem_ctas_m, problem_ctas_n);
     if (max_swizzle_size >= 8 && min_cta_dim >= 6) {
       return 3;
     }
@@ -370,6 +370,18 @@ struct PersistentTileSchedulerSm90StreamKParams {
     Nondeterministic
   };
 
+  // Strategies for decomposing the problem
+  enum class DecompositionMode {
+    // Use a heuristic to determine whether data-parallel, split-K, or stream-K decomposition should be performed
+    Heuristic,
+    // Force a data-parallel decomposition
+    DataParallel,
+    // Force a split-K decomposition. This should be paired with setting the `splits` parameter
+    SplitK,
+    // Force a stream-K decomposition
+    StreamK
+  };
+
   using UnderlyingParams = PersistentTileSchedulerSm90Params;
   using RasterOrder = UnderlyingParams::RasterOrder;
   using RasterOrderOptions = UnderlyingParams::RasterOrderOptions;
@@ -387,6 +399,17 @@ struct PersistentTileSchedulerSm90StreamKParams {
   // and may be overridden in other decompositions.
   FastDivmodU64 divmod_clusters_mnl_{};
 
+  // We divide up the number of stream-K tiles amongst G groups of stream-K units.
+  // The stream-K units within a group collaborate to comptue over the `sk_tiles / G`
+  // tiles assigned to that group. Non-unit group sizes can help to preserve L2 locality of
+  // partial chunks computed by stream-K units -- units 0 in each group will compute identical K extents
+  // of tiles that would be assigned in the same wave according to the rasterization order of the
+  // data-parallel formulation of the problem.
+  FastDivmodU64 divmod_sk_groups_{};
+
+  // Number of stream-K units in each group
+  FastDivmodU64 divmod_sk_units_per_group_{};
+
   uint64_t units_per_problem_ = 0;
   FastDivmod divmod_tiles_per_output_tile_{};
   int32_t log_swizzle_size_ = 0;
@@ -403,6 +426,9 @@ struct PersistentTileSchedulerSm90StreamKParams {
   // at the granularity of a cluster, we store only the number of big clusters.
   uint32_t big_units_ = 0;
 
+  // The number of groups of stream-K units that will process an extra stream-K tile cluster.
+  uint32_t big_groups_ = 0;
+
   // Workspace for holding partial accumulators to be reduced across stream-K/split-K units
   void* reduction_workspace_ = nullptr;
 
@@ -419,8 +445,53 @@ struct PersistentTileSchedulerSm90StreamKParams {
   // Strategy to use when reducing between collaborating CTAs
   ReductionMode reduction_mode_ = ReductionMode::Deterministic;
 
-  // Minimum number of tiled k that can be assigned to a stream-K unit
-  static constexpr uint32_t min_iters_per_sk_unit_ = 4u;
+  // The number of sub blocks in the kernel epilogue
+  FastDivmodU64 divmod_epilogue_subtile_{};
+
+  // The number of blocks that launched for doing separate reduction
+  uint32_t separate_reduction_units_ = 0;
+
+  // Minimum number of k tiles that can be assigned to a stream-K unit
+  static constexpr uint32_t min_iters_per_sk_unit_ = 8u;
+
+  // Maximum number of groups of stream-K units
+  static constexpr uint32_t max_sk_groups_ = 8u;
+
+  // Divides dividend by the cluster size
+  CUTLASS_HOST_DEVICE
+  uint64_t
+  div_cluster_size(uint64_t dividend) const {
+    // Use each underlying fast divmod rather than performing integer division
+    // by the multiplication of major.divisor * minor.divisor
+    return divmod_cluster_shape_minor_.divide(
+      divmod_cluster_shape_major_.divide(dividend)
+    );
+  }
+
+  CUTLASS_HOST_DEVICE
+  uint64_t
+  get_cluster_size() const {
+    return divmod_cluster_shape_minor_.divisor * divmod_cluster_shape_major_.divisor;
+  }
+
+  // Returns whether the kernel uses separate reduction
+  CUTLASS_HOST_DEVICE
+  bool
+  requires_separate_reduction() const {
+    return separate_reduction_units_ > 0;
+  }
+
+  // Returns the maximum number of peers that can collaborate on a given output tile
+  CUTLASS_HOST_DEVICE
+  static uint32_t
+  max_peers_per_tile(uint64_t sk_units, uint64_t sk_tiles) {
+    // When we can divide up our SK units to SK tiles evenly, the number of peers
+    // per SK tile is exactly (sk_units_ / sk_tiles_). In cases where this division
+    // is not exact, some tiles will need to be covered by additional SK units. Because
+    // the extra work can occur at both the beginning and the end of the SK tile, at
+    // most 2 extra peers will be needed.
+    return static_cast<uint32_t>(sk_units / sk_tiles + 2);
+  }
 
   // Initializes members. This variant of the method should only be used when
   // problem_shape and tile_shape contain modes of only rank 1.
@@ -434,7 +505,9 @@ struct PersistentTileSchedulerSm90StreamKParams {
     int max_swizzle,
     RasterOrderOptions raster_order_option,
     ReductionMode reduction_mode,
-    void* workspace
+    DecompositionMode decomposition_mode,
+    void* workspace,
+    const uint32_t epilogue_subtile = 1
   ) {
     dim3 problem_blocks = UnderlyingParams::get_tiled_cta_shape_mnl(
       problem_shape, tile_shape, cluster_shape);
@@ -451,7 +524,9 @@ struct PersistentTileSchedulerSm90StreamKParams {
       max_swizzle,
       raster_order_option,
       reduction_mode,
-      workspace
+      decomposition_mode,
+      workspace,
+      epilogue_subtile
     );
   }
 
@@ -468,7 +543,9 @@ struct PersistentTileSchedulerSm90StreamKParams {
     int max_swizzle,
     RasterOrderOptions raster_order_option,
     ReductionMode reduction_mode,
-    void* workspace
+    DecompositionMode decomposition_mode,
+    void* workspace,
+    const uint32_t epilogue_subtile = 1
   ) {
     UnderlyingParams underlying_params;
     underlying_params.initialize(
@@ -488,7 +565,8 @@ struct PersistentTileSchedulerSm90StreamKParams {
     // Reduction workspace is at the beginning of the workspace. Lock workspace follows.
     void* reduction_workspace = workspace;
 
-    if (splits > 1) {
+    if (decomposition_mode == DecompositionMode::SplitK ||
+        (decomposition_mode == DecompositionMode::Heuristic && splits > 1)) {
       // Short circuit to basic split-K decomposition
 
       // Don't split by more than the available number of SMs
@@ -531,24 +609,9 @@ struct PersistentTileSchedulerSm90StreamKParams {
     uint64_t ctas_per_wave = grid.x * grid.y;
 
     // The number of output tiles to be computed in stream-K and data-parallel fashion, respectively.
-    uint32_t sk_tiles = get_num_sk_tiles(output_tiles, ctas_per_wave, k_tiles_per_output_tile);
+    uint32_t sk_tiles = get_num_sk_tiles(output_tiles, ctas_per_wave, k_tiles_per_output_tile, decomposition_mode);
     uint64_t dp_tiles = output_tiles - sk_tiles;
 
-    if (sk_tiles == 0) {
-      // Short circuit to basic data-parallel decomposition
-      set_params_basic(
-        underlying_params,
-        problem_blocks_m,
-        problem_blocks_n,
-        problem_blocks_l,
-        /* splits = */ 1,
-        k_tiles_per_output_tile,
-        reduction_workspace,
-        reduction_mode
-      );
-      return;
-    }
-
     // Calculate the number of work units covering the data-parallel and stream-K tiles.
     // A "work unit" is a single index in the linearized ID space used by the scheduler.
     // We distinguish it from a "block," which is typically tied to a hardware unit
@@ -576,12 +639,127 @@ struct PersistentTileSchedulerSm90StreamKParams {
     uint64_t min_sized_sk_units = (k_tiles_sk_total / min_iters_per_sk_unit_);
     min_sized_sk_units = (min_sized_sk_units / cluster_size) * cluster_size;
 
-    uint64_t sk_units = cutlass::platform::min(ctas_per_wave, min_sized_sk_units);
+    uint64_t sk_units = platform::min(ctas_per_wave, min_sized_sk_units);
 
-    // If the number of stream-K units is a multiple of the number of stream-K tiles, then
-    // the problem can leverage a basic split-K decomposition for the stream-K tiles.
-    if (sk_tiles < sk_units && sk_units % sk_tiles == 0) {
-      // Short circuit to basic split-K decomposition
+    if (decomposition_mode == DecompositionMode::DataParallel ||
+        (decomposition_mode == DecompositionMode::Heuristic && sk_tiles == 0) ||
+        sk_units == 0) {
+      // Short circuit to basic data-parallel decomposition
+      set_params_basic(
+        underlying_params,
+        problem_blocks_m,
+        problem_blocks_n,
+        problem_blocks_l,
+        /* splits = */ 1,
+        k_tiles_per_output_tile,
+        reduction_workspace,
+        reduction_mode
+      );
+      return;
+    }
+
+    bool do_separate_reduction = should_perform_separate_reduction(
+      epilogue_subtile, sk_units, sk_tiles, dp_tiles, ctas_per_wave);
+
+    // Determine the number of stream-K groups that will be used. We currently use
+    // max_sk_groups_ unless this extends beyond the extent of the dimension over
+    // which the problem is rasterized. For example, if the tiled problem shape
+    // (in CTA_M x CTA_N representation) when using 1x1 clusters is 4x16,
+    // and we rasterize along the M dimension, we choose 4 groups, rather than 8.
+    // If the cluster shape is 2x1, we choose 2 groups (CTA_M / CLUSTER_M).
+    uint32_t max_groups_problem;
+    if (underlying_params.raster_order_ == RasterOrder::AlongM) {
+      max_groups_problem = problem_blocks_m / cluster_shape.m();
+    }
+    else {
+      max_groups_problem = problem_blocks_n / cluster_shape.n();
+    }
+
+    // Select the number of groups that will be use. We start with the maximum
+    // number of potential groups, and iterate down looking for a group size that
+    // evenly divides the stream-K units and tiles, and for which the resulting
+    // number of K tiles per stream-K unit remains above min_iters_per_sk_unit_
+
+    uint32_t groups = platform::min(max_groups_problem, uint32_t(max_sk_groups_));
+
+    // Grouping is disabled when separate reduction is used
+    if (do_separate_reduction) {
+      groups = 1;
+    }
+
+    uint32_t fallback_groups = 0;
+    auto sk_cluster_tiles = sk_tiles / cluster_size;
+    auto sk_cluster_units = sk_units / cluster_size;
+
+    auto sk_splits_too_small = [&](uint32_t g) {
+      // Check whether the number of K tiles computed per stream-K unit is less
+      // than min_iters_per_sk_unit_
+      auto total_sk_k_tiles = (sk_tiles / g) * k_tiles_per_output_tile;
+      auto k_tiles_per_sk_unit = total_sk_k_tiles / (sk_units / g);
+      return k_tiles_per_sk_unit < min_iters_per_sk_unit_;
+    };
+
+    auto is_ideal_grouping = [&](uint32_t g) {
+      // An ideal grouping will evenly divide stream-K clusters, evenly divide
+      // stream-K tiles, and not result in stream-K splits that are too small.
+      return (sk_cluster_units % g == 0) && (sk_cluster_tiles % g == 0) && !sk_splits_too_small(g);
+    };
+
+    auto is_valid_grouping = [&](uint32_t g) {
+      // A grouping is valid, but not ideal, if it evenly divides the
+      // stream-K clusters and does not result in stream-K splits that are
+      // too small. Such a setting can be used as a fallback option in the
+      // case that an ideal grouping is not achievable
+      return sk_cluster_units % g == 0 && !sk_splits_too_small(g);
+    };
+
+    while (groups > 1 && !is_ideal_grouping(groups)) {
+      if (fallback_groups == 0 && is_valid_grouping(groups)) {
+        // Set fallback groups once in preference for a larger number of groups.
+        fallback_groups = groups;
+      }
+      --groups;
+    }
+
+    // If groups == 1, we did not find a group count that satisfies all criteria. If we have
+    // found a fallback group count, use this instead.
+    if (groups == 1 && fallback_groups > 0) {
+      groups = fallback_groups;
+    }
+
+    auto sk_units_per_group = sk_units / groups;
+
+    // sk_tiles is guaranteed to be divisible by cluster_size because it is calculated as:
+    //    sk_tiles = (waves <= 2) ? total_tiles : (sm_count + (total_tiles % sm_count))
+    // Both total_tiles and sm_count are multiples of cluster size due to padding added
+    // prior to kernel launch.
+    uint64_t sk_clustered_tiles = sk_tiles / cluster_size;
+    uint64_t sk_clustered_tiles_per_group = sk_clustered_tiles / groups;
+    uint64_t sk_tiles_per_group = sk_clustered_tiles_per_group * cluster_size;
+
+    // Groups that will process an extra stream-K tile cluster. These differ from "big_units," which
+    // are stream-K units within a group that process an extra K chunk.
+    uint64_t sk_big_groups = sk_clustered_tiles % groups;
+
+    uint64_t k_tiles_per_group = k_tiles_per_output_tile * sk_tiles_per_group;
+
+    // Number of k tiles computed per stream-K unit
+    uint64_t k_tiles_per_sk_unit = k_tiles_per_group / sk_units_per_group;
+
+    uint32_t reduction_units = 0;
+
+    // Use separate reduction when we have less than one wave of output tiles (dp_tiles == 0)
+    // and when each tile will be operated on by at least two stream-K units (sk_units > 2 * sk_tiles)
+    if (do_separate_reduction) {
+      // Each reduction unit will reduce the partials of an epilogue subtile for
+      // a given output tile and compute the epilogue. Thus, there are as many reduction
+      // units as there are epilogue subtiles.
+      reduction_units = sk_tiles * epilogue_subtile;
+    }
+    else if (decomposition_mode == DecompositionMode::Heuristic && sk_tiles < sk_units && sk_units % sk_tiles == 0) {
+      // If the number of stream-K units is a multiple of the number of stream-K tiles, then
+      // the problem can leverage a basic split-K decomposition for the stream-K tiles.
+      // This case happens when separate reduction is disable.
       uint32_t sk_splits = static_cast<uint32_t>(sk_units / sk_tiles);
       set_params_basic(
         underlying_params,
@@ -595,37 +773,13 @@ struct PersistentTileSchedulerSm90StreamKParams {
       );
       return;
     }
-
-    // Number of k iterations computed per stream-K units
-    uint64_t k_tiles_per_sk_unit = k_tiles_sk_total / sk_units;
-
-    // Number of stream-K units that need to compute extra iterations in order to cover
-    // the residual k iterations. This assumes that each such unit computes one additional
-    // iteration.
-    uint64_t sk_big_units = k_tiles_sk_total - (k_tiles_per_sk_unit * sk_units);
-
-    // The division below is guaranteed to be exact because sk_big_units is guaranteed
-    // to be a multiple of cluster_size. This is useful because
-    // it allows us to use a block's linearized cluster ID  to determine whether it is
-    // a big block. The reasoning behind this guarnatee is explained as follows:
-    //     sk_big_units = k_tiles_sk_total - (k_tiles_per_sk_unit * sk_units);
-    //
-    // - k_tiles_sk_total is a multiple of cluster_size because it is the product
-    //   of number of tail tiles and the number of k iterations per tile. Because
-    //   both the number of output tiles and number of available SMs are rounded
-    //   to be multiples of cluster shape, the number of tail tiles
-    //   (output_tiles % avail_sms) is a multpile of cluster_size.
-    //
-    // - sk_units is a multiple of cluster_size because it is either blocks_per_wave
-    //   or 0, and blocks_per_wave is a multiple of the cluster_size due to the grid-planning
-    //   logic rounding to multiples of cluster dimensions
-    uint64_t sk_big_units_per_cluster = sk_big_units / cluster_size;
-
     divmod_cluster_shape_major_ = underlying_params.divmod_cluster_shape_major_;
     divmod_cluster_shape_minor_ = underlying_params.divmod_cluster_shape_minor_;
     divmod_batch_ = underlying_params.divmod_batch_;
     divmod_tiles_per_output_tile_ = FastDivmod(k_tiles_per_output_tile);
     divmod_cluster_blk_major_ = underlying_params.divmod_cluster_blk_major_;
+    divmod_sk_groups_ = FastDivmodU64(static_cast<uint64_t>(groups));
+    divmod_sk_units_per_group_ = FastDivmodU64(static_cast<uint64_t>(sk_units / groups));
 
     // Override divmod_clusters_mnl_ to be the number of cluster-sized stream-K units.
     // This setting ensures that the use of this divmod for stream-K decompositions
@@ -635,12 +789,19 @@ struct PersistentTileSchedulerSm90StreamKParams {
     log_swizzle_size_ = underlying_params.log_swizzle_size_;
     units_per_problem_ = static_cast<uint32_t>(dp_units + sk_units);
     raster_order_ = underlying_params.raster_order_;
-    big_units_ = static_cast<uint32_t>(sk_big_units_per_cluster);
+
+    // Assign big_units_ assuming that group count == 1. This is unused by stream-K
+    // when group count > 1.
+    big_units_ = static_cast<uint32_t>(k_tiles_per_group % k_tiles_per_sk_unit);
+
+    big_groups_ = static_cast<uint32_t>(sk_big_groups);
     reduction_workspace_ = reduction_workspace;
     sk_tiles_ = sk_tiles;
     sk_units_ = static_cast<uint32_t>(sk_units);
     k_tiles_per_sk_unit_ = static_cast<uint32_t>(k_tiles_per_sk_unit);
     reduction_mode_ = reduction_mode;
+    divmod_epilogue_subtile_ = FastDivmodU64(epilogue_subtile);
+    separate_reduction_units_ = reduction_units;
   }
 
   // Given the inputs, computes the physical grid we should launch.
@@ -696,23 +857,28 @@ struct PersistentTileSchedulerSm90StreamKParams {
   // Returns the number of stream-K tiles that will be computed amongst `output_tiles` total
   // output tiles on a device with `ctas_per_wave` CTAs in each wave.
   static uint32_t
-  get_num_sk_tiles(uint64_t output_tiles, uint64_t ctas_per_wave, uint32_t k_tiles_per_output_tile) {
+  get_num_sk_tiles(uint64_t output_tiles, uint64_t ctas_per_wave, uint32_t k_tiles_per_output_tile, DecompositionMode decomposition_mode) {
     uint32_t full_waves = static_cast<uint32_t>(output_tiles / ctas_per_wave);
     uint32_t total_waves = static_cast<uint32_t>((output_tiles + ctas_per_wave - 1) / ctas_per_wave);
 
-    if (full_waves == total_waves || k_tiles_per_output_tile <= min_iters_per_sk_unit_) {
-      // All tiles will be data-parallel tiles if there is either no quantization
-      // or if there is no work to be split.
+    if (decomposition_mode == DecompositionMode::DataParallel ||
+        decomposition_mode == DecompositionMode::SplitK) {
       return 0;
     }
 
-    //
-    // The final wave is not full. Perform some stream-K work.
-    //
+    if (decomposition_mode == DecompositionMode::Heuristic) {
+      if (full_waves == total_waves || k_tiles_per_output_tile <= min_iters_per_sk_unit_) {
+        // All tiles will be data-parallel tiles if there is either no quantization
+        // or if there is no work to be split.
+        return 0;
+      }
 
-    // Rudimentary heuristic: prefer data-parallel decomposition if we have more than
-    // one wave and the tail wave is more than half full. This is subject to change.
-    if (full_waves != 0) {
+      //
+      // The final wave is not full. Perform some stream-K work.
+      //
+
+      // Rudimentary heuristic: prefer data-parallel decomposition if we have more than
+      // one wave and the tail wave is more than half full. This is subject to change.
       uint64_t tail_tiles = output_tiles - (full_waves * ctas_per_wave);
       if (tail_tiles >= (ctas_per_wave / 2)) {
         return 0;
@@ -729,6 +895,22 @@ struct PersistentTileSchedulerSm90StreamKParams {
     return static_cast<uint32_t>(output_tiles - dp_tiles);
   }
 
+  CUTLASS_HOST_DEVICE
+  static uint64_t
+  get_num_sk_units(GemmCoord cluster_shape, uint64_t ctas_per_wave, uint32_t sk_tiles, uint32_t k_tiles_per_output_tile) {
+    // Number of k iterations computed by the stream-K units as a whole
+    uint64_t k_tiles_sk_total = k_tiles_per_output_tile * sk_tiles;
+
+    // Calculate the number of stream-K units that would be needed if each stream-K unit
+    // computed the minimum allowable k iterations. Truncate this to be in units of clusters.
+    auto cluster_size = cluster_shape.m() * cluster_shape.n();
+    uint64_t min_sized_sk_units = (k_tiles_sk_total / min_iters_per_sk_unit_);
+    min_sized_sk_units = (min_sized_sk_units / cluster_size) * cluster_size;
+
+    uint64_t sk_units = platform::min(ctas_per_wave, min_sized_sk_units);
+    return sk_units;
+  }
+
   // Calculates the size of the workspace needed for holding reduction barriers
   CUTLASS_HOST_DEVICE
   static int
@@ -759,9 +941,11 @@ struct PersistentTileSchedulerSm90StreamKParams {
     int splits,
     int max_swizzle,
     RasterOrderOptions raster_order_option,
+    DecompositionMode decomposition_mode,
     uint32_t mma_warp_groups,
     uint32_t barrier_bits,
-    uint32_t accumulator_bits) {
+    uint32_t accumulator_bits,
+    uint32_t epilogue_subtile = 1) {
 
     auto log_swizzle_size = UnderlyingParams::get_log_swizzle_size(problem_blocks.x, problem_blocks.y, max_swizzle);
     problem_blocks.x = round_up(problem_blocks.x, (1 << log_swizzle_size) * cluster_shape.m());
@@ -771,7 +955,12 @@ struct PersistentTileSchedulerSm90StreamKParams {
     // of output tiles that will be split, and then calculate the workspace needed to cover these.
     uint64_t output_tiles = problem_blocks.x * problem_blocks.y * problem_blocks.z;
 
-    if (splits > 1) {
+    if (decomposition_mode == DecompositionMode::DataParallel) {
+      barrier_workspace_size = 0;
+      reduction_workspace_size = 0;
+    }
+    else if (decomposition_mode == DecompositionMode::SplitK ||
+        (decomposition_mode == DecompositionMode::Heuristic && splits > 1)) {
       // Basic split-K variant requires workspace for all output tiles
       barrier_workspace_size = get_barrier_workspace_size(output_tiles, mma_warp_groups, barrier_bits);
       reduction_workspace_size = get_reduction_workspace_size(output_tiles, tile_shape, accumulator_bits);
@@ -794,14 +983,39 @@ struct PersistentTileSchedulerSm90StreamKParams {
         raster_order_option
       );
       uint64_t ctas_per_wave = grid.x * grid.y;
-      uint32_t sk_tiles = get_num_sk_tiles(output_tiles, ctas_per_wave, static_cast<uint32_t>(k_tiles_per_output_tile));
+      uint32_t sk_tiles = get_num_sk_tiles(output_tiles, ctas_per_wave, static_cast<uint32_t>(k_tiles_per_output_tile), decomposition_mode);
+      uint64_t sk_units = get_num_sk_units(cluster_shape, ctas_per_wave, sk_tiles, k_tiles_per_output_tile);
+      uint64_t dp_tiles = output_tiles - sk_tiles;
 
+      uint64_t reduction_tiles = sk_tiles;
+      if (should_perform_separate_reduction(epilogue_subtile, sk_units, sk_tiles, dp_tiles, ctas_per_wave)) {
+        // In separate reduction, each peer writes to its own location in scratch space.
+        // Thus, for separate reduction, we need as many reduction tiles per output tile
+        // as there are the maximum number of peers that can collaborate on an output tile.
+        reduction_tiles *= max_peers_per_tile(sk_units, sk_tiles);
+      }
+
+      // Though separate reduction requires a larger reduction workspace, only one barrier
+      // is needed per output tile. Each peer will increment the barrier by one once the peer has
+      // written its accumulator to scratch space. The separate reduction unit will only begin
+      // performing the reduction when the barrier has reached the number of peers for the output tile.
       barrier_workspace_size = get_barrier_workspace_size(sk_tiles, mma_warp_groups, barrier_bits);
-      reduction_workspace_size = get_reduction_workspace_size(sk_tiles, tile_shape, accumulator_bits);
+      reduction_workspace_size = get_reduction_workspace_size(reduction_tiles, tile_shape, accumulator_bits);
     }
   }
   #endif // !defined(__CUDACC_RTC__)
 
+  // Returns whether the kernel is configured in a manner for which separate reduction should be used
+  CUTLASS_HOST_DEVICE
+  static bool
+  should_perform_separate_reduction(uint32_t epilogue_subtile, uint64_t sk_units, uint64_t sk_tiles, uint64_t dp_tiles, uint64_t ctas_per_wave) {
+    // We perform separate reduction if we have fewer than one wave of output tiles
+    // and each output tile is covered by at least to stream-K units. When sk_units is
+    // multiple of sk_tiles, will choose basic split-k path instead of separate reduction for now.
+    return (epilogue_subtile != 1) && (dp_tiles == 0) && (sk_units > 2u * sk_tiles) &&
+           (sk_units + sk_tiles * epilogue_subtile <= ctas_per_wave);
+  }
+
   // Get the amount of scratch workspace needed for the kernel. This variant of the method should only be used when
   // problem_shape and tile_shape contain modes of only rank 1.
   static int
@@ -813,9 +1027,11 @@ struct PersistentTileSchedulerSm90StreamKParams {
     int splits,
     int max_swizzle,
     RasterOrderOptions raster_order_option,
+    DecompositionMode decomposition_mode,
     uint32_t mma_warp_groups,
     uint32_t barrier_bits,
-    uint32_t element_accumulator_bits) {
+    uint32_t element_accumulator_bits,
+    uint32_t epilogue_subtile) {
 
     dim3 problem_blocks = UnderlyingParams::get_tiled_cta_shape_mnl(problem_shape, tile_shape, cluster_shape);
     uint32_t k_tiles_per_output_tile = (problem_shape.k() + tile_shape.k() - 1) / tile_shape.k();
@@ -829,9 +1045,11 @@ struct PersistentTileSchedulerSm90StreamKParams {
       splits,
       max_swizzle,
       raster_order_option,
+      decomposition_mode,
       mma_warp_groups,
       barrier_bits,
-      element_accumulator_bits
+      element_accumulator_bits,
+      epilogue_subtile
     );
   }
 
@@ -848,9 +1066,11 @@ struct PersistentTileSchedulerSm90StreamKParams {
     int splits,
     int max_swizzle,
     RasterOrderOptions raster_order_option,
+    DecompositionMode decomposition_mode,
     uint32_t mma_warp_groups,
     uint32_t barrier_bits,
-    uint32_t element_accumulator_bits) {
+    uint32_t element_accumulator_bits,
+    uint32_t epilogue_subtile = 1) {
 
     int barrier_workspace_size = 0;
     int reduction_workspace_size = 0;
@@ -867,9 +1087,11 @@ struct PersistentTileSchedulerSm90StreamKParams {
         splits,
         max_swizzle,
         raster_order_option,
+        decomposition_mode,
         mma_warp_groups,
         barrier_bits,
-        element_accumulator_bits
+        element_accumulator_bits,
+        epilogue_subtile
       );
     #endif
 
@@ -889,9 +1111,11 @@ struct PersistentTileSchedulerSm90StreamKParams {
     int splits,
     int max_swizzle,
     RasterOrderOptions raster_order_option,
+    DecompositionMode decomposition_mode,
     uint32_t mma_warp_groups,
     uint32_t barrier_bits,
-    uint32_t element_accumulator_bits) {
+    uint32_t element_accumulator_bits,
+    uint32_t epilogue_subtile) {
 
     dim3 problem_blocks = UnderlyingParams::get_tiled_cta_shape_mnl(problem_shape, tile_shape, cluster_shape);
     uint32_t k_tiles_per_output_tile = (problem_shape.k() + tile_shape.k() - 1) / tile_shape.k();
@@ -907,9 +1131,11 @@ struct PersistentTileSchedulerSm90StreamKParams {
       splits,
       max_swizzle,
       raster_order_option,
+      decomposition_mode,
       mma_warp_groups,
       barrier_bits,
-      element_accumulator_bits
+      element_accumulator_bits,
+      epilogue_subtile
     );
   }
 
@@ -928,9 +1154,11 @@ struct PersistentTileSchedulerSm90StreamKParams {
     int splits,
     int max_swizzle,
     RasterOrderOptions raster_order_option,
+    DecompositionMode decomposition_mode,
     uint32_t mma_warp_groups,
     uint32_t barrier_bits,
-    uint32_t element_accumulator_bits) {
+    uint32_t element_accumulator_bits,
+    uint32_t epilogue_subtile = 1) {
 
     #if !defined(__CUDACC_RTC__)
       int barrier_workspace_size = 0;
@@ -947,9 +1175,11 @@ struct PersistentTileSchedulerSm90StreamKParams {
         splits,
         max_swizzle,
         raster_order_option,
+        decomposition_mode,
         mma_warp_groups,
         barrier_bits,
-        element_accumulator_bits
+        element_accumulator_bits,
+        epilogue_subtile
       );
 
       if (barrier_workspace_size > 0) {
@@ -982,6 +1212,7 @@ struct PersistentTileSchedulerSm90StreamKParams {
     divmod_cluster_shape_minor_ = underlying_params.divmod_cluster_shape_minor_;
     divmod_batch_ = FastDivmodU64(blocks_m * blocks_n);
     divmod_tiles_per_output_tile_ = FastDivmod(k_tiles_per_output_tile);
+    divmod_sk_groups_ = FastDivmodU64(1u);
     auto cluster_size = underlying_params.divmod_cluster_shape_major_.divisor * underlying_params.divmod_cluster_shape_minor_.divisor;
     divmod_clusters_mnl_ = FastDivmodU64((blocks_m * blocks_n * blocks_l) / cluster_size);
     splits_ = splits;
@@ -997,9 +1228,11 @@ struct PersistentTileSchedulerSm90StreamKParams {
     // No stream-K work is performed for "basic" data-parallel and split-K decompositions
     sk_tiles_ = 0;
     sk_units_ = 0;
+    divmod_sk_units_per_group_ = FastDivmodU64(1u);
+    separate_reduction_units_ = 0;
   }
 
-private:
+  private:
   // Round up number of bytes to the nearest multiple of L2 cache line alignment
   CUTLASS_HOST_DEVICE
   static int
@@ -1009,6 +1242,236 @@ private:
   }
 };
 
+////////////////////////////////////////////////////////////////////////////////
+
+// Parameters for SM90 persistent group scheduler (only used for Grouped Gemms)
+template<class ProblemShape>
+struct PersistentTileSchedulerSm90GroupParams {
+
+  enum class RasterOrder {
+    AlongM,
+    AlongN
+  };
+
+  enum class RasterOrderOptions {
+    Heuristic,
+    AlongM,
+    AlongN
+  };
+
+  FastDivmodU64Pow2 divmod_cluster_shape_major_{};
+  FastDivmodU64Pow2 divmod_cluster_shape_minor_{};
+  FastDivmodU64 divmod_batch_{};
+
+  uint64_t blocks_per_problem_ = 0;
+  int32_t log_swizzle_size_ = 0;
+  RasterOrder raster_order_ = RasterOrder::AlongN;
+
+  int32_t groups_ = 0;
+  ProblemShape* problem_shapes_ = nullptr;
+  GemmCoord cta_shape_;
+
+  // Version of initialize that takes in as input the number of CTAs in the M and N and L dimensions.
+  // This is useful for calculating the tiled shape when a mode of problem and/or CTA shape has rank > 1,
+  // for which using CuTe algebra for calculating tile shapes is easiest.
+  void
+  initialize(
+    dim3 problem_blocks,
+    int32_t groups,
+    ProblemShape* problem_shapes,
+    GemmCoord cta_shape,
+    GemmCoord cluster_shape,
+    KernelHardwareInfo const& hw_info,
+    int max_swizzle_size,
+    RasterOrderOptions raster_order_option
+  ) {
+    
+    CUTLASS_UNUSED(hw_info);
+    
+    // Round up to nearest multiple of swizzle_size along each mode
+    auto log_swizzle_size = get_log_swizzle_size(problem_blocks.x, problem_blocks.y, max_swizzle_size);
+    auto problem_blocks_m = round_up(problem_blocks.x, (1 << log_swizzle_size) * cluster_shape.m());
+    auto problem_blocks_n = round_up(problem_blocks.y, (1 << log_swizzle_size) * cluster_shape.n());
+
+    RasterOrder raster_order = get_rasterization_order(
+      problem_blocks_m,
+      problem_blocks_n,
+      raster_order_option
+    );
+
+    //
+    // Set members
+    //
+    groups_ = groups;
+    problem_shapes_ = problem_shapes;
+    cta_shape_ = cta_shape;
+
+    blocks_per_problem_ = problem_blocks_m * problem_blocks_n * problem_blocks.z;
+    log_swizzle_size_ = log_swizzle_size;
+    raster_order_ = raster_order;
+    divmod_batch_ = FastDivmodU64(problem_blocks_m * problem_blocks_n);
+
+    if (raster_order == RasterOrder::AlongN) {
+      divmod_cluster_shape_major_ = FastDivmodU64Pow2(cluster_shape.n());
+      divmod_cluster_shape_minor_ = FastDivmodU64Pow2(cluster_shape.m());
+    }
+    else {
+      divmod_cluster_shape_major_ = FastDivmodU64Pow2(cluster_shape.m());
+      divmod_cluster_shape_minor_ = FastDivmodU64Pow2(cluster_shape.n());
+    }
+  }
+
+  // Version of get_tiled_cta_shape_mnl that takes in as input the number of CTAs in the M and N dimensions.
+  // This is useful for calculating the tiled shape when a mode of problem and/or CTA shape has rank > 1,
+  // for which using CuTe algebra for calculating tile shapes is easiest.
+  CUTLASS_HOST_DEVICE
+  static dim3
+  get_tiled_cta_shape_mnl(GemmCoord cluster_shape, uint32_t cta_m, uint32_t cta_n) {
+    // Round up to nearest multiple of cluster dim along each mode
+    auto problem_blocks_m = ((cta_m + cluster_shape.m() - 1) / cluster_shape.m()) * cluster_shape.m();
+    auto problem_blocks_n = ((cta_n + cluster_shape.n() - 1) / cluster_shape.n()) * cluster_shape.n();
+
+    return {
+      static_cast<uint32_t>(problem_blocks_m),
+      static_cast<uint32_t>(problem_blocks_n),
+      static_cast<uint32_t>(1) // Only a single batch per group is currently supported
+    };
+  }
+
+  // Version of get_grid_shape that takes in as input the number of CTAs in the M and N and L dimensions.
+  // This is useful for calculating the tiled shape when a mode of problem and/or CTA shape has rank > 1,
+  // for which using CuTe algebra for calculating tile shapes is easiest.
+  CUTLASS_HOST_DEVICE static
+  dim3
+  get_grid_shape(
+    dim3 problem_blocks,
+    GemmCoord cluster_shape,
+    KernelHardwareInfo hw_info,
+    int max_swizzle_size,
+    RasterOrderOptions raster_order_option,
+    bool truncate_by_problem_size=true) {
+
+    int const sm_count = hw_info.sm_count;
+
+    // Round up to nearest multiple of swizzle_size along each mode
+    auto log_swizzle_size = get_log_swizzle_size(problem_blocks.x, problem_blocks.y, max_swizzle_size);
+    auto problem_blocks_m = round_up(problem_blocks.x, (1 << log_swizzle_size) * cluster_shape.m());
+    auto problem_blocks_n = round_up(problem_blocks.y, (1 << log_swizzle_size) * cluster_shape.n());
+
+    int problem_blocks_total = problem_blocks_m * problem_blocks_n * problem_blocks.z;
+
+    RasterOrder raster_order = get_rasterization_order(
+      problem_blocks_m,
+      problem_blocks_n,
+      raster_order_option
+    );
+
+    dim3 launch_grid;
+
+    if (raster_order == RasterOrder::AlongN) {
+      launch_grid = dim3(cluster_shape.m(), 1, 1);
+    }
+    else {
+      launch_grid = dim3(1, cluster_shape.n(), 1);
+    }
+
+    auto possibly_truncate = [&](int x, int y) {
+      if (truncate_by_problem_size) {
+        return platform::min(x, y);
+      }
+      else {
+        return x;
+      }
+    };
+
+    // The else path is generic, however, we can avoid some divs if we know cluster size is 1
+    auto cluster_size = cluster_shape.m() * cluster_shape.n();
+    if (cluster_size == 1) {
+      if (raster_order == RasterOrder::AlongN) {
+        launch_grid.y = possibly_truncate(sm_count, problem_blocks_total);
+      }
+      else {
+        launch_grid.x = possibly_truncate(sm_count, problem_blocks_total);
+      }
+    }
+    else {
+      // Optimal grid size calculation is based on
+      // GH100: 8 GPCs, 72 TPCs (9 TPCs/GPC), 2 SMs/TPC, 144 SMs per full GPU
+      // Hence, maximum SMs per GPC = 18
+      constexpr int max_sm_per_gpc = 18;
+      // Provided SM count could possibly be less than the assumed maximum SMs per GPC
+      auto cluster_size = cluster_shape.m() * cluster_shape.n();
+      int const min_num_gpc = sm_count < max_sm_per_gpc ? 1 : sm_count / max_sm_per_gpc;
+      int const max_cta_occupancy_per_gpc = max_sm_per_gpc - (max_sm_per_gpc % cluster_size);
+      int cta_per_device = min_num_gpc * max_cta_occupancy_per_gpc;
+
+      // The calculation below allows for larger grid size launch for different GPUs.
+      int const num_gpc_residual = sm_count < max_sm_per_gpc ? 0 : sm_count % max_sm_per_gpc;
+      int const max_cta_occupancy_per_residual_gpc = num_gpc_residual - (num_gpc_residual % cluster_size);
+      cta_per_device += max_cta_occupancy_per_residual_gpc;
+
+      cta_per_device = sm_count < cta_per_device ? sm_count : cta_per_device;
+
+      if (raster_order == RasterOrder::AlongN) {
+        launch_grid.y = possibly_truncate(
+            cta_per_device       / cluster_shape.m(),
+            problem_blocks_total / cluster_shape.m());
+      }
+      else {
+        launch_grid.x = possibly_truncate(
+            cta_per_device       / cluster_shape.n(),
+            problem_blocks_total / cluster_shape.n());
+      }
+    }
+    return launch_grid;
+  }
+
+  CUTLASS_HOST_DEVICE
+  static int32_t
+  get_log_swizzle_size(int problem_ctas_m, int problem_ctas_n, int max_swizzle_size) {
+    int min_cta_dim = platform::min(problem_ctas_m, problem_ctas_n);
+    if (max_swizzle_size >= 8 && min_cta_dim >= 6) {
+      return 3;
+    }
+    else if (max_swizzle_size >= 4 && min_cta_dim >= 3) {
+      return 2;
+    }
+    else if (max_swizzle_size >= 2 && min_cta_dim >= 2) {
+      return 1;
+    }
+    else {
+      return 0;
+    }
+  }
+
+  CUTLASS_HOST_DEVICE
+  static RasterOrder
+  get_rasterization_order(
+    uint32_t tiles_m,
+    uint32_t tiles_n,
+    RasterOrderOptions raster_order_option
+  ) {
+
+    if (raster_order_option == RasterOrderOptions::Heuristic) {
+      if (tiles_n > tiles_m) {
+        return RasterOrder::AlongM;
+      }
+      else {
+        return RasterOrder::AlongN;
+      }
+    }
+    else {
+      switch (raster_order_option) {
+        case RasterOrderOptions::AlongN:
+          return RasterOrder::AlongN;
+          break;
+        default:
+          return RasterOrder::AlongM;
+      }
+    }
+  }
+};
+
 ////////////////////////////////////////////////////////////////////////////////
 } // namespace detail
 } // namespace kernel