3.6.0 update (#2005)

* 3.6.0 update

* doc and swap stuff

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
Co-authored-by: Haicheng Wu <haichengw@nvidia.com>

include/cutlass/arch/barrier.h

@@ -47,6 +47,99 @@ namespace cutlass {
 namespace arch {
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
+CUTLASS_DEVICE void fence_view_async_shared();
+
+namespace detail { // namespace detail begin
+
+// Single threaded versions that need to be called in an elect_one region
+template<typename T, uint32_t Stages>
+CUTLASS_DEVICE
+void initialize_barrier_array(T ptr, int arv_cnt) {
+  CUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < Stages; i++) {
+    ptr[i].init(arv_cnt);
+  }
+}
+
+template<typename T, uint32_t Stages>
+CUTLASS_DEVICE
+void initialize_barrier_array(uint64_t *ptr, int arv_cnt) {
+  CUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < Stages; i++) {
+    T::init(&ptr[i], arv_cnt);
+  }
+}
+
+template<typename FullBarrier, typename EmptyBarrier, uint32_t Stages>
+CUTLASS_DEVICE
+void initialize_barrier_array_pair(FullBarrier full_barriers, EmptyBarrier empty_barriers, int full_barrier_arv_cnt, int empty_barrier_arv_cnt) {
+  CUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < Stages; i++) {
+    full_barriers[i].init(full_barrier_arv_cnt);
+    empty_barriers[i].init(empty_barrier_arv_cnt);
+  }
+}
+
+template<typename FullBarrier, typename EmptyBarrier, uint32_t Stages>
+CUTLASS_DEVICE
+void initialize_barrier_array_pair(uint64_t *full_barriers_ptr, uint64_t *empty_barriers_ptr, int full_barrier_arv_cnt, int empty_barrier_arv_cnt) {
+  CUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < Stages; i++) {
+    FullBarrier::init(&full_barriers_ptr[i], full_barrier_arv_cnt);
+    EmptyBarrier::init(&empty_barriers_ptr[i], empty_barrier_arv_cnt);
+  }
+}
+
+// Aligned versions that need to be call warp wide
+template<typename T, uint32_t Stages>
+CUTLASS_DEVICE
+void initialize_barrier_array_aligned(T ptr, int arv_cnt) {
+  if(cute::elect_one_sync()) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Stages; i++) {
+      ptr[i].init(arv_cnt);
+    }
+  }
+}
+
+template<typename T, uint32_t Stages>
+CUTLASS_DEVICE
+void initialize_barrier_array_aligned(uint64_t *ptr, int arv_cnt) {
+  if(cute::elect_one_sync()) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Stages; i++) {
+      T::init(&ptr[i], arv_cnt);
+    }
+  }
+}
+
+template<typename FullBarrier, typename EmptyBarrier, uint32_t Stages>
+CUTLASS_DEVICE
+void initialize_barrier_array_pair_aligned(FullBarrier full_barriers, EmptyBarrier empty_barriers, int full_barrier_arv_cnt, int empty_barrier_arv_cnt) {
+  if(cute::elect_one_sync()) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Stages; i++) {
+      full_barriers[i].init(full_barrier_arv_cnt);
+      empty_barriers[i].init(empty_barrier_arv_cnt);
+    }
+  }
+}
+
+template<typename FullBarrier, typename EmptyBarrier, uint32_t Stages>
+CUTLASS_DEVICE
+void initialize_barrier_array_pair_aligned(uint64_t *full_barriers_ptr, uint64_t *empty_barriers_ptr, int full_barrier_arv_cnt, int empty_barrier_arv_cnt) {
+  if(cute::elect_one_sync()) {
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Stages; i++) {
+      FullBarrier::init(&full_barriers_ptr[i], full_barrier_arv_cnt);
+      EmptyBarrier::init(&empty_barriers_ptr[i], empty_barrier_arv_cnt);
+    }
+  }
+}
+
+} // namespace detail end
+
+
 // Enumerates the reserved named barriers to avoid potential conflicts
 // This enum class specifies the NamedBarriers reserved by CUTLASS.
 enum class ReservedNamedBarriers { 

include/cutlass/arch/config.h

@@ -35,6 +35,8 @@
 
 #pragma once
 
+#include "cutlass/platform/platform.h"
+
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 // SM90
@@ -79,3 +81,5 @@
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
+/////////////////////////////////////////////////////////////////////////////////////////////////
+

include/cutlass/arch/memory_sm75.h

@@ -35,6 +35,7 @@
 #pragma once
 
 #include "cutlass/array.h"
+#include "cutlass/detail/helper_macros.hpp"
 #include "cutlass/layout/matrix.h"
 #include "cute/arch/copy_sm75.hpp"
 #include "cute/arch/util.hpp"
@@ -50,7 +51,7 @@ template <
   /// .x1, .x2, or .x4
   int MatrixCount
 >
-inline __device__ void ldsm(Array<unsigned, MatrixCount> & D, void const* ptr);
+CUTLASS_DEVICE void ldsm(Array<unsigned, MatrixCount> & D, void const* ptr);
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 //
@@ -59,19 +60,19 @@ inline __device__ void ldsm(Array<unsigned, MatrixCount> & D, void const* ptr);
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 /// CUTLASS helper to get SMEM pointer
-inline __device__ unsigned cutlass_get_smem_pointer(void *ptr) {
+CUTLASS_DEVICE unsigned cutlass_get_smem_pointer(void *ptr) {
   return cute::cast_smem_ptr_to_uint(ptr);
 }
 
 /// CUTLASS helper to get SMEM pointer
-inline __device__ unsigned cutlass_get_smem_pointer(void const *ptr) {
+CUTLASS_DEVICE unsigned cutlass_get_smem_pointer(void const *ptr) {
   return cutlass_get_smem_pointer(const_cast<void *>(ptr));
 }
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 template <>
-inline __device__ void ldsm<layout::RowMajor, 1>(
+CUTLASS_DEVICE void ldsm<layout::RowMajor, 1>(
     Array<unsigned, 1> & D,
     void const* ptr) {
 
@@ -95,7 +96,7 @@ inline __device__ void ldsm<layout::RowMajor, 1>(
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 template <>
-inline __device__ void ldsm<layout::RowMajor, 2>(
+CUTLASS_DEVICE void ldsm<layout::RowMajor, 2>(
     Array<unsigned, 2> & D,
     void const* ptr) {
 
@@ -119,7 +120,7 @@ inline __device__ void ldsm<layout::RowMajor, 2>(
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 template <>
-inline __device__ void ldsm<layout::RowMajor, 4>(
+CUTLASS_DEVICE void ldsm<layout::RowMajor, 4>(
     Array<unsigned, 4> & D,
     void const* ptr) {
 
@@ -147,7 +148,7 @@ inline __device__ void ldsm<layout::RowMajor, 4>(
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 template <>
-inline __device__ void ldsm<layout::ColumnMajor, 1>(
+CUTLASS_DEVICE void ldsm<layout::ColumnMajor, 1>(
     Array<unsigned, 1> & D,
     void const* ptr) {
 
@@ -171,7 +172,7 @@ inline __device__ void ldsm<layout::ColumnMajor, 1>(
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 template <>
-inline __device__ void ldsm<layout::ColumnMajor, 2>(
+CUTLASS_DEVICE void ldsm<layout::ColumnMajor, 2>(
     Array<unsigned, 2> & D,
     void const* ptr) {
 
@@ -195,7 +196,7 @@ inline __device__ void ldsm<layout::ColumnMajor, 2>(
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 template <>
-inline __device__ void ldsm<layout::ColumnMajor, 4>(
+CUTLASS_DEVICE void ldsm<layout::ColumnMajor, 4>(
     Array<unsigned, 4> & D,
     void const* ptr) {
 

include/cutlass/arch/mma_sm70.h

@@ -33,11 +33,7 @@
 */
 #pragma once
 
-#if defined(__CUDACC_RTC__)
 #include <cuda/std/cassert>
-#else
-#include <assert.h>
-#endif
 
 #include "mma.h"
 #include "cutlass/layout/matrix.h"

include/cutlass/arch/mma_sm75.h

@@ -34,11 +34,7 @@
 
 #pragma once
 
-#if defined(__CUDACC_RTC__)
 #include <cuda/std/cassert>
-#else
-#include <assert.h>
-#endif
 
 #include "cutlass/arch/wmma.h"
 

include/cutlass/arch/mma_sm80.h

@@ -34,11 +34,7 @@
 
 #pragma once
 
-#if defined(__CUDACC_RTC__)
 #include <cuda/std/cassert>
-#else
-#include <assert.h>
-#endif
 
 #include "cutlass/cutlass.h"
 #include "mma.h"

include/cutlass/arch/mma_sm89.h

@@ -35,11 +35,7 @@
 
 #pragma once
 
-#if defined(__CUDACC_RTC__)
 #include <cuda/std/cassert>
-#else
-#include <assert.h>
-#endif
 
 #include "cutlass/cutlass.h"
 #include "mma.h"

include/cutlass/arch/mma_sm90.h

@@ -34,11 +34,7 @@
 
 #pragma once
 
-#if defined(__CUDACC_RTC__)
 #include <cuda/std/cassert>
-#else
-#include <assert.h>
-#endif
 
 #include "mma.h"
 #include "cutlass/layout/matrix.h"

include/cutlass/arch/mma_sparse_sm80.h

@@ -35,11 +35,7 @@
 
 #pragma once
 
-#if defined(__CUDACC_RTC__)
 #include <cuda/std/cassert>
-#else
-#include <assert.h>
-#endif
 
 #include "mma.h"
 #include "cutlass/layout/matrix.h"

include/cutlass/arch/mma_sparse_sm89.h

@@ -35,11 +35,7 @@
 
 #pragma once
 
-#if defined(__CUDACC_RTC__)
 #include <cuda/std/cassert>
-#else
-#include <assert.h>
-#endif
 
 #include "mma.h"
 #include "cutlass/layout/matrix.h"

include/cutlass/arch/simd.h

@@ -34,8 +34,8 @@
 
 #pragma once
 
-#include "../array.h"
-#include "../numeric_types.h"
+#include "cutlass/arch/array.h"
+#include "cutlass/arch/numeric_types.h"
 
 namespace cutlass {
 namespace arch {

include/cutlass/arch/synclog.hpp

@@ -59,7 +59,7 @@ constexpr uint32_t synclog_cap = 1 << 26;
 inline std::mutex synclog_mutex;
 inline std::vector<uint32_t*> synclog_buf_list;
 #if defined(__NVCC__) || (defined(__clang__) && defined(__CUDA__))
-inline __device__ uint32_t* synclog_buf;
+CUTLASS_DEVICE uint32_t* synclog_buf;
 #endif
 
 CUTLASS_DEVICE

include/cutlass/arch/wmma_sm70.h

@@ -34,11 +34,7 @@
 
 #pragma once
 
-#if defined(__CUDACC_RTC__)
 #include <cuda/std/cassert>
-#else
-#include <assert.h>
-#endif
 #include "cutlass/layout/matrix.h"
 
 ////////////////////////////////////////////////////////////////////////////////

include/cutlass/arch/wmma_sm72.h

@@ -34,11 +34,7 @@
 
 #pragma once
 
-#if defined(__CUDACC_RTC__)
 #include <cuda/std/cassert>
-#else
-#include <assert.h>
-#endif
 #include "cutlass/layout/matrix.h"
 
 ////////////////////////////////////////////////////////////////////////////////

include/cutlass/arch/wmma_sm75.h

@@ -34,11 +34,7 @@
 
 #pragma once
 
-#if defined(__CUDACC_RTC__)
 #include <cuda/std/cassert>
-#else
-#include <assert.h>
-#endif
 #include "cutlass/layout/matrix.h"
 
 ////////////////////////////////////////////////////////////////////////////////

include/cutlass/array.h

@@ -2573,20 +2573,8 @@ Array<T, N> fma(Array<T, N> const &a, Array<T, N> const &b, T c) {
   return op(a, b, c);
 }
 
-
 ////////////////////////////////////////////////////////////////////////////////////////////////////
-
-
-
-} // namespace cutlass
-
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-#include "cutlass/array_subbyte.h"
-
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
-namespace cutlass {
+  
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 // AlignedArray
@@ -2606,9 +2594,10 @@ public:
 
 };
 
-////////////////////////////////////////////////////////////////////////////////////////////////////
-
 } // namespace cutlass
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
+#include "cutlass/array_subbyte.h"
+
+////////////////////////////////////////////////////////////////////////////////////////////////////

include/cutlass/array_subbyte.h

@@ -554,6 +554,8 @@ private:
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
 } // namespace cutlass
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////

include/cutlass/blas3.h

@@ -132,7 +132,7 @@ struct MantissaInBits<double> {
 template <>
 struct MantissaInBits<cutlass::complex<double>> {
   static int constexpr bits = 30;
-  static double constexpr error = 1.0e-15;
+  static double constexpr error = 1.0e-14;
 };
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////

include/cutlass/conv/collective/sm90_implicit_gemm_gmma_ss_warpspecialized.hpp

@@ -189,7 +189,7 @@ private:
             -problem_shape.dilation[NumSpatialDimensions-1-i] :
             problem_shape.dilation[NumSpatialDimensions-1-i];
       }
-  
+
       return make_im2col_tma_copy(
           GmemTiledCopyA{},
           tensor_a,
@@ -225,7 +225,7 @@ private:
       auto lower_corner_whd = detail::compute_lower_corner_whd(problem_shape);
       auto upper_corner_whd = detail::compute_upper_corner_whd(problem_shape);
       auto lower_srt = detail::compute_lower_srt(problem_shape);
-  
+
       return make_im2col_tma_copy(
           GmemTiledCopyB{},
           tensor_b,
@@ -372,6 +372,96 @@ public:
       return false;
     }
 
+    if (is_im2col_A || is_im2col_B) {
+      // Check valid corner values for TMA_LOAD_IM2COL, signed int ranging from [-corner_limit, corner_limit - 1]
+      constexpr int32_t corner_limit = 1 << (16 / NumSpatialDimensions - 1);
+      auto lower_corner_whd = detail::compute_lower_corner_whd(problem_shape);
+      for (int i = 0; i < problem_shape.RankS; ++i) {
+        implementable = implementable && lower_corner_whd[i] >= -corner_limit && lower_corner_whd[i] <= (corner_limit - 1);
+      }
+      auto upper_corner_whd = detail::compute_upper_corner_whd(problem_shape);
+      for (int i = 0; i < problem_shape.RankS; ++i) {
+        implementable = implementable && upper_corner_whd[i] >= -corner_limit && upper_corner_whd[i] <= (corner_limit - 1);
+      }
+
+      if (!implementable) {
+        CUTLASS_TRACE_HOST("  CAN IMPLEMENT: Padding values don't meet requirements for TMA LOAD IM2COL.\n");
+        return false;
+      }
+    }
+
+    // Wgrad kernels don't support non-packed output strides, non-packed tensor A stride (linearized)
+    if constexpr (ConvOp == conv::Operator::kWgrad) {
+#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+      std::ostringstream os;
+#endif        
+      const auto & input_shape  = problem_shape.shape_A;
+      const auto & input_stride  = problem_shape.stride_A;
+
+      implementable &= input_stride[ProblemShape::RankT - 1] == 1;
+      int input_shape_size = 1;
+      for (int i = ProblemShape::RankT - 2; i >= 0; --i) {
+        input_shape_size *= input_shape[i + 1];
+        implementable &= input_stride[i] == input_shape_size;
+#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+        if (input_stride[i] != input_shape_size) {
+          os << "\n    *** input_stride[" << i << "] = " << input_stride[i] << " != input_shape_size = " << input_shape_size << " ***";
+        }
+#endif
+      }
+
+      if (!implementable) {
+#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+        os << "\n    input_shape_size: " << input_shape_size
+           << "\n    input_shape: " << input_shape
+           << "\n    input_stride: " << input_stride
+           << "\n";
+#endif        
+        CUTLASS_TRACE_HOST("  CAN IMPLEMENT: Wgrad kernels don't support non-packed input strides.\n");
+#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+        CUTLASS_TRACE_HOST(os.str());
+#endif        
+        return false;
+      }
+
+      const auto & output_shape  = problem_shape.shape_C;
+      const auto & output_stride  = problem_shape.stride_C;
+
+      implementable &= output_stride[ProblemShape::RankT - 1] == 1;
+      int output_shape_size = 1;
+      for (int i = ProblemShape::RankT - 2; i >= 0; --i) {
+        output_shape_size *= output_shape[i + 1];
+        implementable &= output_stride[i] == output_shape_size;
+#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+        if (output_stride[i] != output_shape_size) {
+          os << "\n    *** output_stride[" << i << "] = " << output_stride[i] << " != output_shape_size = " << output_shape_size << " ***";
+        }
+#endif
+      }
+
+      if (!implementable) {
+#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+        os << "\n    output_shape_size: " << input_shape_size
+           << "\n    output_shape: " << input_shape
+           << "\n    output_stride: " << input_stride
+           << "\n";
+#endif
+        CUTLASS_TRACE_HOST("  CAN IMPLEMENT: Wgrad kernels don't support non-packed output strides.\n");
+#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+        CUTLASS_TRACE_HOST(os.str());
+#endif        
+        return false;
+      }
+    }
+
+    // Conv kernels only support cross correlation mode currently.
+    implementable &= problem_shape.mode == cutlass::conv::Mode::kCrossCorrelation;
+
+    if (!implementable) {
+      CUTLASS_TRACE_HOST("  CAN IMPLEMENT: Conv kernels only support cross correlation mode currently.\n");
+      return false;
+    }
+
     if (problem_shape.groups > 1) {
       CUTLASS_TRACE_HOST("  CAN IMPLEMENT: This kernel does not support conv groups > 1.\n");
       return false;
@@ -516,9 +606,9 @@ public:
     // Issue the epilogue waits
     if (lane_predicate) {
       /* This helps avoid early exit of blocks in Cluster
-       * Waits for all stages to either be released (all 
+       * Waits for all stages to either be released (all
        * Consumer UNLOCKs), or if the stage was never used
-       * then would just be acquired since the phase was 
+       * then would just be acquired since the phase was
        * still inverted from make_producer_start_state
        */
       pipeline.producer_tail(smem_pipe_producer_state);
@@ -645,7 +735,7 @@ public:
     k_tile_count -= prologue_mma_count;
 
     smem_pipe_release.advance(k_tile_count);
-    
+
     // Wait on all GMMAs to complete
     warpgroup_wait<0>();
 

include/cutlass/conv/convnd_problem_shape.hpp

@@ -319,6 +319,7 @@ struct ConvProblemShape {
   // |   ShapeB     | KTRSC  | KTRSC  | NDHWC |
   // |   ShapeC     | NZPQK  | NDHWC  | KTRSC |
   //
+  // Input comes from calculate_xformed_act, which does NOT depend on ConvOp.
   CUTLASS_HOST_DEVICE
   constexpr void
   set_shape_stride_ABC(
@@ -328,6 +329,31 @@ struct ConvProblemShape {
     TensorStride stride_flt,
     TensorExtent shape_xformed_act,
     TensorStride stride_xformed_act) {
+#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+    printf("*** set_shape_stride_ABC ***");
+    printf("\n  shape_act: ");
+    print(shape_act);
+    printf("\n  stride_act: ");
+    print(stride_act);
+    printf("\n  shape_flt: ");
+    print(shape_flt);
+    printf("\n  stride_flt: ");
+    print(stride_flt);
+    printf("\n  shape_xformed_act: ");
+    print(shape_xformed_act);
+    printf("\n  stride_xformed_act: ");
+    print(stride_xformed_act);
+    if constexpr (ConvOp == cutlass::conv::Operator::kFprop) {
+      printf("\n  ConvOp: Fprop");
+    }
+    if constexpr (ConvOp == cutlass::conv::Operator::kDgrad) {
+      printf("\n  ConvOp: Dgrad");
+    }
+    if constexpr (ConvOp == cutlass::conv::Operator::kWgrad) {
+      printf("\n  ConvOp: Wgrad");
+    }
+    printf("\n");
+#endif
 
     if constexpr (ConvOp == cutlass::conv::Operator::kFprop) {
       shape_A = shape_act;
@@ -353,6 +379,20 @@ struct ConvProblemShape {
       shape_C = shape_flt;
       stride_C = stride_flt;
     }
+#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && (CUTLASS_DEBUG_TRACE_LEVEL > 1)
+    printf("\n  shape_A: ");
+    print(shape_A);
+    printf("\n  stride_A: ");
+    print(stride_A);
+    printf("\n  shape_B: ");
+    print(shape_B);
+    printf("\n  stride_B: ");
+    print(stride_B);
+    printf("\n  shape_C: ");
+    print(shape_C);
+    printf("\n  stride_C: ");
+    print(stride_C);
+#endif
   }
 
   // Get A extents.

include/cutlass/conv/kernel/direct_convolution.h

@@ -40,6 +40,7 @@
 #include "cutlass/array.h"
 #include "cutlass/numeric_types.h"
 #include "cutlass/matrix_shape.h"
+#include "cutlass/platform/platform.h"
 #include "cutlass/semaphore.h"
 #include "cutlass/tensor_ref.h"
 #include "cutlass/layout/tensor.h"
@@ -155,7 +156,7 @@ struct DirectConvolutionParams {
     swizzle_log_tile = threadblock_swizzle.get_log_tile(grid_tiled_shape);
 
     // Dynamic SMEM usage because stride and dilation are runtime params.
-    smem_size_ = (max(iterator_A.activation_size, int(sizeof(typename Epilogue::SharedStorage))) * kStages + iterator_B.filter_size);
+    smem_size_ = (cutlass::platform::max(iterator_A.activation_size, int(sizeof(typename Epilogue::SharedStorage))) * kStages + iterator_B.filter_size);
   }
 
   CUTLASS_HOST_DEVICE

include/cutlass/coord.h

@@ -37,7 +37,7 @@
 #if defined(__CUDACC_RTC__)
 #include <cuda/std/cstdint>
 #else
-#include <stdint.h>
+#include <cstdint>
 #endif
 
 #include "cutlass/cutlass.h"

include/cutlass/cuda_host_adapter.hpp

@@ -85,7 +85,11 @@ namespace cutlass {
 
 #if !defined(__CUDACC_RTC__)
 
+#if ((__CUDACC_VER_MAJOR__ >= 12) ||                               \
+    ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 8)))
 #include <cudaTypedefs.h>
+#endif // (__CUDACC_VERSION__ >= 11.8)
+
 #include <driver_types.h>
 
 #define CUTLASS_CUDA_DRIVER_STRINGIFY(tok) #tok
@@ -100,7 +104,8 @@ namespace cutlass {
 
 #else // defined(CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL)
 
-#if (__CUDACC_VER_MAJOR__ >= 12 && __CUDACC_VER_MINOR__ >= 5)
+#if ((__CUDACC_VER_MAJOR__ >= 13) ||                               \
+    ((__CUDACC_VER_MAJOR__ == 12) && (__CUDACC_VER_MINOR__ >= 5))) \
 
 #define CUTLASS_CUDA_DRIVER_WRAPPER_DECL(func, ver)             \
   template <typename... Args>                                   \
@@ -138,7 +143,7 @@ namespace cutlass {
     return reinterpret_cast<PFN_##func>(pfn)(args...);          \
   }
 
-#endif // (__CUDACC_VER_MAJOR__ >= 12 && __CUDACC_VER_MINOR__ >= 5)
+#endif // (__CUDACC_VERSION__ >= 12.5)
 
 #endif // defined(CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL)
 

include/cutlass/detail/collective.hpp

@@ -31,6 +31,7 @@
 #pragma once
 
 #include "cute/container/tuple.hpp"
+#include "cute/layout.hpp" // cute::size(shape)
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 namespace cutlass::gemm::collective {

include/cutlass/detail/collective/mixed_input_utils.hpp

@@ -237,7 +237,7 @@ struct LayoutAwareConvertImpl<
   }
 };
 
-// Specialization for UINT4 -> FPF16 with [02461357] value order
+// Specialization for UINT4 -> FP16 with [02461357] value order
 template <>
 struct LayoutAwareConvertImpl<
   cutlass::uint4b_t,
@@ -754,7 +754,6 @@ public:
     cute::tuple<Ts...>& partitioned_extra_info,
     int const k_block) {
 
-    
     static_assert(is_rmem<EngineIn>::value, "Input tensor for A conversion must come from registers");
     static_assert(is_rmem<EngineOut>::value, "Output tensor for A conversion must come from registers");
     static_assert(cosize_v<LayoutIn> == cosize_v<LayoutOut>);
@@ -805,14 +804,15 @@ public:
         {
           auto&& scale_neg_ = reinterpret_cast<cutlass::Array<uint32_t, 2> const&>(scales_neg_vm_(i));
           auto&& scale_pos_ = reinterpret_cast<cutlass::Array<uint32_t, 2>      &>(scales_pos_vm_(i));
+          constexpr uint32_t immLut = (0xf0 & 0xcc) ^ 0xaa;
           asm volatile(
               "{\n"
-              "  and  .b32 %0, %2, %4             ;\n" \
-              "  and  .b32 %1, %3, %5             ;\n" \
+              "  lop3 .b32 %0, %2, %4, %5, %6;\n" \
+              "  xor  .b32 %1, %3, %5;        \n" \
               "}\n"
               : "=r"(scale_pos_[0]), "=r"(scale_pos_[1])
-              : "r"(scale_neg_[0]), "r"(scale_neg_[1]), "n"(0x7F7F7F00), "n"(0x7F7F7F7F)
-              );
+              : "r"(scale_neg_[0]), "r"(scale_neg_[1]), "n"(0xFFFFFF00), "n"(0x80808080), "n"(immLut)
+            );
         }
       }
       CUTLASS_PRAGMA_UNROLL

include/cutlass/detail/helper_macros.hpp

@@ -57,6 +57,12 @@
 #define CUTLASS_DEVICE inline
 #endif
 
+#if ! defined(_MSC_VER)
+#define CUTLASS_LAMBDA_FUNC_INLINE __attribute__((always_inline))
+#else
+#define CUTLASS_LAMBDA_FUNC_INLINE [[msvc::forceinline]]
+#endif
+
 #define CUTLASS_HOST __host__
 #define CUTLASS_GLOBAL __global__ static
 
@@ -74,11 +80,11 @@ CUTLASS_HOST_DEVICE void __CUTLASS_UNUSED(T const &)
 
 #ifdef _MSC_VER
 // Provides support for alternative operators 'and', 'or', and 'not'
-#include <iso646.h>
+#include <ciso646>
 #endif // _MSC_VER
 
 #if !defined(__CUDACC_RTC__)
-#include <assert.h>
+#include <cassert>
 #endif
 
 #if defined(__CUDA_ARCH__)

include/cutlass/detail/mainloop_fusion_helper_scale_factor.hpp

@@ -0,0 +1,75 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2024 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.
+ *
+ **************************************************************************************************/
+
+/*! \file
+    \brief Mainloop Fusion configs specific for scale factors
+*/
+
+#pragma once
+
+#include <cute/util/type_traits.hpp> // cute::void_t
+
+namespace cutlass::detail {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+template <typename CollectiveMainloop, typename = void>
+struct ElementSFType {
+  using type = void;
+};
+
+template <typename CollectiveMainloop>
+struct ElementSFType<CollectiveMainloop, cute::void_t<typename CollectiveMainloop::ElementSF>> {
+  using type = typename CollectiveMainloop::ElementSF;
+};
+
+template <typename CollectiveMainloop, typename = void>
+struct LayoutSFAType {
+  using type = void;
+};
+
+template <typename CollectiveMainloop>
+struct LayoutSFAType<CollectiveMainloop, cute::void_t<typename CollectiveMainloop::LayoutSFA>> {
+  using type = typename CollectiveMainloop::LayoutSFA;
+};
+
+template <typename CollectiveMainloop, typename = void>
+struct LayoutSFBType {
+  using type = void;
+};
+
+template <typename CollectiveMainloop>
+struct LayoutSFBType<CollectiveMainloop, cute::void_t<typename CollectiveMainloop::LayoutSFB>> {
+  using type = typename CollectiveMainloop::LayoutSFB;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace cutlass::detail

include/cutlass/device_kernel.h

@@ -34,8 +34,11 @@
 
 #pragma once
 
+#include <cutlass/detail/helper_macros.hpp> // CUTLASS_HOST_DEVICE
+#include <cutlass/platform/platform.h> // uint64_t
+
 // __grid_constant__ was introduced in CUDA 11.7.
-#if ((__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 7)))
+#if ((__CUDACC_VER_MAJOR__ >= 12) || ((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 7))) && !CUTLASS_CLANG_CUDA
 #  define CUTLASS_GRID_CONSTANT_SUPPORTED
 #endif
 

include/cutlass/epilogue/collective/builders/sm90_builder.inl

@@ -422,7 +422,8 @@ struct CollectiveBuilder<
     Schedule,
     fusion::LinearCombination<ElementD,ElementCompute,ElementC_,ElementCompute,RoundStyle>,
     cute::enable_if_t<cute::is_same_v<Schedule, NoSmemWarpSpecialized> ||
-                      cute::is_same_v<Schedule, PtrArrayNoSmemWarpSpecialized> >> {
+                      cute::is_same_v<Schedule, PtrArrayNoSmemWarpSpecialized> ||
+                      cute::is_same_v<Schedule, PtrArrayNoSmemWarpSpecializedTransposed> >> {
 
   // Passing void C disables source load
   using ElementC = cute::conditional_t<cute::is_void_v<ElementC_>,

include/cutlass/epilogue/collective/default_epilogue_array.hpp

@@ -86,7 +86,7 @@ public:
   static const int kOutputAlignment = ThreadEpilogueOp::kCount;
   using AlignmentType = typename cute::uint_bit<sizeof_bits<ElementOutput>::value * kOutputAlignment>::type;
 
-  static_assert(cute::is_same_v<EpilogueSchedule, PtrArrayNoSmemWarpSpecialized> || cute::is_same_v<EpilogueSchedule, PtrArrayDefault>, "Incompatible epilogue schedule.");
+  static_assert(cute::is_same_v<EpilogueSchedule, PtrArrayNoSmemWarpSpecialized> || cute::is_same_v<EpilogueSchedule, PtrArrayDefault> || cute::is_same_v<EpilogueSchedule, PtrArrayNoSmemWarpSpecializedTransposed>, "Incompatible epilogue schedule.");
   static_assert(rank(InternalStrideC{}) == 3, "StrideCD must be rank-3: [M, N, L]");
   static_assert(rank(InternalStrideD{}) == 3, "StrideCD must be rank-3: [M, N, L]");
 
@@ -198,20 +198,30 @@ public:
       assert(0);
     }
 
-    InternalStrideC stride_c;
-    InternalStrideD stride_d;
-    if constexpr (!cute::is_same_v<InternalStrideC, StrideC>) {
-      // If grouped gemm
-      if (epilogue_op.is_source_needed()) {
-        stride_c = detail::get_epilogue_stride<EpilogueSchedule>(params.dC[l_coord]);
+    auto [stride_c, stride_d] = [&, l = l_coord]() {
+      if constexpr (!cute::is_same_v<InternalStrideC, StrideC>) {
+        // If grouped gemm
+        if (epilogue_op.is_source_needed()) {
+            return make_tuple(
+                detail::get_epilogue_stride<EpilogueSchedule>(params.dC[l]),
+                detail::get_epilogue_stride<EpilogueSchedule>(params.dD[l])
+            );
+        } 
+        else {
+          return make_tuple(
+              InternalStrideC{}, 
+              detail::get_epilogue_stride<EpilogueSchedule>(params.dD[l])
+          );
+        }
+      } 
+      else {
+        return make_tuple(
+            detail::get_epilogue_stride<EpilogueSchedule>(params.dC),
+            detail::get_epilogue_stride<EpilogueSchedule>(params.dD)
+        );
       }
-      stride_d = detail::get_epilogue_stride<EpilogueSchedule>(params.dD[l_coord]);
-    }
-    else {
-      stride_c = detail::get_epilogue_stride<EpilogueSchedule>(params.dC);
-      stride_d = detail::get_epilogue_stride<EpilogueSchedule>(params.dD);
-    }
-
+    }();
+    
     // Represent the full output tensor
     ElementC const* ptr_C_l = nullptr;
     if (epilogue_op.is_source_needed()) {

include/cutlass/epilogue/collective/detail.hpp

@@ -157,7 +157,8 @@ struct EmptyStorage {
 template<class EpilogueSchedule, class Stride>
 CUTLASS_HOST_DEVICE
 auto get_epilogue_stride(Stride stride){
-  if constexpr (cute::is_base_of_v<cutlass::gemm::EpilogueTransposed, EpilogueSchedule>) {
+  if constexpr (cute::is_base_of_v<cutlass::gemm::EpilogueTransposed, EpilogueSchedule>||
+                cute::is_base_of_v<cutlass::epilogue::PtrArrayNoSmemWarpSpecializedTransposed, EpilogueSchedule>) {
     return cute::make_stride(cute::get<1>(stride), cute::get<0>(stride), cute::get<2>(stride));
   }
   else {
@@ -464,7 +465,7 @@ public:
   tensormaps_fence_acquire([[maybe_unused]] cute::TmaDescriptor const* tensormap) { }
 };
 
-// SFINAE helpers for detecting beta/beta_ptr in EVT arguments.
+// SFINAE helpers for detecting beta/beta_ptr/beta_ptr_array in EVT arguments.
 template <class Arguments, class = void>
 struct has_beta {
   static constexpr bool value = false;
@@ -485,6 +486,16 @@ struct has_beta_ptr<Arguments, cute::void_t<decltype(Arguments{}.thread.beta_ptr
   static constexpr bool value = true;
 };
 
+template <class Arguments, class = void>
+struct has_beta_ptr_array {
+  static constexpr bool value = false;
+};
+
+template <class Arguments>
+struct has_beta_ptr_array<Arguments, cute::void_t<decltype(Arguments{}.thread.beta_ptr_array)>> {
+  static constexpr bool value = true;
+};
+
 } // namespace detail
 } // namespace collective
 } // namespace epilogue

include/cutlass/epilogue/collective/sm90_epilogue_array_tma_warpspecialized.hpp

@@ -328,7 +328,7 @@ public:
     }
 
     uint32_t transaction_bytes = TmaTransactionBytes;
-    typename Params::TMA_C tma_load_c = {};
+    typename Params::TMA_C tma_load_c{};
     if constexpr (is_source_supported) {
       ElementC const* ptr_C_first_batch = reinterpret_cast<ElementC const*>(args.ptr_C); 
       Tensor tensor_c = make_tensor(ptr_C_first_batch, make_layout(make_shape(init_M,init_N,init_L), append<3>(stride_c, _0{})));
@@ -409,7 +409,7 @@ public:
           implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_D>(cute::make_shape(M,N,L), InternalStrideD{});
         }
 
-        if constexpr (not cute::is_void_v<ElementC>) {
+        if constexpr (is_source_supported) {
           constexpr int tma_alignment_bits_C = cutlass::detail::get_input_alignment_bits<ElementC>();
           constexpr int min_tma_aligned_elements_C = tma_alignment_bits_C / cutlass::sizeof_bits<ElementC>::value;
           implementable = implementable && cutlass::detail::check_alignment<min_tma_aligned_elements_C>(cute::make_shape(M,N,L), InternalStrideC{});
@@ -432,13 +432,16 @@ public:
 
     bool beta_implementable = true;
 
-    if constexpr (cute::is_void_v<ElementC>) {
+    if (cute::is_void_v<ElementC> || args.ptr_C == nullptr) {
       if constexpr (detail::has_beta<Arguments>::value) {
         beta_implementable = args.thread.beta == 0.0;
       }
       if constexpr (detail::has_beta_ptr<Arguments>::value) {
         beta_implementable = beta_implementable && args.thread.beta_ptr == nullptr;
       }
+      if constexpr (detail::has_beta_ptr_array<Arguments>::value) {
+        beta_implementable = beta_implementable && args.thread.beta_ptr_array == nullptr;
+      }
     }
 
     if (!beta_implementable) {
@@ -775,7 +778,7 @@ public:
                       tRS_rC,
                       thread_idx
                     };
-    auto cst_callbacks = fusion_callbacks.get_consumer_store_callbacks<RefSrc>(cst_args);
+    auto cst_callbacks = fusion_callbacks.template get_consumer_store_callbacks<RefSrc>(cst_args);
     bool is_producer_load_needed = fusion_callbacks.is_producer_load_needed();
     bool is_C_load_needed = is_source_supported && fusion_callbacks.is_C_load_needed();
 
@@ -1017,7 +1020,7 @@ public:
     Tensor gmem_tensormap = make_tensor(params.tensormaps, desc_layout);                      // (SMs, NumInputTensors)
 
     if constexpr (IsLoad) {
-      if (not cute::is_void_v<ElementC>) {
+      if (is_source_supported) {
         constexpr int C_tensormap_index = NumEpilogueWarpGroups;
         Tensor pC_tensormap = make_tensor(params.tma_load_c.get_tma_descriptor(), Int<1>{}, Int<1>{});
         Tensor sC_tensormap = make_tensor(make_smem_ptr(&shared_tensormaps.smem_tensormap_C), Int<1>{}, Int<1>{});
@@ -1058,8 +1061,10 @@ public:
     // Replacing global_address for the next batch
     if constexpr (IsLoad) {
       if constexpr (is_source_supported) {
-        cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormaps.smem_tensormap_C,
-                                                        params.ptr_C[next_batch]);
+        if (params.ptr_C != nullptr) {
+          cute::tma_descriptor_replace_addr_in_shared_mem(shared_tensormaps.smem_tensormap_C,
+                                                          params.ptr_C[next_batch]);
+        }
       }
     }
     else if constexpr (is_destination_supported) {
@@ -1087,18 +1092,20 @@ public:
 
     if constexpr (IsLoad) {
       if constexpr (is_source_supported) {
-        ElementC const* ptr_C = nullptr;
-        Tensor tensor_c = make_tensor(ptr_C, make_layout(make_shape(M,N,Int<1>{}), params.dC[next_group]));
+        if (params.dC != nullptr) {
+          ElementC const* ptr_C = nullptr;
+          Tensor tensor_c = make_tensor(ptr_C, make_layout(make_shape(M,N,Int<1>{}), params.dC[next_group]));
 
-        cute::detail::fill_tma_gmem_shape_stride(params.tma_load_c, tensor_c, 
-                                                 prob_shape, prob_stride);
-        // Convert strides to byte strides
-        for (uint64_t& stride : prob_stride) {
-          stride = (stride * sizeof_bits_v<ElementC>) / 8;
+          cute::detail::fill_tma_gmem_shape_stride(params.tma_load_c, tensor_c, 
+                                                  prob_shape, prob_stride);
+          // Convert strides to byte strides
+          for (uint64_t& stride : prob_stride) {
+            stride = (stride * sizeof_bits_v<ElementC>) / 8;
+          }
+          cute::tma_descriptor_replace_dims_strides_in_shared_mem(shared_tensormaps.smem_tensormap_C,
+                                                                  prob_shape,
+                                                                  prob_stride);
         }
-        cute::tma_descriptor_replace_dims_strides_in_shared_mem(shared_tensormaps.smem_tensormap_C,
-                                                                prob_shape,
-                                                                prob_stride);
       }
     }
     else if constexpr (is_destination_supported) {
@@ -1166,7 +1173,7 @@ public:
   void
   tensormaps_fence_acquire(cute::TmaDescriptor const* tensormap) {
     if constexpr (IsLoad) {
-      if constexpr (not cute::is_void_v<ElementC>) {
+      if constexpr (is_source_supported) {
         cute::tma_descriptor_fence_acquire(tensormap);
       }
     } 

include/cutlass/epilogue/collective/sm90_epilogue_tma_warpspecialized.hpp

@@ -94,7 +94,7 @@ class CollectiveEpilogue<
     SmemLayoutAtomD_,
     CopyOpR2S_,
     CopyAtomC_,
-    CopyOpR2R_,
+    CopyOpR2R_
 > {
 public:
   //
@@ -136,6 +136,9 @@ private:
   static_assert(not cute::is_void_v<NonVoidElementD>, "SmemElementD is void");
   using NonVoidElementC = cute::conditional_t<not is_source_supported,NonVoidElementD,ElementC>; // prevents void ref breakages
 
+  using TmaElementD = cute::conditional_t<cute::is_same_v<NonVoidElementD, cutlass::complex<float>>, uint64_t, NonVoidElementD>;
+  using TmaElementC = cute::conditional_t<cute::is_same_v<NonVoidElementC, cutlass::complex<float>>, uint64_t, NonVoidElementC>;
+
   using SmemElementC = typename cutlass::detail::get_unpacked_element_type<NonVoidElementC>::type;
   using SmemElementD = typename cutlass::detail::get_unpacked_element_type<NonVoidElementD>::type;
 
@@ -239,14 +242,14 @@ public:
   struct Params {
     using TMA_C = decltype(make_tma_copy(
         CopyOpG2S{},
-        make_tensor(make_gmem_ptr(static_cast<NonVoidElementC const*>(nullptr)),
+        make_tensor(make_gmem_ptr<TmaElementC const>(nullptr),
             repeat_like(StrideC{}, int32_t(0)), StrideC{}),
         take<0,2>(SmemLayoutC{}),
         EpilogueTile{},
         _1{}));
     using TMA_D = decltype(make_tma_copy(
         CopyOpS2G{},
-        make_tensor(make_gmem_ptr(static_cast<NonVoidElementD const*>(nullptr)),
+        make_tensor(make_gmem_ptr<TmaElementD>(nullptr),
             repeat_like(StrideD{}, int32_t(0)), StrideD{}),
         take<0,2>(SmemLayoutD{}),
         EpilogueTile{},
@@ -273,9 +276,9 @@ public:
     auto [M, N, K, L] = problem_shape_MNKL;
 
     uint32_t transaction_bytes = TmaTransactionBytes;
-    typename Params::TMA_C tma_load_c = {};
+    typename Params::TMA_C tma_load_c{};
     if constexpr (is_source_supported) {
-      Tensor tensor_c = make_tensor(make_gmem_ptr(args.ptr_C), make_layout(make_shape(M,N,L), args.dC));
+      Tensor tensor_c = make_tensor(make_gmem_ptr<TmaElementC const>(args.ptr_C), make_layout(make_shape(M,N,L), args.dC));
       tma_load_c = make_tma_copy_C_sm90(
           CopyOpG2S{},
           tensor_c,
@@ -285,7 +288,7 @@ public:
 
     typename Params::TMA_D tma_store_d;
     if constexpr (is_destination_supported) {
-      Tensor tensor_d = make_tensor(make_gmem_ptr(args.ptr_D), make_layout(make_shape(M,N,L), args.dD));
+      Tensor tensor_d = make_tensor(make_gmem_ptr<TmaElementD>(args.ptr_D), make_layout(make_shape(M,N,L), args.dD));
       tma_store_d = make_tma_copy_C_sm90(
           CopyOpS2G{},
           tensor_d,
@@ -644,7 +647,18 @@ public:
     // Absolute coordinate tensors (dynamic)
     Tensor mD_crd = make_identity_tensor(make_shape(M,N));                                                     // (M,N)
     Tensor cD_mn = local_tile(mD_crd, take<0,2>(CtaTileMNK{}), make_coord(m_coord, n_coord));          // (CTA_M,CTA_N)
-    Tensor tRS_cD_mn = thread_r2s.partition_S(flat_divide(cD_mn, EpilogueTile{}));     // (R2S,R2S_M,R2S_N,EPI_M,EPI_N)
+    Tensor tRS_cD_mn = [&]() {
+      if constexpr (IsUseR2R) {
+        // (t)hread-partition for ConsumerStoreCallbacks. 
+        TiledCopy tiled_cst = make_tiled_copy_S(Copy_Atom<CopyOpR2S,SmemElementC>{}, tiled_copy_C_atom);
+        ThrCopy thread_cst = tiled_cst.get_slice(thread_idx);
+
+        return thread_cst.partition_S(flat_divide(cD_mn, EpilogueTile{}));             // (R2S,R2S_M,R2S_N,EPI_M,EPI_N)
+      }
+      else {
+        return thread_r2s.partition_S(flat_divide(cD_mn, EpilogueTile{}));             // (R2S,R2S_M,R2S_N,EPI_M,EPI_N)
+      }
+    }();
     // Relative coordinate tensors (static)
     Tensor cD = make_counting_tensor(cD_mn.layout());                                                  // (CTA_M,CTA_N)
     Tensor tRS_cD = make_counting_tensor(tRS_cD_mn.layout());                          // (R2S,R2S_M,R2S_N,EPI_M,EPI_N)

include/cutlass/epilogue/dispatch_policy.hpp

@@ -50,6 +50,7 @@ struct EpilogueSimtVectorized {};
 struct EpiloguePtrArraySimtVectorized {};
 struct NoSmemWarpSpecialized {};
 struct PtrArrayNoSmemWarpSpecialized {};
+struct PtrArrayNoSmemWarpSpecializedTransposed {};
 struct PtrArrayPlanarComplexNoSmemWarpSpecialized {};
 struct TmaWarpSpecialized {};
 struct TmaWarpSpecializedCooperative {};

include/cutlass/epilogue/fusion/operations.hpp

@@ -34,6 +34,7 @@
 #include <cutlass/numeric_conversion.h>
 #include <cutlass/layout/matrix.h>
 #include <cute/numeric/numeric_types.hpp>
+#include <cute/numeric/integral_constant.hpp> // cute::false_type
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
@@ -60,9 +61,12 @@ struct FusionOperation {
   static constexpr int AlignmentScalar = 0;
   static constexpr bool IsScaleFactorSupported = false;
   static constexpr bool IsPerRowScaleSupported = false;
+  static constexpr bool IsPerColScaleSupported = false;
+
   using ElementBias = void;
   static constexpr int AlignmentBias = 0;
   static constexpr bool IsPerRowBiasSupported = false;
+  static constexpr bool IsPerColBiasSupported = false;
   static constexpr bool IsDePerRowBiasSupported = false;
 
   using ActivationFn = void;
@@ -190,6 +194,24 @@ struct LinCombPerRowBiasEltAct
   static constexpr bool IsEltActSupported = true;
 };
 
+// D = activation(alpha * acc + beta * C + per-column bias)
+template<
+  template <class> class ActivationFn_,
+  class ElementOutput_,
+  class ElementCompute_,
+  class ElementBias_ = ElementOutput_,
+  class ElementSource_ = ElementOutput_,
+  class ElementScalar_ = ElementCompute_,
+  int AlignmentBias_ = 128 / cute::sizeof_bits_v<ElementBias_>,
+  FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
+>
+struct LinCombPerColBiasEltAct
+    : LinCombPerColBias<ElementOutput_, ElementCompute_,
+        ElementBias_, ElementSource_, ElementScalar_, AlignmentBias_, RoundStyle_> {
+  using ActivationFn = ActivationFn_<ElementCompute_>;
+  static constexpr bool IsEltActSupported = true;
+};
+
 // D = activation(alpha * acc + beta * C + per-row bias)
 // aux = alpha * acc + beta * C + per-row bias
 template<
@@ -214,6 +236,30 @@ struct LinCombPerRowBiasEltActAux
   static constexpr bool IsAuxOutSupported = true;
 };
 
+// D = activation(alpha * acc + beta * C + per-col bias)
+// aux = alpha * acc + beta * C + per-col bias
+template<
+  class GmemLayoutTagAux_,
+  template <class> class ActivationFn_,
+  class ElementOutput_,
+  class ElementCompute_,
+  class ElementAux_ = ElementOutput_,
+  class ElementBias_ = ElementOutput_,
+  class ElementSource_ = ElementOutput_,
+  class ElementScalar_ = ElementCompute_,
+  int AlignmentAux_ = 128 / cute::sizeof_bits_v<ElementAux_>,
+  int AlignmentBias_ = 128 / cute::sizeof_bits_v<ElementBias_>,
+  FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
+>
+struct LinCombPerColBiasEltActAux
+    : LinCombPerColBiasEltAct<ActivationFn_, ElementOutput_, ElementCompute_,
+        ElementBias_, ElementSource_, ElementScalar_, AlignmentBias_, RoundStyle_> {
+  using ElementAux = ElementAux_;
+  using GmemLayoutTagAux = GmemLayoutTagAux_;
+  static constexpr int AlignmentAux = AlignmentAux_;
+  static constexpr bool IsAuxOutSupported = true;
+};
+
 // D = activation(per-row alpha * acc + per-row beta * C + per-row bias)
 template<
   template <class> class ActivationFn_,
@@ -233,6 +279,46 @@ struct PerRowLinCombPerRowBiasEltAct
   static constexpr bool IsPerRowScaleSupported = true;
 };
 
+// D = per-column alpha * per-row alpha * acc + beta * C
+template<
+  class ElementOutput_,
+  class ElementCompute_,
+  class ElementSource_ = ElementCompute_,
+  class ElementScalar_ = ElementCompute_,
+  int AlignmentScalar_ = 128 / cute::sizeof_bits_v<ElementScalar_>,
+  FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
+>
+struct OuterProdLinComb : FusionOperation {
+  using ElementOutput = ElementOutput_;
+  using ElementCompute = ElementCompute_;
+  using ElementSource = ElementSource_;
+  using ElementScalar = ElementScalar_;
+  static constexpr int AlignmentScalar = AlignmentScalar_;
+  static constexpr auto RoundStyle = RoundStyle_;
+  static constexpr bool IsSourceSupported = true;
+  static constexpr bool IsPerRowScaleSupported = true;
+  static constexpr bool IsPerColScaleSupported = true;
+};
+
+// D = activation(per-col alpha * acc + per-col beta * C + per-column bias)
+template<
+  template <class> class ActivationFn_,
+  class ElementOutput_,
+  class ElementCompute_,
+  class ElementBias_ = ElementOutput_,
+  class ElementSource_ = ElementOutput_,
+  class ElementScalar_ = ElementCompute_, // per-row alpha/beta
+  int AlignmentBias_ = 128 / cute::sizeof_bits_v<ElementBias_>,
+  int AlignmentScalar_ = 128 / cute::sizeof_bits_v<ElementScalar_>,
+  FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
+>
+struct PerColLinCombPerColBiasEltAct
+    : LinCombPerColBiasEltAct<ActivationFn_, ElementOutput_, ElementCompute_,
+        ElementBias_, ElementSource_, ElementScalar_, AlignmentBias_, RoundStyle_> {
+  static constexpr int AlignmentScalar = AlignmentScalar_;
+  static constexpr bool IsPerColScaleSupported = true;
+};
+
 // Z = scale_a * scale_b * alpha * acc + beta * scale_c * C + per-row bias
 // if D is fp8 
 //   D = scale_d * activation(Z)
@@ -254,6 +340,27 @@ struct ScaledLinCombPerRowBiasEltAct
   static constexpr bool IsScaleFactorSupported = true;
 };
 
+// Z = scale_a * scale_b * alpha * acc + beta * scale_c * C + per-col bias
+// if D is fp8 
+//   D = scale_d * activation(Z)
+// else
+//   D = activation(Z)
+template<
+  template <class> class ActivationFn_,
+  class ElementOutput_,
+  class ElementCompute_,
+  class ElementBias_ = ElementOutput_,
+  class ElementSource_ = ElementOutput_,
+  class ElementScalar_ = ElementCompute_,
+  int AlignmentBias_ = 128 / cute::sizeof_bits_v<ElementBias_>,
+  FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
+>
+struct ScaledLinCombPerColBiasEltAct
+    : LinCombPerColBiasEltAct<ActivationFn_, ElementOutput_, ElementCompute_,
+        ElementBias_, ElementSource_, ElementScalar_, AlignmentBias_, RoundStyle_> {
+  static constexpr bool IsScaleFactorSupported = true;
+};
+
 // Z = scale_a * scale_b * alpha * acc + scale_c * beta * C + per-row bias
 // if D is fp8 
 //   amax_d = max(abs(elements in activation(Z)))
@@ -291,6 +398,43 @@ struct ScaledLinCombPerRowBiasEltActAmaxAux
   static constexpr bool IsAuxOutSupported = true;
 };
 
+// Z = scale_a * scale_b * alpha * acc + scale_c * beta * C + per-col bias
+// if D is fp8 
+//   amax_d = max(abs(elements in activation(Z)))
+//   D = scale_d * activation(Z)
+// else
+//   D = activation(Z)
+// if Aux is fp8 
+//   amax_aux = max(abs(elements in Z))
+//   Aux = scale_aux * Z
+// else
+//   Aux = Z
+template<
+  class GmemLayoutTagAux_,
+  template <class> class ActivationFn_,
+  class ElementOutput_,
+  class ElementCompute_,
+  class ElementAux_ = ElementOutput_,
+  class ElementAmax_ = ElementCompute_,
+  class ElementBias_ = ElementOutput_,
+  class ElementSource_ = ElementOutput_,
+  class ElementScalar_ = ElementCompute_,
+  int AlignmentAux_ = 128 / cute::sizeof_bits_v<ElementAux_>,
+  int AlignmentBias_ = 128 / cute::sizeof_bits_v<ElementBias_>,
+  FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
+>
+struct ScaledLinCombPerColBiasEltActAmaxAux
+    : ScaledLinCombPerColBiasEltAct<ActivationFn_, ElementOutput_, ElementCompute_,
+        ElementBias_, ElementSource_, ElementScalar_, AlignmentBias_, RoundStyle_> {
+  using ElementAmax = ElementAmax_;
+  static constexpr bool IsAbsMaxSupported = true;
+
+  using ElementAux = ElementAux_;
+  using GmemLayoutTagAux = GmemLayoutTagAux_;
+  static constexpr int AlignmentAux = AlignmentAux_;
+  static constexpr bool IsAuxOutSupported = true;
+};
+
 // Z = Aux
 // dY = alpha * acc + beta * C
 // D = d_activation(dY, Z)

include/cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp

@@ -708,6 +708,105 @@ struct FusionCallbacks<
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
+// D = activation(alpha * acc + beta * C + per-column bias)
+template<
+  int StagesC,
+  class CtaTileShapeMNK,
+  class EpilogueTile,
+  template <class> class ActivationFn,
+  class ElementOutput,
+  class ElementCompute,
+  class ElementBias = ElementOutput,
+  class ElementSource = ElementOutput,
+  class ElementScalar = ElementCompute,
+  int AlignmentBias = 128 / sizeof_bits_v<ElementBias>,
+  FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
+>
+using Sm90LinCombPerColBiasEltAct =
+  Sm90EVT<Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>,
+    Sm90LinCombPerColBias<StagesC, CtaTileShapeMNK, EpilogueTile, ElementCompute, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle>
+  >;
+
+template <
+  int StagesC,
+  int StagesD,
+  int FragmentSize,
+  bool ReuseSmemC,
+  bool DelayTmaStore,
+  template <class> class ActivationFn,
+  class ElementOutput,
+  class ElementCompute,
+  class ElementBias,
+  class ElementSource,
+  class ElementScalar,
+  int AlignmentBias,
+  FloatRoundStyle RoundStyle,
+  class CtaTileShapeMNK,
+  class EpilogueTile
+>
+struct FusionCallbacks<
+    epilogue::Sm90TmaWarpSpecialized<StagesC, StagesD, FragmentSize, ReuseSmemC, DelayTmaStore>,
+    fusion::LinCombPerColBiasEltAct<
+      ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle
+    >,
+    CtaTileShapeMNK,
+    EpilogueTile
+> : Sm90LinCombPerColBiasEltAct<
+      StagesC, CtaTileShapeMNK, EpilogueTile, ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle
+    > {
+
+  using Impl =
+    Sm90LinCombPerColBiasEltAct<
+      StagesC, CtaTileShapeMNK, EpilogueTile, ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle
+    >;
+  using Operation =
+    fusion::LinCombPerColBiasEltAct<
+      ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle
+    >;
+
+  struct Arguments {
+    ElementScalar alpha = ElementScalar(1);
+    ElementScalar beta = ElementScalar(0);
+    ElementScalar const* alpha_ptr = nullptr;
+    ElementScalar const* beta_ptr = nullptr;
+
+    using StrideAlpha = Stride<_0,_0,int64_t>;
+    using StrideBeta  = Stride<_0,_0,int64_t>;
+    StrideAlpha dAlpha = {_0{}, _0{}, 0};
+    StrideBeta  dBeta  = {_0{}, _0{}, 0};
+
+    using StrideBias = Stride<_0,_1,int64_t>;
+    ElementBias const* bias_ptr = nullptr;
+    StrideBias dBias = {};
+
+    using ActivationArguments = typename Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>::Arguments;
+    ActivationArguments activation = ActivationArguments();
+
+    operator typename Impl::Arguments() const {
+      return
+        {    // unary op : activation(beta * C + (alpha * acc + bias))
+          {    // ternary op : beta * C + (alpha * acc + bias)
+            {{beta}, {beta_ptr}, {dBeta}}, // leaf args : beta
+            {},                   // leaf args : C
+            {                     // ternary op : alpha * acc + bias
+              {{alpha}, {alpha_ptr}, {dAlpha}}, // leaf args : alpha
+              {},                     // leaf args : acc
+              {bias_ptr, ElementBias(0), dBias}, // leaf args : bias
+              {}                  // ternary args : multiply_add
+            },                    // end ternary op
+            {} // ternary args : multiply_add
+          },   // end ternary op
+          activation // unary args : activation
+        };   // end unary op
+    }
+  };
+
+  // Ctor inheritance
+  using Impl::Impl;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
 // D = activation(alpha * acc + beta * C + per-row bias)
 // Aux = alpha * acc + beta * C + per-row bias)
 template<
@@ -832,6 +931,132 @@ struct FusionCallbacks<
 };
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
+// D = activation(alpha * acc + beta * C + per_col bias)
+// Aux = alpha * acc + beta * C + per_col bias)
+template<
+  int StagesC,
+  class CtaTileShapeMNK,
+  class EpilogueTile,
+  int Stages,
+  class StrideAux,
+  class SmemLayoutAtom,
+  class CopyOpR2S,
+  template <class> class ActivationFn,
+  class ElementOutput,
+  class ElementCompute,
+  class ElementAux = ElementOutput,
+  class ElementBias = ElementOutput,
+  class ElementSource = ElementOutput,
+  class ElementScalar = ElementCompute,
+  int AlignmentAux = 128 / sizeof_bits_v<ElementAux>,
+  int AlignmentBias = 128 / sizeof_bits_v<ElementBias>,
+  FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
+>
+using Sm90LinCombPerColBiasEltActAux =
+  Sm90EVT<Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>,
+    Sm90EVT<Sm90AuxStore<Stages, EpilogueTile, ElementAux, RoundStyle, StrideAux, SmemLayoutAtom, CopyOpR2S, AlignmentAux>,
+      Sm90LinCombPerColBias<StagesC, CtaTileShapeMNK, EpilogueTile, ElementCompute, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle>
+    >
+  >;
+
+template <
+  int StagesC,
+  int StagesD,
+  int FragmentSize,
+  bool ReuseSmemC,
+  bool DelayTmaStore,
+  class GmemLayoutTagAux,
+  template <class> class ActivationFn,
+  class ElementOutput,
+  class ElementCompute,
+  class ElementAux,
+  class ElementBias,
+  class ElementSource,
+  class ElementScalar,
+  int AlignmentAux,
+  int AlignmentBias,
+  FloatRoundStyle RoundStyle,
+  class CtaTileShapeMNK,
+  class EpilogueTile,
+  class SmemLayoutAtom,
+  class CopyOpR2S
+>
+struct FusionCallbacks<
+    epilogue::Sm90TmaWarpSpecialized<StagesC, StagesD, FragmentSize, ReuseSmemC, DelayTmaStore>,
+    fusion::LinCombPerColBiasEltActAux<
+      GmemLayoutTagAux, ActivationFn, ElementOutput, ElementCompute,
+      ElementAux, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
+    >,
+    CtaTileShapeMNK,
+    EpilogueTile,
+    SmemLayoutAtom,
+    CopyOpR2S
+> : Sm90LinCombPerColBiasEltActAux<
+      StagesC, CtaTileShapeMNK, EpilogueTile, StagesD, cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>, SmemLayoutAtom, CopyOpR2S, ActivationFn,
+      ElementOutput, ElementCompute, ElementAux, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
+    > {
+
+  using Impl =
+    Sm90LinCombPerColBiasEltActAux<
+      StagesC, CtaTileShapeMNK, EpilogueTile, StagesD, cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>, SmemLayoutAtom, CopyOpR2S, ActivationFn,
+      ElementOutput, ElementCompute, ElementAux, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
+    >;
+  using Operation =
+    fusion::LinCombPerColBiasEltActAux<
+      GmemLayoutTagAux, ActivationFn,
+      ElementOutput, ElementCompute, ElementAux, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
+    >;
+
+  struct Arguments {
+    ElementScalar alpha = ElementScalar(1);
+    ElementScalar beta = ElementScalar(0);
+    ElementScalar const* alpha_ptr = nullptr;
+    ElementScalar const* beta_ptr = nullptr;
+
+    using StrideAlpha = Stride<_0,_0,int64_t>;
+    using StrideBeta  = Stride<_0,_0,int64_t>;
+    StrideAlpha dAlpha = {_0{}, _0{}, 0};
+    StrideBeta  dBeta  = {_0{}, _0{}, 0};
+
+    using StrideBias = Stride<_0,_1,int64_t>;
+    ElementBias const* bias_ptr = nullptr;
+    StrideBias dBias = {};
+
+    using ActivationArguments = typename Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>::Arguments;
+    ActivationArguments activation = ActivationArguments();
+
+    using StrideAux = cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>;
+    ElementAux* aux_ptr = nullptr;
+    StrideAux dAux = {};
+
+    operator typename Impl::Arguments() const {
+      return
+        {    // unary op : activation(store(beta * C + (alpha * acc + bias)))
+          {                 // unary op : store(beta * C + (alpha * acc + bias))
+            {                  // ternary op : beta * C + (alpha * acc + bias)
+              {{beta}, {beta_ptr}, {dBeta}}, // leaf args : beta
+              {},                   // leaf args : C
+              {                     // ternary op : alpha * acc + bias
+                {{alpha}, {alpha_ptr}, {dAlpha}}, // leaf args : alpha
+                {},                     // leaf args : acc
+                {bias_ptr, ElementBias(0), dBias}, // leaf args : bias
+                {}                  // ternary args : multiply_add
+              },                    // end ternary op
+              {}               // ternary args : multiply_add
+            },                 // end ternary op
+            {aux_ptr, dAux} // unary args : store
+          },                // end unary op
+          activation // unary args : activation
+        };   // end unary op
+    }
+  };
+
+  // Ctor inheritance
+  using Impl::Impl;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
 // D = per-row alpha * acc + per-row beta * C + per-row bias
 template<
   class CtaTileShapeMNK,
@@ -954,6 +1179,133 @@ struct FusionCallbacks<
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
+// D = per-col alpha * acc + per-col beta * C + per-column bias
+template<
+  int StagesC,
+  class CtaTileShapeMNK,
+  class EpilogueTile,
+  class ElementOutput,
+  class ElementCompute,
+  class ElementBias = ElementOutput,
+  class ElementSource = ElementOutput,
+  class ElementScalar = ElementCompute,
+  int AlignmentBias = 128 / sizeof_bits_v<ElementBias>,
+  int AlignmentScalar = 128 / sizeof_bits_v<ElementScalar>,
+  FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
+>
+using Sm90PerColLinCombPerColBias =
+  Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementOutput, ElementCompute, RoundStyle>, // beta * C + (alpha * acc + bias)
+    Sm90RowBroadcast<0, CtaTileShapeMNK, ElementScalar, ElementCompute, Stride<_0,bool,int64_t>, AlignmentScalar>, // beta, dynamic scalar/vector broadcast
+    Sm90SrcFetch<ElementSource>, // C
+    Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementCompute, ElementCompute, RoundStyle>, // alpha * acc + bias
+      Sm90RowBroadcast<0, CtaTileShapeMNK, ElementScalar, ElementCompute, Stride<_0,bool,int64_t>, AlignmentScalar>, // alpha, dynamic scalar/vector broadcast
+      Sm90AccFetch, // acc
+      Sm90RowBroadcast<0, CtaTileShapeMNK, ElementBias, ElementCompute, Stride<_0,_1,int64_t>, AlignmentBias> // bias
+    >
+  >;
+
+// D = activation(per-col alpha * acc + per-col beta * C + per-column bias)
+template<
+  int StagesC,
+  class CtaTileShapeMNK,
+  class EpilogueTile,
+  template <class> class ActivationFn,
+  class ElementOutput,
+  class ElementCompute,
+  class ElementBias = ElementOutput,
+  class ElementSource = ElementOutput,
+  class ElementScalar = ElementCompute,
+  int AlignmentBias = 128 / sizeof_bits_v<ElementBias>,
+  int AlignmentScalar = 128 / sizeof_bits_v<ElementScalar>,
+  FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
+>
+using Sm90PerColLinCombPerColBiasEltAct =
+  Sm90EVT<Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>,
+    Sm90PerColLinCombPerColBias<StagesC, CtaTileShapeMNK, EpilogueTile, ElementCompute, ElementCompute,
+                                ElementBias, ElementSource, ElementScalar, AlignmentBias, AlignmentScalar, RoundStyle>
+  >;
+
+template <
+  int StagesC,
+  int StagesD,
+  int FragmentSize,
+  bool ReuseSmemC,
+  bool DelayTmaStore,
+  template <class> class ActivationFn,
+  class ElementOutput,
+  class ElementCompute,
+  class ElementBias,
+  class ElementSource,
+  class ElementScalar,
+  int AlignmentBias,
+  int AlignmentScalar,
+  FloatRoundStyle RoundStyle,
+  class CtaTileShapeMNK,
+  class EpilogueTile
+>
+struct FusionCallbacks<
+    epilogue::Sm90TmaWarpSpecialized<StagesC, StagesD, FragmentSize, ReuseSmemC, DelayTmaStore>,
+    fusion::PerColLinCombPerColBiasEltAct<
+      ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, AlignmentScalar, RoundStyle
+    >,
+    CtaTileShapeMNK,
+    EpilogueTile
+> : Sm90PerColLinCombPerColBiasEltAct<
+      StagesC, CtaTileShapeMNK, EpilogueTile, ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, AlignmentScalar, RoundStyle
+    > {
+
+  using Impl =
+    Sm90PerColLinCombPerColBiasEltAct<
+      StagesC, CtaTileShapeMNK, EpilogueTile, ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, AlignmentScalar, RoundStyle
+    >;
+  using Operation =
+    fusion::PerColLinCombPerColBiasEltAct<
+      ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, AlignmentScalar, RoundStyle
+    >;
+
+  struct Arguments {
+    ElementScalar alpha = ElementScalar(1);
+    ElementScalar beta = ElementScalar(0);
+    ElementScalar const* alpha_ptr = nullptr;
+    ElementScalar const* beta_ptr = nullptr;
+
+    using StrideAlpha = Stride<_0,bool,int64_t>;
+    using StrideBeta  = Stride<_0,bool,int64_t>;
+    StrideAlpha dAlpha = {_0{}, bool(1), 0};
+    StrideBeta  dBeta  = {_0{}, bool(1), 0};
+
+    using StrideBias = Stride<_0,_1,int64_t>;
+    ElementBias const* bias_ptr = nullptr;
+    StrideBias dBias = {};
+
+    using ActivationArguments = typename Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>::Arguments;
+    ActivationArguments activation = ActivationArguments();
+
+    operator typename Impl::Arguments() const {
+      return
+        {    // unary op : activation(beta * C + (alpha * acc + bias))
+          {    // ternary op : beta * C + (alpha * acc + bias)
+            {beta_ptr, beta, dBeta}, // leaf args : beta
+            {},               // leaf args : C
+            {                 // ternary op : alpha * acc + bias
+              {alpha_ptr, alpha, dAlpha},   // leaf args : alpha
+              {},                   // leaf args : acc
+              {bias_ptr, ElementBias(0), dBias}, // leaf args : bias
+              {}              // ternary args : multiply_add
+            },                // end ternary op
+            {} // ternary args : multiply_add
+          },   // end ternary op
+          activation // unary args : activation
+        };   // end unary op
+    }
+  };
+
+  // Ctor inheritance
+  using Impl::Impl;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
 namespace detail {
 
 template <typename T>
@@ -1120,6 +1472,154 @@ struct FusionCallbacks<
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
+// D = scale_a * scale_b * alpha * acc + scale_c * beta * C + per-col bias
+template<
+  class CtaTileShapeMNK,
+  class ElementOutput,
+  class ElementCompute,
+  class ElementBias = ElementOutput,
+  class ElementSource = ElementOutput,
+  class ElementScalar = ElementCompute,
+  int AlignmentBias = 128 / sizeof_bits_v<ElementBias>,
+  FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
+>
+using Sm90ScaledLinCombPerColBias =
+  Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementOutput, ElementCompute, RoundStyle>, // beta * C + (alpha * acc + bias)
+    Sm90ScalarBroadcast<ElementScalar, Stride<_0,_0,int64_t>, 2>, // scale_c * beta
+    Sm90SrcFetch<ElementSource>, // C
+    Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementCompute, ElementCompute, RoundStyle>, // alpha * acc + bias
+      Sm90ScalarBroadcast<ElementScalar, Stride<_0,_0,int64_t>, 3>, // scale_a * scale_b * alpha
+      Sm90AccFetch, // acc
+      Sm90RowBroadcast<0, CtaTileShapeMNK, ElementBias, ElementCompute, Stride<_0,_1,int64_t>, AlignmentBias> // bias
+    >
+  >;
+
+// Z = scale_a * scale_b * alpha * acc + beta * scale_c * C + per-col bias
+// if D is fp8 
+//   D = scale_d * activation(Z)
+// else
+//   D = activation(Z)
+template<
+  class CtaTileShapeMNK,
+  template <class> class ActivationFn,
+  class ElementOutput,
+  class ElementCompute,
+  class ElementBias = ElementOutput,
+  class ElementSource = ElementOutput,
+  class ElementScalar = ElementCompute,
+  int AlignmentBias = 128 / sizeof_bits_v<ElementBias>,
+  FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
+>
+using Sm90ScaledLinCombPerColBiasEltAct =
+  Sm90EVT<Sm90Compute<detail::ScaleOutOp<ElementOutput>::template Op, ElementOutput, ElementCompute, RoundStyle>, // activation(Z) * scale_d
+    Sm90EVT<Sm90Compute<ActivationFn, ElementCompute, ElementCompute, RoundStyle>, // activation(Z)
+      // Z = scale_a * scale_b * alpha * acc + beta * scale_c * C + per-row bias
+      Sm90ScaledLinCombPerColBias<CtaTileShapeMNK, ElementCompute, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle>
+    >,
+    Sm90ScalarBroadcast<ElementScalar> // scale_d
+  >;
+
+template <
+  int StagesC,
+  int StagesD,
+  int FragmentSize,
+  bool ReuseSmemC,
+  bool DelayTmaStore,
+  template <class> class ActivationFn,
+  class ElementOutput,
+  class ElementCompute,
+  class ElementBias,
+  class ElementSource,
+  class ElementScalar,
+  int AlignmentBias,
+  FloatRoundStyle RoundStyle,
+  class CtaTileShapeMNK,
+  class EpilogueTile
+>
+struct FusionCallbacks<
+    epilogue::Sm90TmaWarpSpecialized<StagesC, StagesD, FragmentSize, ReuseSmemC, DelayTmaStore>,
+    fusion::ScaledLinCombPerColBiasEltAct<
+      ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle
+    >,
+    CtaTileShapeMNK,
+    EpilogueTile
+> : Sm90ScaledLinCombPerColBiasEltAct<
+      CtaTileShapeMNK, ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle
+    > {
+
+  using Impl =
+    Sm90ScaledLinCombPerColBiasEltAct<
+      CtaTileShapeMNK, ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle
+    >;
+  using Operation =
+    fusion::ScaledLinCombPerColBiasEltAct<
+      ActivationFn, ElementOutput, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle
+    >;
+
+  struct Arguments {
+    ElementScalar alpha = ElementScalar(1);
+    ElementScalar beta = ElementScalar(0);
+    ElementScalar const* alpha_ptr = nullptr;
+    ElementScalar const* beta_ptr = nullptr;
+
+    ElementScalar scale_a = ElementScalar(1);
+    ElementScalar scale_b = ElementScalar(1);
+    ElementScalar scale_c = ElementScalar(1);
+    ElementScalar scale_d = ElementScalar(1);
+    ElementScalar const* scale_a_ptr = nullptr;
+    ElementScalar const* scale_b_ptr = nullptr;
+    ElementScalar const* scale_c_ptr = nullptr;
+    ElementScalar const* scale_d_ptr = nullptr;
+
+    using StrideAlpha = Stride<_0,_0,int64_t>;
+    using StrideBeta  = Stride<_0,_0,int64_t>;
+    StrideAlpha dAlpha = {_0{}, _0{}, 0};
+    StrideBeta  dBeta  = {_0{}, _0{}, 0};
+
+    using StrideBias = Stride<_0,_1,int64_t>;
+    ElementBias const* bias_ptr = nullptr;
+    StrideBias dBias = {};
+
+    using ActivationArguments = typename Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>::Arguments;
+    ActivationArguments activation = ActivationArguments();
+
+    operator typename Impl::Arguments() const {
+      return
+        {    // binary op : activation((scale_c * beta) * C + ((scale_a * scale_b * alpha) * acc + bias)) * scale_d
+          {    // unary op : activation((scale_c * beta) * C + ((scale_a * scale_b * alpha) * acc + bias))
+            {    // ternary op : (scale_c * beta) * C + ((scale_a * scale_b * alpha) * acc + bias)
+              {{beta, scale_c},
+               {beta_ptr, scale_c_ptr},
+               {dBeta, {_0{}, _0{}, 0}}
+               },  // leaf args : (scale_c * beta)
+              {},  // leaf args : C
+              {    // ternary op : (scale_a * scale_b * alpha) * acc + bias
+                {{alpha, scale_a, scale_b}, 
+                 {alpha_ptr, scale_a_ptr, scale_b_ptr},
+                 {dAlpha, {_0{}, _0{}, 0}, {_0{}, _0{}, 0}}
+                 },                   // leaf args : (scale_a * scale_b * alpha)
+                {},                   // leaf args : acc
+                {bias_ptr, ElementBias(0), dBias}, // leaf args : bias
+                {} // ternary args : multiply_add
+              },   // end ternary op
+              {} // ternary args : multiply_add
+            },   // end ternary op
+            activation // unary args : activation
+          },   // end unary op
+          {{scale_d},
+           {scale_d_ptr}
+           },   // leaf args : scale_d
+          {} // binary args : multiplies or first
+        };   // end binary op
+    }
+  };
+
+  // Ctor inheritance
+  using Impl::Impl;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
 // Z = scale_a * scale_b * alpha * acc + scale_c * beta * C + per-row bias
 // if D is fp8 
 //   amax_d = max(abs(elements in activation(Z)))
@@ -1440,6 +1940,326 @@ struct FusionCallbacks<
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
+// Z = scale_a * scale_b * alpha * acc + scale_c * beta * C + per-col bias
+// if D is fp8 
+//   amax_d = max(abs(elements in activation(Z)))
+//   D = scale_d * activation(Z)
+// else
+//   D = activation(Z)
+// if Aux is fp8 
+//   amax_aux = max(abs(elements in Z))
+//   Aux = scale_aux * Z
+// else
+//   Aux = Z
+
+// fp8 aux specialization
+template<
+  class CtaTileShapeMNK,
+  class EpilogueTile,
+  int StagesD,
+  class StrideAux,
+  class SmemLayoutAtom,
+  class CopyOpR2S,
+  template <class> class ActivationFn,
+  class ElementOutput,
+  class ElementCompute,
+  class ElementAux = ElementOutput,
+  class ElementAmax = ElementCompute,
+  class ElementBias = ElementOutput,
+  class ElementSource = ElementOutput,
+  class ElementScalar = ElementCompute,
+  int AlignmentAux = 128 / sizeof_bits_v<ElementAux>,
+  int AlignmentBias = 128 / sizeof_bits_v<ElementBias>,
+  FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
+>
+using Sm90ScaledLinCombPerColBiasEltActAmaxAuxFp8 =
+  Sm90SplitTreeVisitor<
+    // Z = scale_a * scale_b * alpha * acc + scale_c * beta * C + per-col bias
+    Sm90ScaledLinCombPerColBias<CtaTileShapeMNK, ElementCompute, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle>,
+    // D = activation(Z) * scale_d, amax_d = max(abs(elements in D))
+    Sm90EVT<Sm90Compute<detail::ScaleOutOp<ElementOutput>::template Op, ElementOutput, ElementCompute, RoundStyle>, // activation(Z) * scale_d
+      Sm90EVT<Sm90ScalarReduction<detail::amax, atomic_maximum, ElementAmax, ElementCompute, RoundStyle>, // amax_d
+        Sm90EVT<Sm90Compute<ActivationFn, ElementCompute, ElementCompute, RoundStyle>, // activation(Z)
+          Sm90SplitTreeFetch // Z
+        >
+      >,
+      Sm90ScalarBroadcast<ElementScalar> // scale_d
+    >,
+    // Aux = Z * scale_aux, amax_aux = max(abs(elements in Aux))
+    Sm90EVT<Sm90AuxStore<StagesD, EpilogueTile, ElementAux, RoundStyle, StrideAux, SmemLayoutAtom, CopyOpR2S, AlignmentAux>, // store(Aux)
+      Sm90EVT<Sm90Compute<cutlass::multiplies, ElementCompute, ElementCompute, RoundStyle>, // Z * scale_aux
+        Sm90EVT<Sm90ScalarReduction<detail::amax, atomic_maximum, ElementAmax, ElementCompute, RoundStyle>, // amax_aux
+          Sm90SplitTreeFetch // Z
+        >,
+        Sm90ScalarBroadcast<ElementScalar> // scale_aux
+      >
+    >
+  >;
+
+// non-fp8 aux specialization
+// lets us use some EVT specializations such as relu + uint1b_t aux
+template<
+  class CtaTileShapeMNK,
+  class EpilogueTile,
+  int StagesD,
+  class StrideAux,
+  class SmemLayoutAtom,
+  class CopyOpR2S,
+  template <class> class ActivationFn,
+  class ElementOutput,
+  class ElementCompute,
+  class ElementAux = ElementOutput,
+  class ElementAmax = ElementCompute,
+  class ElementBias = ElementOutput,
+  class ElementSource = ElementOutput,
+  class ElementScalar = ElementCompute,
+  int AlignmentAux = 128 / sizeof_bits_v<ElementAux>,
+  int AlignmentBias = 128 / sizeof_bits_v<ElementBias>,
+  FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
+>
+using Sm90ScaledLinCombPerColBiasEltActAmaxAuxNotFp8 =
+  // D = activation(Z) * scale_d, amax_d = max(abs(elements in D))
+  Sm90EVT<Sm90Compute<detail::ScaleOutOp<ElementOutput>::template Op, ElementOutput, ElementCompute, RoundStyle>, // activation(Z) * scale_d
+    Sm90EVT<Sm90ScalarReduction<detail::amax, atomic_maximum, ElementAmax, ElementCompute, RoundStyle>, // amax_d
+      Sm90EVT<Sm90Compute<ActivationFn, ElementCompute, ElementCompute, RoundStyle>, // activation(Z)
+        Sm90EVT<Sm90AuxStore<StagesD, EpilogueTile, ElementAux, RoundStyle, StrideAux, SmemLayoutAtom, CopyOpR2S, AlignmentAux>, // Aux = Z
+          // Z = scale_a * scale_b * alpha * acc + scale_c * beta * C + per-row bias
+          Sm90ScaledLinCombPerColBias<CtaTileShapeMNK, ElementCompute, ElementCompute, ElementBias, ElementSource, ElementScalar, AlignmentBias, RoundStyle>
+        >
+      >
+    >,
+    Sm90ScalarBroadcast<ElementScalar> // scale_d
+  >;
+
+// dispatcher
+template<
+  class CtaTileShapeMNK,
+  class EpilogueTile,
+  int StagesD,
+  class StrideAux,
+  class SmemLayoutAtom,
+  class CopyOpR2S,
+  template <class> class ActivationFn,
+  class ElementOutput,
+  class ElementCompute,
+  class ElementAux = ElementOutput,
+  class ElementAmax = ElementCompute,
+  class ElementBias = ElementOutput,
+  class ElementSource = ElementOutput,
+  class ElementScalar = ElementCompute,
+  int AlignmentAux = 128 / sizeof_bits_v<ElementAux>,
+  int AlignmentBias = 128 / sizeof_bits_v<ElementBias>,
+  FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
+>
+using Sm90ScaledLinCombPerColBiasEltActAmaxAux = conditional_t<detail::is_fp8_v<ElementAux>,
+  Sm90ScaledLinCombPerColBiasEltActAmaxAuxFp8<
+    CtaTileShapeMNK, EpilogueTile, StagesD, StrideAux, SmemLayoutAtom, CopyOpR2S, ActivationFn,
+    ElementOutput, ElementCompute, ElementAux, ElementAmax, ElementBias, ElementSource, ElementScalar,AlignmentAux, AlignmentBias, RoundStyle
+  >,
+  Sm90ScaledLinCombPerColBiasEltActAmaxAuxNotFp8<
+    CtaTileShapeMNK, EpilogueTile, StagesD, StrideAux, SmemLayoutAtom, CopyOpR2S, ActivationFn,
+    ElementOutput, ElementCompute, ElementAux, ElementAmax, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
+  >
+>;
+
+
+template <
+  int StagesC,
+  int StagesD,
+  int FragmentSize,
+  bool ReuseSmemC,
+  bool DelayTmaStore,
+  class GmemLayoutTagAux,
+  template <class> class ActivationFn,
+  class ElementOutput,
+  class ElementCompute,
+  class ElementAux,
+  class ElementAmax,
+  class ElementBias,
+  class ElementSource,
+  class ElementScalar,
+  int AlignmentAux,
+  int AlignmentBias,
+  FloatRoundStyle RoundStyle,
+  class CtaTileShapeMNK,
+  class EpilogueTile,
+  class SmemLayoutAtom,
+  class CopyOpR2S
+>
+struct FusionCallbacks<
+    epilogue::Sm90TmaWarpSpecialized<StagesC, StagesD, FragmentSize, ReuseSmemC, DelayTmaStore>,
+    fusion::ScaledLinCombPerColBiasEltActAmaxAux<
+      GmemLayoutTagAux, ActivationFn, ElementOutput, ElementCompute,
+      ElementAux, ElementAmax, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
+    >,
+    CtaTileShapeMNK,
+    EpilogueTile,
+    SmemLayoutAtom,
+    CopyOpR2S
+> : Sm90ScaledLinCombPerColBiasEltActAmaxAux<
+      CtaTileShapeMNK, EpilogueTile, StagesD, cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>,
+      SmemLayoutAtom, CopyOpR2S, ActivationFn,
+      ElementOutput, ElementCompute, ElementAux, ElementAmax, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
+    > {
+
+  using Impl =
+    Sm90ScaledLinCombPerColBiasEltActAmaxAux<
+      CtaTileShapeMNK, EpilogueTile, StagesD, cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>,
+      SmemLayoutAtom, CopyOpR2S, ActivationFn,
+      ElementOutput, ElementCompute, ElementAux, ElementAmax, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
+    >;
+  using Operation =
+    fusion::ScaledLinCombPerColBiasEltActAmaxAux<
+      GmemLayoutTagAux, ActivationFn, ElementOutput, ElementCompute,
+      ElementAux, ElementAmax, ElementBias, ElementSource, ElementScalar, AlignmentAux, AlignmentBias, RoundStyle
+    >;
+
+  struct Arguments {
+    ElementScalar alpha = ElementScalar(1);
+    ElementScalar beta = ElementScalar(0);
+    ElementScalar const* alpha_ptr = nullptr;
+    ElementScalar const* beta_ptr = nullptr;
+
+    ElementScalar scale_a = ElementScalar(1);
+    ElementScalar scale_b = ElementScalar(1);
+    ElementScalar scale_c = ElementScalar(1);
+    ElementScalar scale_d = ElementScalar(1);
+    ElementScalar const* scale_a_ptr = nullptr;
+    ElementScalar const* scale_b_ptr = nullptr;
+    ElementScalar const* scale_c_ptr = nullptr;
+    ElementScalar const* scale_d_ptr = nullptr;
+
+    ElementScalar scale_aux = ElementScalar(1);
+    ElementScalar const* scale_aux_ptr = nullptr;
+
+    using StrideAlpha = Stride<_0,_0,int64_t>;
+    using StrideBeta  = Stride<_0,_0,int64_t>;
+    StrideAlpha dAlpha = {_0{}, _0{}, 0};
+    StrideBeta  dBeta  = {_0{}, _0{}, 0};
+
+    using StrideBias = Stride<_0,_1,int64_t>;
+    ElementBias const* bias_ptr = nullptr;
+    StrideBias dBias = {};
+
+    using ActivationArguments = typename Sm90Compute<ActivationFn, ElementOutput, ElementCompute, RoundStyle>::Arguments;
+    ActivationArguments activation = ActivationArguments();
+
+    ElementAmax* amax_D_ptr = nullptr;
+    ElementAmax* amax_aux_ptr = nullptr;
+
+    using StrideAux = cutlass::gemm::TagToStrideC_t<GmemLayoutTagAux>;
+    ElementAux* aux_ptr = nullptr;
+    StrideAux dAux = {};
+
+    operator typename Impl::Arguments() const {
+      // Only compute amax_d if D is fp8
+      ElementAmax* amax_D_ptr_ = nullptr;
+      if constexpr (detail::is_fp8_v<ElementOutput>) {
+        amax_D_ptr_ = amax_D_ptr;
+      }
+
+      // Aux is fp8 -> DAG arguments
+      if constexpr (detail::is_fp8_v<ElementAux>) {
+        typename Impl::Arguments args;
+        // always use structured binding to unpack DAG args since it may or may not be a tuple
+        auto& [Z_args, aux_args, D_args] = args;
+
+        Z_args =
+          {    // ternary op : (scale_c * beta) * C + ((scale_a * scale_b * alpha) * acc + bias)
+            {{beta, scale_c},
+             {beta_ptr, scale_c_ptr},
+             {dBeta, {_0{}, _0{}, 0}}
+             },  // leaf args : (scale_c * beta)
+            {},  // leaf args : C
+            {    // ternary op : (scale_a * scale_b * alpha) * acc + bias
+              {{alpha, scale_a, scale_b}, 
+               {alpha_ptr, scale_a_ptr, scale_b_ptr},
+               {dAlpha, {_0{}, _0{}, 0}, {_0{}, _0{}, 0}}
+               },                   // leaf args : (scale_a * scale_b * alpha)
+              {},                   // leaf args : acc
+              {bias_ptr, ElementBias(0), dBias}, // leaf args : bias
+              {} // ternary args : multiply_add
+            },   // end ternary op
+            {} // ternary args : multiply_add
+          };   // end ternary op
+
+        D_args =
+          {    // binary op : activation(Z) * scale_d or activation(Z)
+            {    // unary op : reduce(activation(Z))
+              {             // unary op : activation(Z)
+                {},             // leaf args : Z
+                activation      // unary args : activation
+              },                // end unary op
+              {amax_D_ptr_} // unary args : reduce
+            },              // end unary op
+            {{scale_d},
+             {scale_d_ptr}
+             },  // leaf args : scale_d
+            {} // binary args : multiplies or first
+          };   // end binary op
+
+        aux_args =
+          {    // unary op : store(Aux)
+            {    // binary op : Z * scale_d or Z
+              {    // unary op : reduce(Z)
+                {},            // leaf args : Z
+                {amax_aux_ptr} // unary args : reduce
+              },   // end unary op
+              {{scale_aux},
+               {scale_aux_ptr}
+               },  // leaf args : scale_d
+              {} // binary args : multiplies
+            },   // end binary op
+            {aux_ptr, dAux} // unary args : store
+          };   // end unary op
+
+        return args;
+      }
+
+      // Aux is not fp8 -> Tree arguments
+      else {
+        return
+          {  // binary op : activation(Z) * scale_d or activation(Z)
+            {  // unary op : reduce(activation(Z))
+              {  // unary op : activation(Z)
+                {  // unary op : store(Z)
+                  {  // ternary op : (scale_c * beta) * C + ((scale_a * scale_b * alpha) * acc + bias)
+                    {{beta, scale_c},
+                    {beta_ptr, scale_c_ptr},
+                    {dBeta, {_0{}, _0{}, 0}}
+                    },  // leaf args : (scale_c * beta)
+                    {},               // leaf args : C
+                    {                 // ternary op : (scale_a * scale_b * alpha) * acc + bias
+                      {{alpha, scale_a, scale_b}, 
+                       {alpha_ptr, scale_a_ptr, scale_b_ptr},
+                       {dAlpha, {_0{}, _0{}, 0}, {_0{}, _0{}, 0}}
+                      },                // leaf args : (scale_a * scale_b * alpha)
+                      {},               // leaf args : acc
+                      {bias_ptr, ElementBias(0), dBias
+                      },                // leaf args : bias
+                      {}              // ternary args : multiply_add
+                    },                // end ternary op
+                    {}              // ternary args : multiply_add
+                  },                // end ternary op
+                  {aux_ptr, dAux} // unary args : store
+                },                // end unary op
+                activation     // unary args : activation
+              },               // end unary op
+              {amax_D_ptr_} // unary args : reduce
+            },              // end unary op
+            {{scale_d},{scale_d_ptr}}, // leaf args : scale_d
+            {} // binary args : multiplies or first
+          };   // end binary op
+      }
+    }
+  };
+
+  // Ctor inheritance
+  using Impl::Impl;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
 template<
   class CtaTileShapeMNK,
   class EpilogueTile,
@@ -1679,6 +2499,87 @@ struct FusionCallbacks<
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
+// D = per-column alpha * per-row alpha * acc + beta * c
+template<
+  class CtaTileShapeMNK,
+  class ElementOutput,
+  class ElementCompute,
+  class ElementSource = ElementOutput,
+  class ElementScalar = ElementCompute,
+  int AlignmentScalar = 128 / sizeof_bits_v<ElementScalar>, // Alignment of per-column and per-row scaling vectors
+  FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
+>
+using Sm90OuterProdLinComb =
+  Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementOutput, ElementCompute, RoundStyle>, // c(beta) * c(C) + c(alpha * acc)
+    Sm90ScalarBroadcast<ElementScalar, Stride<_0,_0,int>>, // beta
+    Sm90SrcFetch<ElementSource>, // C
+    Sm90EVT<Sm90Compute<multiplies, ElementCompute, ElementCompute, RoundStyle>, // c(alpha) * c(acc)
+      Sm90OuterProduct<0, CtaTileShapeMNK, ElementScalar, Stride<_1,_0,int>, Stride<_0,_1,int>, AlignmentScalar>, // alpha_col * alpha_row
+      Sm90AccFetch // acc
+    >
+  >;
+
+template <
+  int StagesC,
+  int StagesD,
+  int FragmentSize,
+  bool ReuseSmemC,
+  bool DelayTmaStore,
+  class ElementOutput,
+  class ElementCompute,
+  class ElementSource,
+  class ElementScalar,
+  int AlignmentScalar,
+  FloatRoundStyle RoundStyle,
+  class CtaTileShapeMNK,
+  class EpilogueTile
+>
+struct FusionCallbacks<
+    epilogue::Sm90TmaWarpSpecialized<StagesC, StagesD, FragmentSize, ReuseSmemC, DelayTmaStore>,
+    OuterProdLinComb<ElementOutput, ElementCompute, ElementSource, ElementScalar, AlignmentScalar, RoundStyle>,
+    CtaTileShapeMNK,
+    EpilogueTile
+> : Sm90OuterProdLinComb<CtaTileShapeMNK, ElementOutput, ElementCompute, ElementSource, ElementScalar, AlignmentScalar, RoundStyle> {
+  using Impl = Sm90OuterProdLinComb<CtaTileShapeMNK, ElementOutput, ElementCompute, ElementSource, ElementScalar, AlignmentScalar, RoundStyle>;
+  using Operation = OuterProdLinComb<ElementOutput, ElementCompute, ElementSource, ElementScalar, AlignmentScalar, RoundStyle>;
+
+  struct Arguments {
+
+    // Give a name and flat ordering to the fusion callback args
+    using StrideCol  = Stride<_1,_0,int>;
+    using StrideRow  = Stride<_0,_1,int>;
+    using StrideBeta = Stride<_0,_0,int>;
+    ElementScalar const* alpha_ptr_col = nullptr;
+    ElementScalar const* alpha_ptr_row = nullptr;
+    ElementScalar        beta = static_cast<ElementScalar>(0);
+    ElementScalar const* beta_ptr = nullptr;
+    StrideCol  dAlphaCol = {};
+    StrideRow  dAlphaRow = {};
+    StrideBeta dBeta     = {};
+
+    // Conversion to the args expected by the visitor implementation
+    // to_underlying_arguments will implicitly call this
+    operator typename Impl::Arguments() const {
+      return
+        {
+          {beta, beta_ptr, dBeta}, // leaf args : beta
+          {},                      // leaf args : C
+          {
+            { alpha_ptr_col, alpha_ptr_row, dAlphaCol, dAlphaRow }, // leaf args : alpha cols / rows
+            {},                                                     // leaf args : acc
+            {}
+          },
+          {}
+        };
+    }
+  };
+
+  // Ctor inheritance
+  using Impl::Impl;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
 // D = softmax(top_k(alpha * acc + beta * C))
 template<
   int TopK,

include/cutlass/epilogue/fusion/sm90_visitor_compute_tma_warpspecialized.hpp

@@ -266,8 +266,8 @@ struct Sm90TreeVisitor<
     auto const& scale_op = get<0>(Impl::ops);
     auto const& added_op = get<2>(Impl::ops);
     if constexpr (detail::IsScalarBroadcast<InputScaleOp>::value && not is_void_v<ElementSource>) {
-      return (get<2>(scale_op.params_ptr->dScalar[0]) != 0 && scale_op.params_ptr->scalar_ptrs[0] != nullptr) || 
-              is_C_load_needed() || 
+      return (get<2>(scale_op.params_ptr->dScalar[0]) != 0 && scale_op.params_ptr->scalar_ptrs[0] != nullptr) ||
+              is_C_load_needed() ||
               added_op.is_producer_load_needed();
     }
     else {
@@ -408,8 +408,9 @@ template <
 >
 struct Sm90TreeVisitor<
   Sm90Compute<Activation, ElementOutput, ElementCompute, RoundStyle,
-              cute::enable_if_t<cute::is_same_v<Activation<ElementCompute>, cutlass::epilogue::thread::ReLu<ElementCompute>> ||
-                                cute::is_same_v<Activation<ElementCompute>, cutlass::epilogue::thread::Clamp<ElementCompute>>  >>,
+              cute::enable_if_t<cute::is_same_v<Activation<ElementCompute>, cutlass::epilogue::thread::ReLu<ElementCompute>>  ||
+                                cute::is_same_v<Activation<ElementCompute>, cutlass::epilogue::thread::Clamp<ElementCompute>> ||
+                                cute::is_same_v<Activation<ElementCompute>, cutlass::epilogue::thread::ThresholdReLU<ElementCompute>> >>,
   Sm90TreeVisitor<
     Sm90AuxStore<
       Stages,
@@ -503,7 +504,8 @@ struct Sm90TreeVisitor<
       CUTLASS_PRAGMA_UNROLL
       for (int i = 0; i < FragmentSize; ++i) {
         ElementCompute pre_relu = frg_compute[i];
-        if constexpr (cute::is_same_v<Activation<ElementCompute>, cutlass::epilogue::thread::Clamp<ElementCompute>>) {
+        if constexpr (cute::is_same_v<Activation<ElementCompute>, cutlass::epilogue::thread::Clamp<ElementCompute>> ||
+                      cute::is_same_v<Activation<ElementCompute>, cutlass::epilogue::thread::ThresholdReLU<ElementCompute>>) {
           frg_compute[i] = relu(frg_compute[i], params_compute);
         }
         else {

include/cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp

@@ -734,11 +734,12 @@ private:
 // Supports reduction over multiple broadcasts to support fusions such as fp8 scaling factors
 template<
   class Element,
-  class StrideMNL = Stride<_0,_0,_0>,
+  class StrideMNL_ = Stride<_0,_0,_0>,
   int BroadcastCount = 1,
   template <class> class ReductionFn = multiplies
 >
 struct Sm90ScalarBroadcastPtrArray {
+  using StrideMNL = StrideMNL_;
   static_assert(is_static_v<decltype(take<0,2>(StrideMNL{}))>); // batch stride can be dynamic or static
   static_assert(take<0,2>(StrideMNL{}) == Stride<_0,_0>{});
 
@@ -780,8 +781,8 @@ struct Sm90ScalarBroadcastPtrArray {
 
   CUTLASS_DEVICE bool
   is_producer_load_needed() const {
-    // producer load is needed if Element is not void and we have multiple scalars
-    return !cute::is_void_v<Element> and size<2>(params_ptr->dScalar[0]) != 0;
+    // producer load is needed if Element is not void
+    return !cute::is_void_v<Element>;
   }
 
   CUTLASS_DEVICE bool
@@ -814,7 +815,7 @@ struct Sm90ScalarBroadcastPtrArray {
   CUTLASS_DEVICE auto
   get_producer_load_callbacks(ProducerLoadArgs<Args...> const& args) {
     // Get the scalar for batched broadcast
-    if (get<2>(params_ptr->dScalar[0]) != 0) {
+    if (size<2>(params_ptr->dScalar[0]) != 0) {
       auto [m_coord, n_coord, k_coord, l_coord] = args.tile_coord_mnkl;
       update_scalar(l_coord);
     }
@@ -1377,6 +1378,171 @@ struct Sm90ColBroadcast {
   }
 };
 
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Do outer product from the column and row loaded
+//
+template<
+  int Stages,
+  class CtaTileShapeMNK,
+  class ElementScalar,
+  class StrideColMNL_ = Stride<_1,_0,int64_t>, /// NOTE: Batched scaling untested for now
+  class StrideRowMNL_ = Stride<_0,_1,int64_t>,
+  int Alignment = 128 / sizeof_bits_v<ElementScalar>,
+  bool EnableNullptr = false // Fallback scalar broadcast for nullptr params
+>
+struct Sm90OuterProduct {
+  using StrideColMNL = StrideColMNL_;
+  using StrideRowMNL = StrideRowMNL_;
+  static_assert(Stages == 0, "OuterProduct doesn't support smem usage");
+  static_assert(Alignment * sizeof_bits_v<ElementScalar> % 128 == 0, "sub-16B alignment not supported yet");
+  static_assert(!EnableNullptr, "Nullptr fallback not implemented");
+  static_assert(is_static_v<decltype(take<0,2>(StrideColMNL{}))> &&
+                is_static_v<decltype(take<0,2>(StrideRowMNL{}))>, "Only batch stride can be dynamic");
+  static_assert(take<0,2>(StrideColMNL{}) == Stride<_1,_0>{} &&
+                take<0,2>(StrideRowMNL{}) == Stride<_0,_1>{}, "Row and column incorrectly formatted");
+
+  // Accumulator distributes col/row elements evenly amongst threads so we can just directly load from gmem
+  struct SharedStorage { };
+
+  struct Arguments {
+    ElementScalar const* ptr_col = nullptr;
+    ElementScalar const* ptr_row = nullptr;
+    StrideColMNL dCol = {};
+    StrideRowMNL dRow = {};
+  };
+
+  using Params = Arguments;
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
+    return args;
+  }
+
+  template <class ProblemShape>
+  static bool
+  can_implement(ProblemShape const& problem_shape, Arguments const& args) {
+    return true;
+  }
+
+  template <class ProblemShape>
+  static size_t
+  get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
+    return 0;
+  }
+
+  template <class ProblemShape>
+  static cutlass::Status
+  initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream,
+    CudaHostAdapter* cuda_adapter = nullptr) {
+    return cutlass::Status::kSuccess;
+  }
+
+  CUTLASS_DEVICE bool
+  is_producer_load_needed() const {
+    return false;
+  }
+
+  CUTLASS_DEVICE bool
+  is_C_load_needed() const {
+    return false;
+  }
+
+  CUTLASS_DEVICE bool
+  is_zero() const {
+    return false;
+  }
+
+  CUTLASS_HOST_DEVICE
+  Sm90OuterProduct() { }
+
+  CUTLASS_HOST_DEVICE
+  Sm90OuterProduct(Params const& params, SharedStorage const& shared_storage)
+  : params(params) { }
+
+  Params params;
+
+  template <class... Args>
+  CUTLASS_DEVICE auto
+  get_producer_load_callbacks(ProducerLoadArgs<Args...> const& args) {
+    return EmptyProducerLoadCallbacks{};
+  }
+
+  template<
+    class GTensorCol, class RTensorCol,
+    class GTensorRow, class RTensorRow
+  >
+  struct ConsumerStoreCallbacks : EmptyConsumerStoreCallbacks {
+    CUTLASS_DEVICE
+    ConsumerStoreCallbacks(GTensorCol&& tCgCol, RTensorCol&& tCrCol,
+                           GTensorRow&& tCgRow, RTensorRow&& tCrRow,
+                           Params const& params)
+      : tCgCol(cute::forward<GTensorCol>(tCgCol))
+      , tCrCol(cute::forward<RTensorCol>(tCrCol))
+      , tCgRow(cute::forward<GTensorRow>(tCgRow))
+      , tCrRow(cute::forward<RTensorRow>(tCrRow))
+      , params(params) {}
+
+    GTensorCol tCgCol;                                                                 // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
+    RTensorCol tCrCol;                                                                 // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
+    GTensorRow tCgRow;                                                                 // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
+    RTensorRow tCrRow;                                                                 // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
+    Params const& params;
+
+    CUTLASS_DEVICE void
+    begin() {
+
+      // Filter so we don't issue redundant copies over stride-0 modes
+      copy(filter(tCgCol), filter(tCrCol));
+      copy(filter(tCgRow), filter(tCrRow));
+    }
+
+    template <typename ElementAccumulator, int FragmentSize>
+    CUTLASS_DEVICE Array<ElementScalar, FragmentSize>
+    visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) {
+      Array<ElementScalar, FragmentSize> frg_colrow;
+      Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n);
+      Tensor tCrRow_mn = tCrRow(_,_,_,epi_m,epi_n);
+
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < FragmentSize; ++i) {
+        frg_colrow[i] = static_cast<ElementScalar>(tCrCol_mn(epi_v * FragmentSize + i) * tCrRow_mn(epi_v * FragmentSize + i));
+      }
+      return frg_colrow;
+    }
+
+  };
+
+  template <
+    bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
+    class... Args
+  >
+  CUTLASS_DEVICE auto
+  get_consumer_store_callbacks(ConsumerStoreArgs<Args...> const& args) {
+
+    auto [M, N, K, L] = args.problem_shape_mnkl;
+    Tensor mCol = make_tensor(make_gmem_ptr(params.ptr_col), make_shape(M,N,L), params.dCol);
+    Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row), make_shape(M,N,L), params.dRow);
+    Tensor tCgCol = sm90_partition_for_epilogue<ReferenceSrc>(                         // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
+      mCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
+    Tensor tCgRow = sm90_partition_for_epilogue<ReferenceSrc>(                         // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
+      mRow, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
+    Tensor tCrCol = make_tensor_like(tCgCol);                                          // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
+    Tensor tCrRow = make_tensor_like(tCgRow);                                          // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
+
+    return ConsumerStoreCallbacks<
+      decltype(tCgCol), decltype(tCrCol),
+      decltype(tCgRow), decltype(tCrRow)
+    >(
+      cute::move(tCgCol), cute::move(tCrCol),
+      cute::move(tCgRow), cute::move(tCrRow),
+      params
+    );
+  }
+
+};
+
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 // Batch matrix broadcast

include/cutlass/epilogue/fusion/sm90_visitor_store_tma_warpspecialized.hpp

@@ -293,11 +293,11 @@ template <
   class LayoutOrStrideMNL,
   class SmemLayoutAtom, // Unused
   class CopyOpR2S,      // Unused
-  int Alignment, 
+  int Alignment,
   bool EnableNullptr
 >
 struct Sm90AuxStore<
-  0, EpilogueTile, Element, RoundStyle, LayoutOrStrideMNL, 
+  0, EpilogueTile, Element, RoundStyle, LayoutOrStrideMNL,
   SmemLayoutAtom, CopyOpR2S, Alignment, EnableNullptr
 > {
   using ElementAux = Element;
@@ -343,7 +343,7 @@ struct Sm90AuxStore<
   CUTLASS_HOST_DEVICE
   Sm90AuxStore(Params const& params, SharedStorage const& shared_storage)
     : params_ptr(&params) { }
-  
+
   Params const* params_ptr;
 
   CUTLASS_DEVICE bool
@@ -381,7 +381,7 @@ struct Sm90AuxStore<
         tC_cAux(cute::forward<CTensorR2G>(tC_cAux)),
         problem_shape_mnl(problem_shape_mnl),
         params_ptr(params_ptr) {}
-    
+
     GTensorR2G tC_gAux;
     RTensor tC_rAux;
     CTensorR2G tC_cAux;
@@ -414,7 +414,7 @@ struct Sm90AuxStore<
 
       Tensor tC_cAux_mn = tC_cAux(_,_,_,epi_m,epi_n);
       Tensor tC_cAux_vec = tensor<1>(zipped_divide(coalesce(tC_cAux_mn), MCL.compose(Int<V>{})));
-      
+
       Tensor tC_gAux_vec = recast<Array<Element, V>>(coalesce(tC_gAux(_,_,_,epi_m,epi_n)));
       Tensor tC_rAux_vec = recast<Array<Element, V>>(coalesce(tC_rAux));
 
@@ -451,7 +451,7 @@ struct Sm90AuxStore<
     // Predication support
     Tensor coordAux = make_identity_tensor(shape(mAux));
     Tensor tC_cAux = sm90_partition_for_epilogue<ReferenceSrc>(
-                      coordAux, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);   
+                      coordAux, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
 
     return ConsumerStoreCallbacks<decltype(tC_gAux), decltype(tC_rAux), decltype(tC_cAux), decltype(problem_shape_mnl)>(
       cute::move(tC_gAux),
@@ -703,7 +703,6 @@ public:
     else if constexpr (FinalReduction) {
       auto problem_shape_mnkl = append<4>(problem_shape, 1);
       auto [M, N, K, L] = problem_shape_mnkl;
-
       auto [tile_M, tile_N, tile_K] = CtaTileShapeMNK{};
       size_t tile_counters_offset = product(ceil_div(make_shape(size<>(M), size<>(N), L), make_shape(tile_M, tile_N))) * tile_N * sizeof(ElementCompute);
       tile_counters_offset = round_nearest(tile_counters_offset, MinWorkspaceAlignment);
@@ -753,19 +752,18 @@ public:
   static cutlass::Status
   initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream,
     CudaHostAdapter* cuda_adapter = nullptr) {
-#if !defined(CUTLASS_SKIP_REDUCTION_INIT)
-    auto problem_shape_mnkl = append<4>(problem_shape, 1);
-    auto [M, N, K, L] = problem_shape_mnkl;
     if constexpr (IsAtomic) {
+      auto problem_shape_mnkl = append<4>(problem_shape, 1);
+      auto [M, N, K, L] = problem_shape_mnkl;
       Layout mRow_layout = make_layout(make_shape(size<>(M),size<>(N),size<>(L)), args.dRow);
       if (args.ptr_row != nullptr) {
         return fill_workspace(args.ptr_row, ElementOutput(args.reduction_identity), cosize(mRow_layout), stream, cuda_adapter);
       }
       return Status::kSuccess;
     }
-    else
-#endif 
-    if constexpr (FinalReduction) {
+    else if constexpr (FinalReduction) {
+      auto problem_shape_mnkl = append<4>(problem_shape, 1);
+      auto [M, N, K, L] = problem_shape_mnkl;
       auto [tile_M, tile_N, tile_K] = CtaTileShapeMNK{};
       size_t tile_counters_offset = product(ceil_div(make_shape(size<>(M),size<>(N),L), make_shape(tile_M, tile_N))) * tile_N * sizeof(ElementCompute);
       tile_counters_offset = round_nearest(tile_counters_offset, MinWorkspaceAlignment);
@@ -939,7 +937,7 @@ public:
             for (int v = 0; v < size(frg_A); ++v) {
               // Step1: swap
               if (not (lane_m & m)) { // the first half of threads swap fragments from the first half of data to the second
-                swap(frg_A(v), frg_B(v));
+                cutlass::swap(frg_A(v), frg_B(v));
               }
 
               // Step2: shuffle
@@ -1023,9 +1021,7 @@ public:
         }
         else {
           if (is_reduced_lane) {
-            // Filter so we don't issue redundant copies over stride-0 modes
-            // (only works if 0-strides are in same location, which is by construction)
-            copy_aligned(filter(tCrRow), recast<ElementGmem>(filter(tCgBuf)));
+            copy_aligned(tCrRow, recast<ElementGmem>(tCgBuf));
           }
         }
         sync_fn();
@@ -1054,9 +1050,7 @@ public:
         }
         else {
           if (is_reduced_lane) {
-            // Filter so we don't issue redunant copies over stride-0 modes
-            // (only works if 0-strides are in same location, which is by construction)
-            copy_aligned(filter(tCrRow), filter(tCsBuf));
+            copy_aligned(tCrRow, tCsBuf);
           }
         }
         sync_fn();
@@ -1296,7 +1290,6 @@ public:
     else if constexpr (FinalReduction) {
       auto problem_shape_mnkl = append<4>(problem_shape, 1);
       auto [M, N, K, L] = problem_shape_mnkl;
-
       auto [tile_M, tile_N, tile_K] = CtaTileShapeMNK{};
       size_t tile_counters_offset = product(ceil_div(make_shape(M,N,L), make_shape(tile_M, tile_N))) * tile_M * sizeof(ElementCompute);
       tile_counters_offset = round_nearest(tile_counters_offset, MinWorkspaceAlignment);
@@ -1348,19 +1341,18 @@ public:
   static cutlass::Status
   initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream,
     CudaHostAdapter* cuda_adapter = nullptr) {
-#if !defined(CUTLASS_SKIP_REDUCTION_INIT)
-    auto problem_shape_mnkl = append<4>(problem_shape, 1);
-    auto [M, N, K, L] = problem_shape_mnkl;
     if constexpr (IsAtomic) {
+      auto problem_shape_mnkl = append<4>(problem_shape, 1);
+      auto [M, N, K, L] = problem_shape_mnkl;
       Layout mCol_layout = make_layout(make_shape(size<>(M),size<>(N),size<>(L)), args.dCol);
       if (args.ptr_col != nullptr) {
         return fill_workspace(args.ptr_col, ElementOutput(args.reduction_identity), cosize(mCol_layout), stream, cuda_adapter);
       }
       return Status::kSuccess;
     }
-    else
-#endif 
-    if constexpr (FinalReduction) {
+    else if constexpr (FinalReduction) {
+      auto problem_shape_mnkl = append<4>(problem_shape, 1);
+      auto [M, N, K, L] = problem_shape_mnkl;
       auto [tile_M, tile_N, tile_K] = CtaTileShapeMNK{};
       size_t tile_counters_offset = product(ceil_div(make_shape(M,N,L), make_shape(tile_M, tile_N))) * tile_M * sizeof(ElementCompute);
       tile_counters_offset = round_nearest(tile_counters_offset, MinWorkspaceAlignment);
@@ -1522,9 +1514,7 @@ public:
         using ElementGmem = cute::conditional_t<FinalReduction, ElementCompute volatile, ElementCompute>;
         Tensor tCgBuf = sm90_partition_for_epilogue<ReferenceSrc>(gBuf_nl(_,_,n,l), epi_tile, tiled_copy, thread_idx);
         if (is_reduced_lane) {
-          // Filter so we don't issue redundant copies over stride-0 modes
-          // (only works if 0-strides are in same location, which is by construction)
-          copy_aligned(filter(tCrCol), recast<ElementGmem>(filter(tCgBuf)));
+          copy_aligned(tCrCol, recast<ElementGmem>(tCgBuf));
         }
         sync_fn();
       }
@@ -1542,9 +1532,7 @@ public:
         // Dump warp reduction to smem workspace
         Tensor tCsBuf = sm90_partition_for_epilogue<ReferenceSrc>(sBuf(_,_,get<1>(warp_mn)), epi_tile, tiled_copy, thread_idx);
         if (is_reduced_lane) {
-          // Filter so we don't issue redunant copies over stride-0 modes
-          // (only works if 0-strides are in same location, which is by construction)
-          copy_aligned(filter(tCrCol), filter(tCsBuf));
+          copy_aligned(tCrCol, tCsBuf);
         }
         sync_fn();
 

include/cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp

@@ -300,7 +300,6 @@ struct Sm90VisitorImplBase {
   tuple<Ops...> ops;
 };
 
-
 template <class... Ops>
 struct Sm90VisitorImpl : Sm90VisitorImplBase<Ops...> {
 
@@ -658,7 +657,6 @@ struct Sm90SplitTreeVisitor : Sm90VisitorImpl<InputTree, AuxOutTrees..., OutputT
     return ConsumerStoreCallbacks<decltype(callbacks_tuple)>(std::move(callbacks_tuple));
   }
 };
-
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 template<

include/cutlass/epilogue/thread/activation.h

@@ -258,6 +258,54 @@ struct LeakyReLU<Array<T, N> > {
   }
 };
 
+// Y = min((X <= threshold ? 0 : X), upper_bound)
+template <typename T>
+struct ThresholdReLU {
+  static constexpr bool kIsHeavy = false;
+
+  struct Arguments {
+    T threshold = T(0);
+    T upper_bound = CUTLASS_STL_NAMESPACE::numeric_limits<T>::max();
+  };
+
+  CUTLASS_HOST_DEVICE
+  T operator()(T value, T threshold, T upper_bound) const {
+    minimum_with_nan_propagation<T> mn;
+    
+    return mn((value <= threshold ? T(0) : value), upper_bound);
+  }
+
+  CUTLASS_HOST_DEVICE
+  T operator()(T value, Arguments const& args = Arguments()) const {
+    return operator()(value, args.threshold, args.upper_bound);
+  }
+};
+
+template <typename T, int N>
+struct ThresholdReLU<Array<T,N>> {
+  static constexpr bool kIsHeavy = false;
+
+  using Arguments = typename ThresholdReLU<T>::Arguments;
+
+  CUTLASS_HOST_DEVICE
+  Array<T,N> operator()(Array<T,N> const& values, T threshold, T upper_bound) const {
+    ThresholdReLU<T> relu;
+
+    Array<T,N> retvals;
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      retvals[i] = relu(values[i], threshold, upper_bound);    
+    }
+
+    return retvals;
+  }
+
+  CUTLASS_HOST_DEVICE
+  Array<T,N> operator()(Array<T,N> const& values, Arguments const& args = Arguments()) const {
+    return operator()(values, args.threshold, args.upper_bound);
+  }
+};
+
 // Tanh operator
 template <typename T>
 struct Tanh {
@@ -311,26 +359,7 @@ struct Sigmoid {
 };
 
 template <typename T, int N>
-struct Sigmoid<Array<T, N> > {
-  static const bool kIsHeavy = true;
-
-  CUTLASS_HOST_DEVICE
-  Array<T, N> operator()(Array<T, N> const &value) const {
-    Array<T, N> y;
-    Sigmoid<T> sigmoid_op;
-
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 0; i < N; ++i) {
-      y[i] = sigmoid_op(value[i]);
-    }
-
-    return y;
-  }
-};
-
-template <int N>
-struct Sigmoid<Array<half_t, N>> {
-  using T = half_t;
+struct Sigmoid<Array<T, N>> {
   static const bool kIsHeavy = true;
 
   CUTLASS_HOST_DEVICE
@@ -450,6 +479,9 @@ struct HardSwish<Array<half_t, N> > {
   }
 };
 
+template <typename T>
+using ScaledHardSwish = Scale<HardSwish<T>>;
+
 //
 // GELU function definitions implemented as described by
 //   Hendrycks, D., and Gimpel, K. in

include/cutlass/epilogue/thread/linear_combination.h

@@ -169,7 +169,7 @@ public:
 
   /// Constructs the function object, possibly loading from pointers in host memory
   CUTLASS_HOST_DEVICE
-  LinearCombination(Params const &params, int group_idx = 0) {
+  explicit LinearCombination(Params const &params, int group_idx) {
     if (params.alpha_ptr_array != nullptr && params.alpha_ptr_array[group_idx] != nullptr) {
       alpha_ = *(params.alpha_ptr_array[group_idx]);
     }
@@ -190,6 +190,10 @@ public:
     }
   }
 
+  CUTLASS_HOST_DEVICE
+  explicit LinearCombination(const Params & params) 
+  : LinearCombination(params, /* group_idx */ 0) { }
+
   /// Returns true if source is needed
   CUTLASS_HOST_DEVICE
   bool is_source_needed() const {

include/cutlass/epilogue/threadblock/epilogue.h

@@ -39,11 +39,7 @@
 
 #pragma once
 
-#if defined(__CUDACC_RTC__)
 #include <cuda/std/cassert>
-#else
-#include <assert.h>
-#endif
 
 #include "cutlass/cutlass.h"
 #include "cutlass/numeric_types.h"
@@ -478,6 +474,12 @@ public:
     // Iterate over accumulator tile
     //
 
+    #ifdef __clang__
+    #pragma clang diagnostic push
+    #pragma clang diagnostic ignored "-Wcuda-compat"
+    // Turn off clangs warning about loop unroll argument using parens.
+    #endif
+
     #pragma unroll(IterationsUnroll ? OutputTileIterator::kIterations : 1)
     for (int iter = 0; iter < OutputTileIterator::kIterations; ++iter)
     {
@@ -531,6 +533,10 @@ public:
       destination_iterator.store(output_fragment);
       ++destination_iterator;
     }
+    
+    #ifdef __clang__
+    #pragma clang diagnostic pop
+    #endif
   }
 };
 

include/cutlass/epilogue/threadblock/epilogue_base.h

@@ -43,11 +43,7 @@
 #include <utility>
 #endif
 
-#if defined(__CUDACC_RTC__)
 #include <cuda/std/cassert>
-#else
-#include <assert.h>
-#endif
 
 #include "cutlass/cutlass.h"
 #include "cutlass/matrix_shape.h"

include/cutlass/epilogue/threadblock/epilogue_gemm_k_reduction.h

@@ -38,11 +38,7 @@
 
 #pragma once
 
-#if defined(__CUDACC_RTC__)
 #include <cuda/std/cassert>
-#else
-#include <assert.h>
-#endif
 
 #include "cutlass/cutlass.h"
 #include "cutlass/numeric_types.h"

include/cutlass/epilogue/threadblock/epilogue_smem_accumulator.h

@@ -38,11 +38,7 @@
 
 #pragma once
 
-#if defined(__CUDACC_RTC__)
 #include <cuda/std/cassert>
-#else
-#include <assert.h>
-#endif
 
 #include "cutlass/cutlass.h"
 #include "cutlass/numeric_types.h"

include/cutlass/epilogue/threadblock/epilogue_streamk_with_broadcast.h

@@ -39,11 +39,11 @@
 
 #pragma once
 
-#if defined(__CUDACC_RTC__)
 #include <cuda/std/cassert>
+
+#if defined(__CUDACC_RTC__)
 #include <cuda/std/utility>
 #else
-#include <assert.h>
 #include <utility>
 #endif
 

include/cutlass/epilogue/threadblock/epilogue_with_absmax.h

@@ -50,11 +50,11 @@
 
 #pragma once
 
-#if defined(__CUDACC_RTC__)
 #include <cuda/std/cassert>
+
+#if defined(__CUDACC_RTC__)
 #include <cuda/std/utility>
 #else
-#include <assert.h>
 #include <utility>
 #endif
 

include/cutlass/epilogue/threadblock/epilogue_with_broadcast.h

@@ -39,11 +39,11 @@
 
 #pragma once
 
-#if defined(__CUDACC_RTC__)
 #include <cuda/std/cassert>
+
+#if defined(__CUDACC_RTC__)
 #include <cuda/std/utility>
 #else
-#include <assert.h>
 #include <utility>
 #endif
 

include/cutlass/epilogue/threadblock/epilogue_with_reduction.h

@@ -39,11 +39,7 @@
 
 #pragma once
 
-#if defined(__CUDACC_RTC__)
 #include <cuda/std/cassert>
-#else
-#include <assert.h>
-#endif
 
 #include "cutlass/cutlass.h"
 #include "cutlass/array.h"

include/cutlass/epilogue/threadblock/epilogue_with_visitor_callbacks.h

@@ -303,6 +303,12 @@ public:
       // Pipeline Loop
       //
 
+      #ifdef __clang__
+      #pragma clang diagnostic push
+      #pragma clang diagnostic ignored "-Wcuda-compat"
+      // Turn off clang warning about loop unroll argument using parens.
+      #endif
+
       #pragma unroll(IterationsUnroll ? kIterations : 1)
       for (int iter_idx = 1; iter_idx < kIterations + 1; ++iter_idx) {
 
@@ -377,8 +383,19 @@ public:
 
         callbacks.end_step(iter_idx-1);
       }
+
+      #ifdef __clang__
+      #pragma clang diagnostic pop
+      #endif
+
     } else {
 
+      #ifdef __clang__
+      #pragma clang diagnostic push
+      #pragma clang diagnostic ignored "-Wcuda-compat"
+      // Turn off clang warning about loop unroll argument using parens.
+      #endif
+
       #pragma unroll(IterationsUnroll ? kIterations : 1)
       for (int iter_idx = 0; iter_idx < kIterations; ++iter_idx) {
 
@@ -459,6 +476,11 @@ public:
 
         callbacks.end_step(iter_idx);
       }
+
+      #ifdef __clang__
+      #pragma clang diagnostic pop
+      #endif
+
     }
 
     callbacks.end_epilogue();

include/cutlass/epilogue/threadblock/fusion/visitor_load.hpp

@@ -335,7 +335,8 @@ struct VisitorAuxLoad{
 template<
   class ThreadMap,
   class Element,
-  class StrideMNL
+  class StrideMNL,
+  bool EnableNullptr = true // Fallback scalar broadcast for nullptr params
 >
 struct VisitorRowBroadcast {
 
@@ -399,6 +400,16 @@ struct VisitorRowBroadcast {
 
     CUTLASS_DEVICE void
     begin_epilogue() {
+      if constexpr (EnableNullptr) {
+        if (params_ptr->ptr_row == nullptr) {
+          auto tC_rRow_vec = recast<Array<Element, VecLength>>(coalesce(tC_rRow));
+          CUTLASS_PRAGMA_UNROLL
+          for (int i = 0; i < size(tC_rRow_vec); ++i) {
+            tC_rRow_vec[i].fill(params_ptr->null_default);
+          }
+          return;
+        }
+      }
       clear(tC_rRow);
       auto src_v = filter(tC_gRow);
       auto coord_v = filter(tC_cRow);
@@ -406,7 +417,7 @@ struct VisitorRowBroadcast {
       CUTLASS_PRAGMA_UNROLL
       for (int i = 0; i < size(src_v); ++i) {
         bool guard = get<1>(coord_v(i)) < n;
-        cutlass::arch::global_load<VecType, sizeof(VecType)>(dst_v(i), (void const*)&src_v(i), guard);
+        cutlass::arch::global_load<VecType, sizeof(VecType)>(dst_v(i), (void const *)&src_v(i), guard);
       }
     }
 
@@ -464,7 +475,8 @@ struct VisitorRowBroadcast {
 template<
   class ThreadMap,
   class Element,
-  class StrideMNL = Stride<_1,_0,_0>
+  class StrideMNL = Stride<_1,_0,_0>,
+  bool EnableNullptr = true // Fallback scalar broadcast for nullptr params
 >
 struct VisitorColBroadcast {
 
@@ -523,6 +535,12 @@ struct VisitorColBroadcast {
 
     CUTLASS_DEVICE void
     begin_epilogue() {
+      if constexpr (EnableNullptr) {
+        if (params_ptr->ptr_col == nullptr) {
+          fill(tC_rCol, params_ptr->null_default);
+          return;
+        }
+      }
       clear(tC_rCol);
       Tensor pred = make_tensor<bool>(shape(tC_gCol));
       CUTLASS_PRAGMA_UNROLL

include/cutlass/epilogue/threadblock/fusion/visitor_store.hpp

@@ -519,10 +519,7 @@ struct VisitorRowReduction {
       // Guard against uses of the existing SMEM tile
       __syncthreads();
 
-      CUTLASS_PRAGMA_UNROLL
-      for (int i = 0; i < size(tRS_rSrc); ++i) {
-        copy_vec<VecType>(filter(tRS_rSrc), filter(tRS_sRows));
-      }
+      copy(tRS_rSrc, tRS_sRows);
 
       __syncthreads();
 

include/cutlass/epilogue/threadblock/output_tile_thread_map.h

@@ -391,7 +391,7 @@ struct OutputTileOptimalThreadMap {
     1>;
 
   /// Initial offset function
-  CUTLASS_DEVICE
+  CUTLASS_HOST_DEVICE
   static MatrixCoord initial_offset(int thread_idx) {
 
 //    int warp_idx = __shfl_sync(0xffffffff, thread_idx / kWarpSize, 0);
@@ -462,7 +462,7 @@ struct OutputTileOptimalThreadMap {
     static int const kThreads = Threads;
 
     /// Function to compute each thread's initial offset
-    CUTLASS_DEVICE
+    CUTLASS_HOST_DEVICE
     static MatrixCoord initial_offset(int thread_idx) {
 
 //      int warp_idx = __shfl_sync(0xffffffff, thread_idx / kWarpSize, 0);

include/cutlass/epilogue/warp/tile_iterator_tensor_op_mixed.h

@@ -212,15 +212,23 @@ public:
     // When the optimization is enabled, small tiles require separate logic.
     bool kN32_optimization = (WarpShape::kN * Detail::kLanesInQuad * Policy::kElementsPerAccess * sizeof_bits<Element>::value) % 1024 == 0;
     if (kN32_optimization) {
+      
       int ptr_idx = ((warp_column_ * sizeof_bits<Element>::value) / 1024) % Detail::kPointerCount;
+      
       if (ptr_idx == 0) {
         ptr = pointers_[0];
       } else if (ptr_idx == 1) {
-        ptr = pointers_[1];
+	if constexpr (AccessType::kElements >= 2) {
+          ptr = pointers_[1];
+	}
       } else if (ptr_idx == 2) {
-        ptr = pointers_[2];
+	if constexpr (AccessType::kElements >= 3) {
+          ptr = pointers_[2];
+	}
       } else if (ptr_idx == 3) {
-        ptr = pointers_[3];
+	if constexpr (AccessType::kElements >= 4) {
+          ptr = pointers_[3];
+	}
       }
     }
 

include/cutlass/fast_math.h

@@ -38,7 +38,7 @@
 #include <cmath>
 #include <type_traits>
 #endif
-
+#include <cuda/std/utility>
 #include "cutlass/cutlass.h"
 #include "cutlass/array.h"
 #include "cutlass/uint128.h"
@@ -54,12 +54,7 @@ namespace cutlass {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-template <typename T>
-CUTLASS_HOST_DEVICE void swap(T &lhs, T &rhs) {
-  T tmp = lhs;
-  lhs = rhs;
-  rhs = tmp;
-}
+using ::cuda::std::swap;
 
 /******************************************************************************
  * Static math utilities

include/cutlass/float8.h

@@ -1053,8 +1053,8 @@ float_e5m2_t::float_e5m2_t(float_e4m3_t x) {
 /// datatype in runtime argument list.
 ///
 /// Currently supported runtime datatypes compatible with type_erased_dynamic_float8_t:
-///   QMMAFormat::E5M2
-///   QMMAFormat::E4M3
+///   MXF8F6F4Format::E5M2
+///   MXF8F6F4Format::E4M3
 ///
 ///////////////////////////////////////////////////////////////
 

include/cutlass/floating_point_nvrtc.h

@@ -35,6 +35,12 @@
 
 #pragma once
 
+#include <cutlass/detail/helper_macros.hpp> // CUTLASS_HOST_DEVICE
+#include <cutlass/platform/platform.h> // uint32_t
+#if !defined(__CUDACC_RTC__)
+#include <cstring> // std::memcpy
+#endif
+
 namespace cutlass {
 
 ///////////////////////////////////////////////////////////////////////////////////////////////////

include/cutlass/functional.h

@@ -50,7 +50,7 @@
 
 #ifdef _MSC_VER
 // Provides support for alternate operators such as 'and', 'or', ...
-#include <iso646.h>
+#include <ciso646>
 #endif // _MSC_VER
 
 namespace cutlass {

include/cutlass/gemm/collective/builders/sm90_gmma_builder.inl

@@ -35,6 +35,8 @@
 #include "cutlass/pipeline/sm90_pipeline.hpp"
 #include "cutlass/gemm/collective/collective_mma_decl.hpp"
 #include "cutlass/gemm/collective/collective_builder_decl.hpp"
+#include "cute/arch/cluster_sm90.hpp"
+#include "cute/tensor.hpp"
 
 // SM90 Collective Builders should be used only starting CUDA 12.0
 #if (__CUDACC_VER_MAJOR__ >= 12)
@@ -236,8 +238,9 @@ struct CollectiveBuilder<
       GmmaMajorA, ElementAMma, decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
   using SmemLayoutAtomB = decltype(detail::ss_smem_selector<
       GmmaMajorB, ElementBMma, decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
-
-  static constexpr int Sm90ReducedSmemCapacityBytes = detail::sm90_smem_capacity_bytes;
+ 
+  static constexpr int Sm90ReducedSmemCapacityBytes =
+    detail::sm90_smem_capacity_bytes;
 
   static constexpr int PipelineStages = detail::compute_stage_count_or_override<Sm90ReducedSmemCapacityBytes,
       ElementAMma, ElementBMma, TileShape_MNK>(StageCountType{});
@@ -343,7 +346,7 @@ public:
       return t;
     }
     else {
-      return cute::stride(t);
+        return cute::stride(t);
     }
   }
 
@@ -415,15 +418,15 @@ public:
   static constexpr int KernelSmemCarveout = static_cast<int>(TensorMapStorage);
   static constexpr int Sm90ReducedSmemCapacityBytes = detail::sm90_smem_capacity_bytes - KernelSmemCarveout;
 
-  static constexpr int PipelineStages = IsMixedInput ? 
-      detail::compute_stage_count_or_override_single_affine_transformed_input<detail::sm90_smem_capacity_bytes,
-          RealElementA, RealElementB, ElementScale, ElementZero, TileShape_MNK, StageCountType::bytes, SmemAlignment>(StageCountType{}) :
-      detail::compute_stage_count_or_override<detail::sm90_smem_capacity_bytes,
-        ElementAMma, ElementBMma, TileShape_MNK, StageCountType::bytes, SmemAlignment>(StageCountType{});
+  static constexpr int PipelineStages = IsMixedInput ?
+        detail::compute_stage_count_or_override_single_affine_transformed_input<detail::sm90_smem_capacity_bytes,
+          RealElementA, RealElementB, ElementScale, ElementZero, TileShape_MNK, StageCountType::bytes, SmemAlignment>(StageCountType{})
+      : detail::compute_stage_count_or_override<detail::sm90_smem_capacity_bytes,
+          ElementAMma, ElementBMma, TileShape_MNK, StageCountType::bytes, SmemAlignment>(StageCountType{});
       
   using DispatchPolicy = cute::conditional_t<IsMixedInput,
-      MainloopSm90TmaGmmaRmemAWarpSpecializedMixedInput<PipelineStages, ClusterShape_MNK, KernelScheduleType>,
-      MainloopSm90TmaGmmaRmemAWarpSpecialized<PipelineStages, ClusterShape_MNK, KernelScheduleType>>;
+        MainloopSm90TmaGmmaRmemAWarpSpecializedMixedInput<PipelineStages, ClusterShape_MNK, KernelScheduleType>
+      , MainloopSm90TmaGmmaRmemAWarpSpecialized<PipelineStages, ClusterShape_MNK, KernelScheduleType>>;
 
   using SmemCopyAtomA = cute::conditional_t<SwapAB, void, Copy_Atom<cute::AutoVectorizingCopy, ElementA>>;
   using SmemCopyAtomB = cute::conditional_t<SwapAB, Copy_Atom<cute::AutoVectorizingCopy, ElementB>, void>;
@@ -761,13 +764,13 @@ struct CollectiveBuilder<
   static constexpr int NumLoadWarpGroups = cute::is_same_v<KernelScheduleType, KernelCpAsyncWarpSpecialized> ? 2 : 1;
 
   using AlignmentTypeA = cute::uint_byte_t<static_cast<int>(sizeof(ElementA)) * AlignmentA>;
-  using GmemCopyAtomA = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<AlignmentTypeA>, ElementA>;
+  using GmemCopyAtomA = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<AlignmentTypeA>, ElementA>;
   using GmemTiledCopyA = decltype(detail::make_simt_gmem_tiled_copy<
       GmemCopyAtomA, NumThreadsPerWarpGroup * NumLoadWarpGroups, AlignmentA, TagToStrideA_t<GmemLayoutATag>,
       decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
 
   using AlignmentTypeB = cute::uint_byte_t<static_cast<int>(sizeof(ElementB)) * AlignmentB>;
-  using GmemCopyAtomB = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<AlignmentTypeB>, ElementB>;
+  using GmemCopyAtomB = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<AlignmentTypeB>, ElementB>;
   using GmemTiledCopyB = decltype(detail::make_simt_gmem_tiled_copy<
       GmemCopyAtomB, NumThreadsPerWarpGroup * NumLoadWarpGroups, AlignmentB, TagToStrideB_t<GmemLayoutBTag>,
       decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
@@ -867,13 +870,13 @@ struct CollectiveBuilder<
   static constexpr int NumLoadWarpGroups = 1;
 
   using AlignmentTypeA = cute::uint_byte_t<static_cast<int>(sizeof(ElementA)) * AlignmentA>;
-  using GmemCopyAtomA = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<AlignmentTypeA>, ElementA>;
+  using GmemCopyAtomA = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<AlignmentTypeA>, ElementA>;
   using GmemTiledCopyA = decltype(detail::make_simt_gmem_tiled_copy<
       GmemCopyAtomA, NumThreadsPerWarpGroup * NumLoadWarpGroups, AlignmentA, TagToStrideA_t<GmemLayoutATag>,
       decltype(cute::get<0>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());
 
   using AlignmentTypeB = cute::uint_byte_t<static_cast<int>(sizeof(ElementB)) * AlignmentB>;
-  using GmemCopyAtomB = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<AlignmentTypeB>, ElementB>;  
+  using GmemCopyAtomB = cute::Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS_ZFILL<AlignmentTypeB>, ElementB>;
   using GmemTiledCopyB = decltype(detail::make_simt_gmem_tiled_copy<
       GmemCopyAtomB, NumThreadsPerWarpGroup * NumLoadWarpGroups, AlignmentB, TagToStrideB_t<GmemLayoutBTag>,
       decltype(cute::get<1>(TileShape_MNK{})), decltype(cute::get<2>(TileShape_MNK{}))>());

include/cutlass/gemm/collective/collective_builder_decl.hpp

@@ -54,6 +54,18 @@ struct StageCountAutoCarveout {
   explicit StageCountAutoCarveout(cute::Int<carveout_bytes>) {}
 };
 
+namespace detail {
+
+// Forward Declaration
+template<class CollectiveEpilogue>
+constexpr int
+compute_carveout_from_epi();
+
+} // namespace detail
+
+template<class CollectiveEpilogue>
+struct StageCountAutoCarveoutEpi : StageCountAutoCarveout<detail::compute_carveout_from_epi<CollectiveEpilogue>()> {};
+
 using StageCountAuto = StageCountAutoCarveout<0>;
 
 // Used to automatically let the builder pick the kernel schedule.

include/cutlass/gemm/collective/collective_mma.hpp

@@ -41,9 +41,10 @@
 #include "cutlass/gemm/collective/sm90_mma_multistage_gmma_rs_warpspecialized.hpp"
 #include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss.hpp"
 #include "cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized.hpp"
-#include "cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp"
+#include "cutlass/gemm/collective/sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp" 
 #include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized.hpp"
 #include "cutlass/gemm/collective/sm90_sparse_mma_tma_gmma_ss_warpspecialized.hpp"
 #include "cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized.hpp"
 #include "cutlass/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8.hpp"
+
 /////////////////////////////////////////////////////////////////////////////////////////////////

include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized.hpp

@@ -374,7 +374,7 @@ struct CollectiveMma<
       // Prepare the TMA loads for A and B
       //
 
-      constexpr uint32_t cluster_shape_x = get<0>(DispatchPolicy::ClusterShape());
+      constexpr uint32_t cluster_shape_x = get<0>(typename DispatchPolicy::ClusterShape());
       uint2 cluster_local_block_id = {block_rank_in_cluster % cluster_shape_x, block_rank_in_cluster / cluster_shape_x};
 
       Tensor gA_mkl = get<0>(load_inputs);

include/cutlass/gemm/device/gemm_universal_adapter.h

@@ -85,13 +85,40 @@ class GemmUniversalAdapter;
 ////////////////////////////// CUTLASS 3.x API /////////////////////////////////
 ////////////////////////////////////////////////////////////////////////////////
 
+namespace detail {
+
+// Work-around for some DispatchPolicy types not having a Stages member.
+// In that case, the Stages value is 0.  Most code should static_assert
+// that the number of stages is valid.
+
+// Whether DispatchPolicy::Stages is valid.
+// It should also be convertible to int, but if not, that will show up
+// as a build error when GemmUniversalAdapter attempts to assign it to kStages.
+template <class DispatchPolicy, class Enable = void>
+struct has_Stages : cute::false_type {};
+
+template <class DispatchPolicy>
+struct has_Stages<DispatchPolicy, cute::void_t<decltype(DispatchPolicy::Stages)>> : cute::true_type {};
+
+template<class DispatchPolicy>
+constexpr int stages_member(DispatchPolicy) {
+  if constexpr (has_Stages<DispatchPolicy>::value) {
+    return DispatchPolicy::Stages;
+  }
+  else {
+    return 0;
+  }
+}
+
+} // namespace detail
+
 template <class GemmKernel_>
 class GemmUniversalAdapter<
   GemmKernel_,
-  cute::enable_if_t<gemm::detail::IsCutlass3GemmKernel<GemmKernel_>::value>>
+  cute::enable_if_t<gemm::detail::IsCutlass3GemmKernel<GetUnderlyingKernel_t<GemmKernel_>>::value>>
 {
 public:
-  using GemmKernel = GemmKernel_;
+  using GemmKernel = GetUnderlyingKernel_t<GemmKernel_>;
   using TileShape = typename GemmKernel::TileShape;
   using ElementA = typename GemmKernel::ElementA;
   using ElementB = typename GemmKernel::ElementB;
@@ -158,7 +185,7 @@ public:
       CUTE_STATIC_V(cute::tile_size<1>(typename CollectiveMainloop::TiledMma{})) / WarpsInMmaN,
       CUTE_STATIC_V(cute::tile_size<2>(typename CollectiveMainloop::TiledMma{}))>;
 
-  static int constexpr kStages = CollectiveMainloop::DispatchPolicy::Stages;
+  static int constexpr kStages = detail::stages_member(typename CollectiveMainloop::DispatchPolicy{});
 
   // Inspect TiledCopy for A and B to compute the alignment size
   static int constexpr kAlignmentA = cutlass::detail::get_alignment_count_from_gmem_tiled_copy<
@@ -336,7 +363,7 @@ public:
   }
 
   /// Primary run() entry point API that is static allowing users to create and manage their own params.
-  /// Supplied params struct must be construct by calling GemmKernel::to_underling_arguments()
+  /// Supplied params struct must be construct by calling GemmKernel::to_underlying_arguments()
   static Status
   run(Params& params,
       cudaStream_t stream = nullptr,
@@ -358,10 +385,10 @@ public:
       [[maybe_unused]] constexpr bool is_static_1x1x1 =
         cute::is_static_v<typename GemmKernel::DispatchPolicy::ClusterShape> and
         cute::size(typename GemmKernel::DispatchPolicy::ClusterShape{}) == 1;
-      dim3 cluster(cute::size<0>(typename GemmKernel::DispatchPolicy::ClusterShape{}),
-                   cute::size<1>(typename GemmKernel::DispatchPolicy::ClusterShape{}),
-                   cute::size<2>(typename GemmKernel::DispatchPolicy::ClusterShape{}));
-      void* kernel_params[] = {&params};
+      [[maybe_unused]] dim3 cluster(cute::size<0>(typename GemmKernel::DispatchPolicy::ClusterShape{}),
+        cute::size<1>(typename GemmKernel::DispatchPolicy::ClusterShape{}),
+        cute::size<2>(typename GemmKernel::DispatchPolicy::ClusterShape{}));
+      [[maybe_unused]] void* kernel_params[] = {&params};
 
       if constexpr (kEnableCudaHostAdapter) {
         //
@@ -377,13 +404,23 @@ public:
 #if (CUTLASS_DEBUG_TRACE_LEVEL > 1)
           CUTLASS_TRACE_HOST("GemmUniversal::run: Launching kernel with CUDA host adapter");
 #endif
-          launch_result = cuda_adapter->launch(grid,
-                                               cluster,
-                                               block,
-                                               smem_size,
-                                               stream,
-                                               kernel_params,
-                                               0);
+          if constexpr (is_static_1x1x1) {
+            launch_result = cuda_adapter->launch(grid,
+                                                block,
+                                                smem_size,
+                                                stream,
+                                                kernel_params,
+                                                0);
+          }
+          else {
+            launch_result = cuda_adapter->launch(grid,
+                                                cluster,
+                                                block,
+                                                smem_size,
+                                                stream,
+                                                kernel_params,
+                                                0);
+          }
         }
         else {
           CUTLASS_TRACE_HOST("GemmUniversal::run: kEnableCudaHostAdapter is true, but CUDA host adapter is null");
@@ -392,8 +429,10 @@ public:
       }
       else {
         CUTLASS_ASSERT(cuda_adapter == nullptr);
-        void const* kernel = (void const*) device_kernel<GemmKernel>;
-        if constexpr (GemmKernel::ArchTag::kMinComputeCapability == 90) {
+        [[maybe_unused]] void const* kernel = (void const*) device_kernel<GemmKernel>;
+        static constexpr bool kClusterLaunch = GemmKernel::ArchTag::kMinComputeCapability == 90
+                                        ;
+        if constexpr (kClusterLaunch) {
           if constexpr (is_static_1x1x1) {
 #if (CUTLASS_DEBUG_TRACE_LEVEL > 1)
             CUTLASS_TRACE_HOST("GemmUniversal::run: Launching static 1x1x1 kernel");
@@ -526,11 +565,11 @@ public:
 template <class GemmKernel_>
 class GemmUniversalAdapter<
   GemmKernel_,
-  cute::enable_if_t<not gemm::detail::IsCutlass3GemmKernel<GemmKernel_>::value>>
+  cute::enable_if_t<not gemm::detail::IsCutlass3GemmKernel<GetUnderlyingKernel_t<GemmKernel_>>::value>>
 {
 public:
 
-  using GemmKernel = GemmKernel_;
+  using GemmKernel = GetUnderlyingKernel_t<GemmKernel_>;
 
   static bool const kInternalTranspose =
     !cutlass::epilogue::threadblock::detail::is_2x_evt_v<typename GemmKernel::Epilogue> &&  // 2.x EVT does not require internal transpose

include/cutlass/gemm/dispatch_policy.hpp

@@ -105,7 +105,8 @@ struct KernelCpAsyncWarpSpecializedPingpong { };
 struct KernelCpAsyncWarpSpecializedCooperative { };
 struct KernelTma { };
 struct KernelTmaWarpSpecialized { };
-struct KernelTmaWarpSpecializedPingpong { };
+struct KernelTmaWarpSpecializedPingpong { 
+};
 struct KernelTmaWarpSpecializedCooperative { 
 };
 
@@ -247,6 +248,7 @@ struct MainloopSm90TmaGmmaRmemAWarpSpecialized {
     "KernelSchedule must be one of the warp specialized policies");
 };
 
+
 template<
   int Stages_,
   class ClusterShape_ = Shape<_1,_1,_1>,
@@ -310,6 +312,7 @@ struct MainloopSm90TmaGmmaWarpSpecializedSparse {
   using Schedule = KernelSchedule;
 };
 
+
 //////////////////////////////////////////////////////////////////////////////
 
 } // namespace cutlass::gemm

include/cutlass/gemm/group_array_problem_shape.hpp

@@ -69,7 +69,7 @@ struct GroupProblemShape {
   CUTLASS_HOST_DEVICE
   UnderlyingProblemShape const
   get_host_problem_shape(int32_t group_idx) const {
-    return host_problem_shapes[group_idx];
+    return host_problem_shapes != nullptr ? host_problem_shapes[group_idx] : UnderlyingProblemShape{};
   }
 
   CUTLASS_HOST_DEVICE

include/cutlass/gemm/kernel/default_gemm_grouped_per_group_scale.h

@@ -0,0 +1,384 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2024 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.
+ *
+ **************************************************************************************************/
+
+/*! \file
+    \brief 
+      Default kernel-level GEMM definitions combine threadblock-scoped matrix multiply-add with
+      the appropriate threadblock-scoped epilogue.
+  
+      Note, CUTLASS epilogues universally target row-major outputs. Column-major outputs are
+      accommodated by exchanging A and B operands and assuming transposed layouts. Partial
+      specializations here choose 'device::GemmTransposed' to implement this functionality.
+
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+
+#include "cutlass/complex.h"
+#include "cutlass/layout/matrix.h"
+#include "cutlass/numeric_types.h"
+
+#include "cutlass/gemm/kernel/gemm_grouped_per_group_scale.h"
+#include "cutlass/gemm/kernel/gemm_transpose_operands.h"
+#include "cutlass/gemm/kernel/default_gemm.h"
+#include "cutlass/gemm/kernel/default_gemm_complex.h"
+#include "cutlass/gemm/device/default_gemm_configuration.h"
+
+#include "cutlass/layout/permute.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace kernel {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+    /// Element type for A matrix operand
+    typename ElementA_,
+    /// Layout type for A matrix operand
+    typename LayoutA_,
+    /// Complex elementwise transformation on A operand
+    ComplexTransform TransformA,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentA,
+    /// Element type for B matrix operand
+    typename ElementB_,
+    /// Layout type for B matrix operand
+    typename LayoutB_,
+    /// Complex elementwise transformation on B operand
+    ComplexTransform TransformB,
+    /// Access granularity of B matrix in units of elements
+    int kAlignmentB,
+    /// Element type for C and D matrix operands
+    typename ElementC_,
+    /// Layout type for C and D matrix operands
+    typename LayoutC_,
+    /// Element type for internal accumulation
+    typename ElementAccumulator,
+    /// Operator class tag
+    typename OperatorClass,
+    /// Tag indicating architecture to tune for
+    typename ArchTag,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename InstructionShape,
+    /// Epilogue output operator
+    typename EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle,
+    /// Number of stages used in the pipelined mainloop
+    int Stages,
+    /// Whether the schedule of problems to visit has been precomputed
+    GroupScheduleMode GroupScheduleMode_ = GroupScheduleMode::kDeviceOnly,
+    /// Operation performed by GEMM
+    typename Operator = typename device::DefaultGemmConfiguration<
+        OperatorClass, ArchTag, ElementA_, ElementB_, ElementC_,
+        ElementAccumulator>::Operator,
+    /// Use zfill or predicate for out-of-bound cp.async
+    SharedMemoryClearOption SharedMemoryClear = SharedMemoryClearOption::kNone,
+    /// Permute result D
+    typename PermuteDLayout = layout::NoPermute,
+    ///
+    typename Enable = void
+    >
+struct DefaultGemmGroupedPerGroupScale;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Real-valued GEMM kernels
+//
+
+template <
+    /// Element type for A matrix operand
+    typename ElementA,
+    /// Layout type for A matrix operand
+    typename LayoutA,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentA,
+    /// Element type for B matrix operand
+    typename ElementB,
+    /// Layout type for B matrix operand
+    typename LayoutB,
+    /// Access granularity of B matrix in units of elements
+    int kAlignmentB,
+    /// Element type for C and D matrix operands
+    typename ElementC,
+    /// Layout type for C and D matrix operands
+    typename LayoutC,
+    /// Element type for internal accumulation
+    typename ElementAccumulator,
+    /// Operator class tag
+    typename OperatorClass,
+    /// Tag indicating architecture to tune for
+    typename ArchTag,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename InstructionShape,
+    /// Epilogue output operator
+    typename EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle,
+    /// Number of stages used in the pipelined mainloop
+    int Stages,
+    /// Whether the schedule of problems to visit has been precomputed
+    GroupScheduleMode GroupScheduleMode_,
+    /// Operation performed by GEMM
+    typename Operator,
+    /// Use zfill or predicate for out-of-bound cp.async
+    SharedMemoryClearOption SharedMemoryClear,
+    /// Permute result D
+    typename PermuteDLayout
+>
+struct DefaultGemmGroupedPerGroupScale<
+  ElementA,
+  LayoutA,
+  ComplexTransform::kNone,   // transform A
+  kAlignmentA,
+  ElementB,
+  LayoutB,
+  ComplexTransform::kNone,   // transform B
+  kAlignmentB,
+  ElementC,
+  LayoutC,
+  ElementAccumulator,
+  OperatorClass,
+  ArchTag,
+  ThreadblockShape,
+  WarpShape,
+  InstructionShape,
+  EpilogueOutputOp,
+  ThreadblockSwizzle,
+  Stages,
+  GroupScheduleMode_,
+  Operator,
+  SharedMemoryClear,
+  PermuteDLayout,
+  typename platform::enable_if< ! cutlass::is_complex<ElementAccumulator>::value>::type
+> {
+
+  // If true, we must construct a 'transposed-and-exchanged' Mma operator.
+  static bool const kInternalTranspose = platform::is_same<LayoutC, layout::ColumnMajor>::value;
+
+  using MapArguments = kernel::detail::MapArguments<
+    ElementA,
+    LayoutA,
+    ComplexTransform::kNone,
+    kAlignmentA,
+    ElementB,
+    LayoutB,
+    ComplexTransform::kNone,
+    kAlignmentB,
+    LayoutC,
+    kInternalTranspose
+  >;
+
+  // Define the default GEMM kernel
+  using DefaultGemmKernel = typename kernel::DefaultGemm<
+    typename MapArguments::ElementA,
+    typename MapArguments::LayoutA,
+    MapArguments::kAlignmentA,
+    typename MapArguments::ElementB,
+    typename MapArguments::LayoutB,
+    MapArguments::kAlignmentB,
+    ElementC,
+    typename MapArguments::LayoutC,
+    ElementAccumulator,
+    OperatorClass,
+    ArchTag,
+    ThreadblockShape,
+    WarpShape,
+    InstructionShape,
+    EpilogueOutputOp,
+    ThreadblockSwizzle,
+    Stages,
+    true,
+    Operator,
+    SharedMemoryClear,
+    false, /*GatherA*/
+    false, /*GatherB*/
+    false, /*ScatterD*/
+    PermuteDLayout
+  >::GemmKernel;
+
+    /// Define the kernel in terms of the default kernel
+  using GemmKernel = kernel::GemmGroupedPerGroupScale<
+    typename DefaultGemmKernel::Mma,
+    typename DefaultGemmKernel::Epilogue,
+    ThreadblockSwizzle,
+    GroupScheduleMode_,
+    kInternalTranspose
+  >;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+//
+// Complex-valued GEMM kernels
+//
+
+template <
+    /// Element type for A matrix operand
+    typename ElementA,
+    /// Layout type for A matrix operand
+    typename LayoutA,
+    /// Complex elementwise transformation on A operand
+    ComplexTransform TransformA,
+    /// Access granularity of A matrix in units of elements
+    int kAlignmentA,
+    /// Element type for B matrix operand
+    typename ElementB,
+    /// Layout type for B matrix operand
+    typename LayoutB,
+    /// Complex elementwise transformation on B operand
+    ComplexTransform TransformB,
+    /// Access granularity of B matrix in units of elements
+    int kAlignmentB,
+    /// Element type for C and D matrix operands
+    typename ElementC,
+    /// Layout type for C and D matrix operands
+    typename LayoutC,
+    /// Element type for internal accumulation
+    typename ElementAccumulator,
+    /// Operator class tag
+    typename OperatorClass,
+    /// Tag indicating architecture to tune for
+    typename ArchTag,
+    /// Threadblock-level tile size (concept: GemmShape)
+    typename ThreadblockShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename WarpShape,
+    /// Warp-level tile size (concept: GemmShape)
+    typename InstructionShape,
+    /// Epilogue output operator
+    typename EpilogueOutputOp,
+    /// Threadblock-level swizzling operator
+    typename ThreadblockSwizzle,
+    /// Number of stages used in the pipelined mainloop
+    int Stages,
+    /// Whether the schedule of problems to visit has been precomputed
+    GroupScheduleMode GroupScheduleMode_,
+    /// Operation performed by GEMM
+    typename Operator,
+    /// Use zfill or predicate for out-of-bound cp.async
+    SharedMemoryClearOption SharedMemoryClear
+  >
+struct DefaultGemmGroupedPerGroupScale<
+  ElementA,
+  LayoutA,
+  TransformA,
+  kAlignmentA,
+  ElementB,
+  LayoutB,
+  TransformB,
+  kAlignmentB,
+  ElementC,
+  LayoutC,
+  ElementAccumulator,
+  OperatorClass,
+  ArchTag,
+  ThreadblockShape,
+  WarpShape,
+  InstructionShape,
+  EpilogueOutputOp,
+  ThreadblockSwizzle,
+  Stages,
+  GroupScheduleMode_,
+  Operator,
+  SharedMemoryClear,
+  layout::NoPermute, /*PermuteDLayout*/
+  typename platform::enable_if<cutlass::is_complex<ElementAccumulator>::value>::type
+> {
+
+  // If true, we must construct a 'transposed-and-exchanged' Mma operator.
+  static bool const kInternalTranspose = platform::is_same<LayoutC, layout::ColumnMajor>::value;
+
+  using MapArguments = kernel::detail::MapArguments<
+    ElementA,
+    LayoutA,
+    TransformA,
+    kAlignmentA,
+    ElementB,
+    LayoutB,
+    TransformB,
+    kAlignmentB,
+    LayoutC,
+    kInternalTranspose
+  >;
+
+  using DefaultGemmKernel = typename kernel::DefaultGemmComplex<
+    typename MapArguments::ElementA,
+    typename MapArguments::LayoutA,
+    typename MapArguments::ElementB,
+    typename MapArguments::LayoutB,
+    ElementC,
+    typename MapArguments::LayoutC,
+    ElementAccumulator,
+    OperatorClass,
+    ArchTag,
+    ThreadblockShape,
+    WarpShape,
+    InstructionShape,
+    EpilogueOutputOp,
+    ThreadblockSwizzle,
+    Stages,
+    MapArguments::kTransformA,
+    MapArguments::kTransformB,
+    Operator,
+    false
+  >::GemmKernel;
+
+  /// Define the kernel in terms of the default kernel
+  using GemmKernel = kernel::GemmGroupedPerGroupScale<
+    typename DefaultGemmKernel::Mma,
+    typename DefaultGemmKernel::Epilogue, 
+    ThreadblockSwizzle,
+    GroupScheduleMode_,
+    kInternalTranspose
+  >;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace kernel
+}  // namespace gemm
+}  // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

include/cutlass/gemm/kernel/ell_gemm.h

@@ -691,7 +691,7 @@ struct EllGemm<Mma_, Epilogue_, ThreadblockSwizzle_, SplitKSerial, false> {
         static int const kAlignmentA = Mma::IteratorA::AccessType::kElements;
         static int const kAlignmentB = Mma::IteratorB::AccessType::kElements;
         static int const kAlignmentC = Epilogue::OutputTileIterator::kElementsPerAccess;
-        constexpr bool is_double = (sizeof(Mma::IteratorA::Element) == 8);
+        constexpr bool is_double = (sizeof(typename Mma::IteratorA::Element) == 8);
         constexpr bool is_multiple_alignment =
           (kAlignmentA > 1) && (kAlignmentB > 1) && (kAlignmentC > 1);
         const bool is_specialized_blocksize =
@@ -699,11 +699,11 @@ struct EllGemm<Mma_, Epilogue_, ThreadblockSwizzle_, SplitKSerial, false> {
           && params.ell_blocksize >= Mma::Shape::kK;
         // Compute threadblock-scoped matrix multiply-add
         if ((is_double || is_multiple_alignment) && is_specialized_blocksize) {
-          mma.operator()<false, true>(
+          mma.template operator()<false, true>(
               gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators, ell_iterator);
         }
         else {
-          mma.operator()<false, false>(
+          mma.template operator()<false, false>(
               gemm_k_iterations, accumulators, iterator_A, iterator_B, accumulators, ell_iterator);
         }
       }

include/cutlass/gemm/kernel/gemm_grouped_per_group_scale.h

@@ -0,0 +1,261 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2024 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.
+ *
+ **************************************************************************************************/
+
+/*! \file
+    \brief Problem visitor for grouped GEMMs
+*/
+
+#pragma once
+
+#include "cutlass/cutlass.h"
+#include "cutlass/fast_math.h"
+#include "cutlass/gemm/gemm.h"
+#include "cutlass/matrix_coord.h"
+#include "cutlass/complex.h"
+#include "cutlass/semaphore.h"
+
+#include "cutlass/layout/matrix.h"
+#include "cutlass/trace.h"
+#include "cutlass/gemm/kernel/gemm_transpose_operands.h"
+#include "cutlass/gemm/kernel/gemm_grouped_problem_visitor.h"
+#include "cutlass/epilogue/thread/linear_combination.h"
+#include "cutlass/gemm/kernel/gemm_grouped.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace cutlass {
+namespace gemm {
+namespace kernel {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+template <
+  typename Mma_,                           ///! Threadblock-scoped matrix multiply-accumulate
+  typename Epilogue_,                      ///! Epilogue
+  typename ThreadblockSwizzle_,            ///! Threadblock swizzling function
+  GroupScheduleMode GroupScheduleMode_,    ///! Type of scheduling to perform
+  bool Transposed = false
+>
+struct GemmGroupedPerGroupScale : 
+  public GemmGrouped<Mma_, Epilogue_, ThreadblockSwizzle_, GroupScheduleMode_, Transposed> {
+
+  // Inherit constructors
+  using Base = GemmGrouped<Mma_, Epilogue_, ThreadblockSwizzle_, GroupScheduleMode_, Transposed>;
+
+  // Inherit type definitions
+  using typename Base::Mma;
+  using typename Base::Epilogue;
+  using typename Base::EpilogueOutputOp;
+  using typename Base::ThreadblockSwizzle;
+  using typename Base::Params;
+  using typename Base::SharedStorage;
+
+  // Explicitly inherit the kTransposed constant
+  static bool const kTransposed = Base::kTransposed;
+
+  /// Executes one GEMM
+  CUTLASS_DEVICE
+  void operator()(Params const &params, SharedStorage &shared_storage) {
+
+    //
+    // These types shadow the type-level definitions and support the ability to implement
+    // a 'transposed' GEMM that computes the transposed problems.
+    //
+    using ElementA = typename Mma::IteratorA::Element;
+    using LayoutA = typename Mma::IteratorA::Layout;
+    using ElementB = typename Mma::IteratorB::Element;
+    using LayoutB = typename Mma::IteratorB::Layout;
+    using ElementC = typename Epilogue::OutputTileIterator::Element;
+    using LayoutC = typename Epilogue::OutputTileIterator::Layout;
+
+    //
+    // Problem visitor.
+    //
+    typename Base::ProblemVisitor problem_visitor(
+      params.problem_visitor,
+      shared_storage.problem_visitor,
+      blockIdx.x);
+
+    // Outer 'persistent' loop to iterate over tiles
+    while (problem_visitor.next_tile()) {
+
+      GemmCoord problem_size  = problem_visitor.problem_size();
+      int32_t problem_idx     = problem_visitor.problem_index();
+      int32_t threadblock_idx = int32_t(problem_visitor.threadblock_idx());
+
+      GemmCoord grid_shape = problem_visitor.grid_shape(problem_size);
+
+      cutlass::gemm::GemmCoord threadblock_offset(
+        int(threadblock_idx / grid_shape.n()) * Mma::Shape::kM,
+        int(threadblock_idx % grid_shape.n()) * Mma::Shape::kN,
+        0);
+
+      // Load element pointers. Exchange pointers and strides if working on the transpose
+      ElementA *ptr_A = reinterpret_cast<ElementA *>((kTransposed ? params.ptr_B[problem_idx] : params.ptr_A[problem_idx]));
+      typename LayoutA::LongIndex ldm_A = (kTransposed ? params.ldb[problem_idx] : params.lda[problem_idx]);
+
+      ElementB *ptr_B = reinterpret_cast<ElementB *>((kTransposed ? params.ptr_A[problem_idx] : params.ptr_B[problem_idx]));
+      typename LayoutB::LongIndex ldm_B = (kTransposed ? params.lda[problem_idx] : params.ldb[problem_idx]);
+
+      // Compute initial location in logical coordinates
+      cutlass::MatrixCoord tb_offset_A{
+        threadblock_offset.m(),
+        0,
+      };
+
+      cutlass::MatrixCoord tb_offset_B{
+        0,
+        threadblock_offset.n()
+      };
+
+      // Compute position within threadblock
+      int thread_idx = threadIdx.x;
+
+      // Construct iterators to A and B operands
+      typename Mma::IteratorA iterator_A(
+        LayoutA(ldm_A),
+        ptr_A,
+        {problem_size.m(), problem_size.k()},
+        thread_idx,
+        tb_offset_A);
+
+      typename Mma::IteratorB iterator_B(
+        LayoutB(ldm_B),
+        ptr_B,
+        {problem_size.k(), problem_size.n()},
+        thread_idx,
+        tb_offset_B);
+
+      typename Mma::FragmentC accumulators;
+
+      accumulators.clear();
+      
+      // Broadcast the warp_id computed by lane 0 to ensure dependent code
+      // is compiled as warp-uniform.
+      int warp_idx = canonical_warp_idx_sync();
+
+      int lane_idx = threadIdx.x % 32;
+
+      //
+      // Matrix multiply phase
+      //
+
+      // Construct thread-scoped matrix multiply
+      Mma mma(shared_storage.kernel.main_loop, thread_idx, warp_idx, lane_idx);
+
+      // Compute threadblock-scoped matrix multiply-add
+      int gemm_k_iterations = (problem_size.k() + Mma::Shape::kK - 1) / Mma::Shape::kK;
+
+      // Wait for all threads to finish their epilogue phases from the previous tile.
+      __syncthreads();
+
+      // Compute threadblock-scoped matrix multiply-add
+      mma(
+        gemm_k_iterations, 
+        accumulators, 
+        iterator_A, 
+        iterator_B, 
+        accumulators);
+
+      //
+      // Epilogue
+      //
+
+      ElementC *ptr_C = params.ptr_C[problem_idx];
+      ElementC *ptr_D = params.ptr_D[problem_idx];
+
+      LayoutC layout_C(params.ldc[problem_idx]);
+      LayoutC layout_D(params.ldd[problem_idx]);
+
+      typename Epilogue::OutputTileIterator::Params params_C(layout_C);
+      typename Epilogue::OutputTileIterator::Params params_D(layout_D);
+
+      // Tile iterator loading from source tensor.
+      typename Epilogue::OutputTileIterator iterator_C(
+        params_C,
+        ptr_C,
+        problem_size.mn(),
+        thread_idx,
+        threadblock_offset.mn()
+      );
+
+      // Tile iterator writing to destination tensor.
+      typename Epilogue::OutputTileIterator iterator_D(
+        params_D,
+        ptr_D,
+        problem_size.mn(),
+        thread_idx,
+        threadblock_offset.mn()
+      );
+
+      Epilogue epilogue(
+        shared_storage.kernel.epilogue, 
+        thread_idx, 
+        warp_idx, 
+        lane_idx);
+
+      // The if branch is for the per-group scaling epilogue. The customized epilogue operator scales each gemm output by a scalar value.
+      // This branch is only enabled if EpilogueOutputOp is LinearCombination.
+      if constexpr (platform::is_same<EpilogueOutputOp,
+                              ::cutlass::epilogue::thread::LinearCombination<typename EpilogueOutputOp::ElementOutput,
+                                  EpilogueOutputOp::kCount, typename EpilogueOutputOp::ElementAccumulator,
+                                  typename EpilogueOutputOp::ElementCompute, EpilogueOutputOp::kScale,
+                                  EpilogueOutputOp::kRound>>::value)
+      {
+        EpilogueOutputOp output_op(params.output_op, problem_idx);
+        // Execute the epilogue operator to update the destination tensor.
+        epilogue(
+            output_op, 
+            iterator_D, 
+            accumulators, 
+            iterator_C); 
+      } else {
+        EpilogueOutputOp output_op(params.output_op);
+        // Execute the epilogue operator to update the destination tensor.
+        epilogue(
+            output_op, 
+            iterator_D, 
+            accumulators, 
+            iterator_C); 
+      }
+
+      // Next tile
+      problem_visitor.advance(gridDim.x);
+    }
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace kernel
+} // namespace gemm
+} // namespace cutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////

include/cutlass/gemm/kernel/gemm_grouped_problem_visitor.h

@@ -68,7 +68,7 @@ struct GemmGroupedProblemSizeHelper {
   CUTLASS_HOST_DEVICE
   static void possibly_transpose_problem(cutlass::gemm::GemmCoord& problem) {
     if (kTransposed) {
-      swap(problem.m(), problem.n());
+      cutlass::swap(problem.m(), problem.n());
     }
   }
 

include/cutlass/gemm/kernel/gemm_universal_with_visitor_streamk.h

@@ -437,7 +437,7 @@ protected:
 
     int m_begin = tile_work.tiled_coord.m() * Mma::Shape::kM;
     int m_end = params.block_mapping.problem_size.m();
-    return Mma::IteratorA(
+    return typename Mma::IteratorA(
         params.params_A,
         ptr_A,
         { m_end, tile_work.k_end },
@@ -466,7 +466,7 @@ protected:
 
     int n_begin = tile_work.tiled_coord.n() * Mma::Shape::kN;
     int n_end = params.block_mapping.problem_size.n();
-    return Mma::IteratorB(
+    return typename Mma::IteratorB(
         params.params_B,
         ptr_B,
         { tile_work.k_end, n_end },

include/cutlass/gemm/kernel/grouped_problem_visitor.h

@@ -66,10 +66,10 @@ struct BaseGroupedProblemVisitor {
     int32_t problem_idx;
     int32_t problem_start;
 
-    CUTLASS_DEVICE
+    CUTLASS_HOST_DEVICE
     ProblemInfo() : problem_idx(kNoPrefetchEntry), problem_start(kNoPrefetchEntry) {}
 
-    CUTLASS_DEVICE
+    CUTLASS_HOST_DEVICE
     ProblemInfo(int32_t problem_idx_, int32_t problem_start_) :
       problem_idx(problem_idx_), problem_start(problem_start_) {}
   };

include/cutlass/gemm/kernel/params_universal_base.h

@@ -182,7 +182,7 @@ struct UniversalParamsBase
       CUTLASS_TRACE_HOST("  Initialize " << workspace_bytes << " workspace bytes");
 
       cudaError_t result = cudaMemsetAsync(
-        semaphore,
+        static_cast<int *>(workspace),
         0,
         workspace_bytes,
         stream);

include/cutlass/gemm/kernel/rank_2k_grouped.h

@@ -479,14 +479,14 @@ public:
 
       // Construct iterators to A and B operands for Mma1
       typename Mma1::IteratorA iterator_A(
-        Mma1::IteratorA::Params(ldm_A),
+        typename Mma1::IteratorA::Params(ldm_A),
         ptr_A,
         {problem_size.m(), problem_size_k},
         thread_idx,
         tb_offset_MxK);
 
       typename Mma1::IteratorB iterator_BT(
-        Mma1::IteratorB::Params(ldm_B),
+        typename Mma1::IteratorB::Params(ldm_B),
         ptr_B,
         {problem_size_k, problem_size.n()},
         thread_idx,
@@ -494,14 +494,14 @@ public:
 
       // Construct iterators to A and B operands for Mma2
       typename Mma2::IteratorA iterator_B(
-        Mma2::IteratorA::Params(ldm_B),
+        typename Mma2::IteratorA::Params(ldm_B),
         ptr_B,
         {problem_size.m(), problem_size_k},
         thread_idx,
         tb_offset_MxK);
 
       typename Mma2::IteratorB iterator_AT(
-        Mma2::IteratorB::Params(ldm_A),
+        typename Mma2::IteratorB::Params(ldm_A),
         ptr_A,
         {problem_size_k, problem_size.n()},
         thread_idx,
@@ -560,7 +560,7 @@ public:
 
         // Tile iterator loading from source tensor.
         typename Epilogue::OutputTileIterator iterator_C(
-          Epilogue::OutputTileIterator::Params(params.ldc[problem_idx]),
+          typename Epilogue::OutputTileIterator::Params(params.ldc[problem_idx]),
           ptr_C,
           problem_size.mn(),
           thread_idx,
@@ -570,7 +570,7 @@ public:
 
         // Tile iterator writing to destination tensor.
         typename Epilogue::OutputTileIterator iterator_D(
-          Epilogue::OutputTileIterator::Params(params.ldd[problem_idx]),
+          typename Epilogue::OutputTileIterator::Params(params.ldd[problem_idx]),
           ptr_D,
           problem_size.mn(),
           thread_idx,
@@ -634,7 +634,7 @@ public:
 
       // Tile iterator loading from source tensor.
       typename Epilogue::OutputTileIterator iterator_C(
-        Epilogue::OutputTileIterator::Params(params.ldc[problem_idx]),
+        typename Epilogue::OutputTileIterator::Params(params.ldc[problem_idx]),
         ptr_C,
         problem_size.mn(),
         thread_idx,
@@ -644,7 +644,7 @@ public:
 
       // Tile iterator writing to destination tensor.
       typename Epilogue::OutputTileIterator iterator_D(
-        Epilogue::OutputTileIterator::Params(params.ldd[problem_idx]),
+        typename Epilogue::OutputTileIterator::Params(params.ldd[problem_idx]),
         ptr_D,
         problem_size.mn(),
         thread_idx,

include/cutlass/gemm/kernel/rank_2k_grouped_problem_visitor.h

@@ -357,7 +357,7 @@ struct Rank2KGroupedProblemVisitor : public GroupedProblemVisitor<
     int32_t macro_col = macro_id - (((macro_row+1) * macro_row)/2);
 
     if (kFillModeC == cutlass::FillMode::kUpper) {
-      swap(macro_row, macro_col);
+      cutlass::swap(macro_row, macro_col);
     }
 
     int32_t row = OffsetHelper::macro_row_to_row(macro_row, threadblock_id);

include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp

@@ -218,11 +218,6 @@ public:
     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, typename ProblemShape::UnderlyingProblemShape{}, 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, sm_count);
     workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
@@ -231,6 +226,11 @@ public:
     workspace_offset += CollectiveMainloop::get_workspace_size(problem_shapes, args.mainloop, sm_count);
     workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
 
+    void* scheduler_workspace = workspace_ptr + workspace_offset;
+    workspace_offset += TileScheduler::template get_workspace_size<typename ProblemShape::UnderlyingProblemShape, ElementAccumulator>(
+      args.scheduler, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups);
+    workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
+
     TileSchedulerParams scheduler;
     if constexpr (IsGroupedGemmKernel) {
       scheduler = TileScheduler::to_underlying_arguments(
@@ -276,10 +276,6 @@ public:
     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, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles);
-    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) {
@@ -294,6 +290,10 @@ public:
     workspace_size += CollectiveMainloop::get_workspace_size(args.problem_shape, args.mainloop, sm_count);
     workspace_size = round_nearest(workspace_size,  MinWorkspaceAlignment);
 
+    workspace_size += TileScheduler::template get_workspace_size<typename ProblemShape::UnderlyingProblemShape, ElementAccumulator>(
+      args.scheduler, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles);
+    workspace_size = round_nearest(workspace_size,  MinWorkspaceAlignment);
+
     return workspace_size;
   }
 
@@ -306,23 +306,25 @@ public:
     constexpr uint32_t NumEpilogueSubTiles = CollectiveEpilogue::get_store_pipe_increment(TileShape{});
     static constexpr uint32_t NumAccumulatorMtxs = 1;
 
-    status = TileScheduler::template initialize_workspace<typename ProblemShape::UnderlyingProblemShape, ElementAccumulator>(
-      args.scheduler, workspace_ptr + workspace_offset, stream, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles, NumAccumulatorMtxs, cuda_adapter);
-    workspace_offset += TileScheduler::template get_workspace_size<typename ProblemShape::UnderlyingProblemShape, ElementAccumulator>(
-      args.scheduler, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles);
+    status = CollectiveEpilogue::initialize_workspace(args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream, cuda_adapter);
+    workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue, args.hw_info.sm_count);
     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, cuda_adapter);
-    workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue, args.hw_info.sm_count);
-    workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
-
     status = CollectiveMainloop::initialize_workspace(args.problem_shape, args.mainloop, workspace_ptr + workspace_offset, stream, cuda_adapter);
     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;
+    }
 
+    status = TileScheduler::template initialize_workspace<typename ProblemShape::UnderlyingProblemShape, ElementAccumulator>(
+      args.scheduler, workspace_ptr + workspace_offset, stream, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles, NumAccumulatorMtxs, cuda_adapter);
+    workspace_offset += TileScheduler::template get_workspace_size<typename ProblemShape::UnderlyingProblemShape, ElementAccumulator>(
+      args.scheduler, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles);
+    workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
     if (status != Status::kSuccess) {
       return status;
     }
@@ -633,7 +635,7 @@ public:
         constexpr bool IsEpiLoad = true;
 
         if (work_tile_info.is_valid()) {
-          collective_epilogue.tensormaps_perform_update<IsEpiLoad>(
+          collective_epilogue.template tensormaps_perform_update<IsEpiLoad>(
             shared_storage.tensormaps.epilogue,
             params.epilogue,
             epi_load_tensormap,
@@ -644,7 +646,7 @@ public:
 
           // Converge before issuing tensormap fence release since fence is aligned
           __syncwarp();
-          collective_epilogue.tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue, epi_load_tensormap, 0);
+          collective_epilogue.template tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue, epi_load_tensormap, 0);
         }
 
         load_order_barrier.wait();
@@ -667,7 +669,7 @@ public:
             auto blk_coord = make_coord(m_coord, n_coord, _, l_coord);
 
             if (did_batch_change) {
-              collective_epilogue.tensormaps_fence_acquire<IsEpiLoad>(epi_load_tensormap);
+              collective_epilogue.template tensormaps_fence_acquire<IsEpiLoad>(epi_load_tensormap);
             }
 
             bool wait = work_tile_info.is_valid() && curr_batch != next_work_tile_info.L_idx;
@@ -697,7 +699,7 @@ public:
 
             // tensormap update
             {
-              collective_epilogue.tensormaps_perform_update<IsEpiLoad>(
+              collective_epilogue.template tensormaps_perform_update<IsEpiLoad>(
                 shared_storage.tensormaps.epilogue,
                 params.epilogue,
                 epi_load_tensormap,
@@ -708,7 +710,7 @@ public:
 
               // Converge before issuing tensormap fence release since fence is aligned
               __syncwarp();
-              collective_epilogue.tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue, epi_load_tensormap, 0);
+              collective_epilogue.template tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue, epi_load_tensormap, 0);
             }
           }
 
@@ -738,7 +740,7 @@ public:
       if (work_tile_info.is_valid()) {
 
         if (warp_idx_in_warp_group == 0) {
-          collective_epilogue.tensormaps_perform_update<IsEpiLoad>(
+          collective_epilogue.template tensormaps_perform_update<IsEpiLoad>(
             shared_storage.tensormaps.epilogue,
             params.epilogue,
             epi_store_tensormap,
@@ -749,8 +751,8 @@ public:
 
           // Converge before issuing tensormap fence release since fence is aligned
           __syncwarp();
-          collective_epilogue.tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue, 
-                                                                     epi_store_tensormap,
+          collective_epilogue.template tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue, 
+                                                                     epi_store_tensormap, 
                                                                      consumer_warp_group_idx);
         }
       }
@@ -805,7 +807,7 @@ public:
           params.scheduler, work_tile_info, accumulators, NumMmaWarpGroups, consumer_warp_group_idx);
 
         if (did_batch_change) {
-          collective_epilogue.tensormaps_fence_acquire<IsEpiLoad>(epi_store_tensormap);
+          collective_epilogue.template tensormaps_fence_acquire<IsEpiLoad>(epi_store_tensormap);
         }
 
         if (TileScheduler::compute_epilogue(work_tile_info, params.scheduler)) {
@@ -843,7 +845,7 @@ public:
             problem_shape_MNKL = append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), 1);
           }
           if (warp_idx_in_warp_group == 0) {
-            collective_epilogue.tensormaps_perform_update<IsEpiLoad>(
+            collective_epilogue.template tensormaps_perform_update<IsEpiLoad>(
               shared_storage.tensormaps.epilogue,
               params.epilogue,
               epi_store_tensormap,
@@ -854,7 +856,7 @@ public:
 
             // Converge before issuing tensormap fence release since fence is aligned
             __syncwarp();
-            collective_epilogue.tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue,
+            collective_epilogue.template tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue, 
                                                                        epi_store_tensormap,
                                                                        consumer_warp_group_idx);
           }

include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp

@@ -226,11 +226,6 @@ public:
     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, typename ProblemShape::UnderlyingProblemShape{}, 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, sm_count);
     workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
@@ -239,6 +234,11 @@ public:
     workspace_offset += CollectiveMainloop::get_workspace_size(problem_shapes, args.mainloop, sm_count);
     workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
 
+    void* scheduler_workspace = workspace_ptr + workspace_offset;
+    workspace_offset += TileScheduler::template get_workspace_size<typename ProblemShape::UnderlyingProblemShape, ElementAccumulator>(
+      args.scheduler, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups);
+    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.
@@ -288,10 +288,6 @@ public:
     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, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles);
-    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) {
@@ -306,6 +302,10 @@ public:
     workspace_size += CollectiveMainloop::get_workspace_size(args.problem_shape, args.mainloop, sm_count);
     workspace_size = round_nearest(workspace_size,  MinWorkspaceAlignment);
 
+    workspace_size += TileScheduler::template get_workspace_size<typename ProblemShape::UnderlyingProblemShape, ElementAccumulator>(
+      args.scheduler, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles);
+    workspace_size = round_nearest(workspace_size,  MinWorkspaceAlignment);
+
     return workspace_size;
   }
 
@@ -318,6 +318,20 @@ public:
     constexpr uint32_t NumEpilogueSubTiles = CollectiveEpilogue::get_store_pipe_increment(TileShape{});
     static constexpr uint32_t NumAccumulatorMtxs = 1;
 
+    status = CollectiveEpilogue::initialize_workspace(args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream, cuda_adapter);
+    workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue, args.hw_info.sm_count);
+    workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    status = CollectiveMainloop::initialize_workspace(args.problem_shape, args.mainloop, workspace_ptr + workspace_offset, stream, cuda_adapter);
+    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;
+    }
+
     status = TileScheduler::template initialize_workspace<typename ProblemShape::UnderlyingProblemShape, ElementAccumulator>(
       args.scheduler, workspace_ptr + workspace_offset, stream, typename ProblemShape::UnderlyingProblemShape{}, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles, NumAccumulatorMtxs, cuda_adapter);
     workspace_offset += TileScheduler::template get_workspace_size<typename ProblemShape::UnderlyingProblemShape, ElementAccumulator>(
@@ -326,19 +340,6 @@ public:
     if (status != Status::kSuccess) {
       return status;
     }
-
-    status = CollectiveEpilogue::initialize_workspace(args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream, cuda_adapter);
-    workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue, args.hw_info.sm_count);
-    workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
-
-    status = CollectiveMainloop::initialize_workspace(args.problem_shape, args.mainloop, workspace_ptr + workspace_offset, stream, cuda_adapter);
-    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;
   }
 
@@ -666,7 +667,7 @@ public:
         constexpr bool IsEpiLoad = true;
 
         if (work_tile_info.is_valid()) {
-          collective_epilogue.tensormaps_perform_update<IsEpiLoad>(
+          collective_epilogue.template tensormaps_perform_update<IsEpiLoad>(
             shared_storage.tensormaps.epilogue,
             params.epilogue,
             epi_load_tensormap,
@@ -677,7 +678,7 @@ public:
 
           // Converge before issuing tensormap fence release since fence is aligned
           __syncwarp();
-          collective_epilogue.tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue, epi_load_tensormap, 0);
+          collective_epilogue.template tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue, epi_load_tensormap, 0);
         }
 
         load_order_barrier.wait();
@@ -700,7 +701,7 @@ public:
             auto blk_coord = make_coord(m_coord, n_coord, _, l_coord);
 
             if (did_batch_change) {
-              collective_epilogue.tensormaps_fence_acquire<IsEpiLoad>(epi_load_tensormap);
+              collective_epilogue.template tensormaps_fence_acquire<IsEpiLoad>(epi_load_tensormap);
             }
 
             bool wait = work_tile_info.is_valid() && curr_batch != next_work_tile_info.L_idx;
@@ -730,7 +731,7 @@ public:
 
             // tensormap update
             {
-              collective_epilogue.tensormaps_perform_update<IsEpiLoad>(
+              collective_epilogue.template tensormaps_perform_update<IsEpiLoad>(
                 shared_storage.tensormaps.epilogue,
                 params.epilogue,
                 epi_load_tensormap,
@@ -741,7 +742,7 @@ public:
 
               // Converge before issuing tensormap fence release since fence is aligned
               __syncwarp();
-              collective_epilogue.tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue, epi_load_tensormap, 0);
+              collective_epilogue.template tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue, epi_load_tensormap, 0);
             }
           }
 
@@ -771,7 +772,7 @@ public:
       if (work_tile_info.is_valid()) {
 
         if (warp_idx_in_warp_group == 0) {
-          collective_epilogue.tensormaps_perform_update<IsEpiLoad>(
+          collective_epilogue.template tensormaps_perform_update<IsEpiLoad>(
             shared_storage.tensormaps.epilogue,
             params.epilogue,
             epi_store_tensormap,
@@ -782,7 +783,7 @@ public:
 
           // Converge before issuing tensormap fence release since fence is aligned
           __syncwarp();
-          collective_epilogue.tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue,
+          collective_epilogue.template tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue,
                                                                      epi_store_tensormap,
                                                                      consumer_warp_group_idx);
         }
@@ -844,7 +845,7 @@ public:
           params.scheduler, work_tile_info, accumulators, NumMmaWarpGroups, consumer_warp_group_idx);
 
         if (did_batch_change) {
-          collective_epilogue.tensormaps_fence_acquire<IsEpiLoad>(epi_store_tensormap);
+          collective_epilogue.template tensormaps_fence_acquire<IsEpiLoad>(epi_store_tensormap);
         }
 
         if (TileScheduler::compute_epilogue(work_tile_info, params.scheduler)) {
@@ -897,7 +898,7 @@ public:
             problem_shape_MNKL = append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), 1);
           }
           if (warp_idx_in_warp_group == 0) {
-            collective_epilogue.tensormaps_perform_update<IsEpiLoad>(
+            collective_epilogue.template tensormaps_perform_update<IsEpiLoad>(
               shared_storage.tensormaps.epilogue,
               params.epilogue,
               epi_store_tensormap,
@@ -908,7 +909,7 @@ public:
 
             // Converge before issuing tensormap fence release since fence is aligned
             __syncwarp();
-            collective_epilogue.tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue,
+            collective_epilogue.template tensormaps_cp_fence_release<IsEpiLoad>(shared_storage.tensormaps.epilogue,
                                                                        epi_store_tensormap,
                                                                        consumer_warp_group_idx);
           }

include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_cooperative.hpp

@@ -51,8 +51,6 @@
 
 namespace cutlass::gemm::kernel {
 
-///////////////////////////////////////////////////////////////////////////////
-
 template <
   class ProblemShape_,
   class CollectiveMainloop_,
@@ -107,7 +105,6 @@ public:
                                           TileShape,
                                           ClusterShape
                                           >::Scheduler;
-
   using TileSchedulerArguments = typename TileScheduler::Arguments;
   using TileSchedulerParams = typename TileScheduler::Params;
   
@@ -122,7 +119,8 @@ public:
   static constexpr uint32_t NumMmaWarpGroups = NumMMAThreads / NumThreadsPerWarpGroup;
   static constexpr uint32_t MaxThreadsPerBlock = NumMMAThreads + (NumLoadWarpGroups * NumThreadsPerWarpGroup);
   static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
-
+  static constexpr uint32_t NumFixupBarriers = NumMmaWarpGroups;
+  
   /// Register requirement for Load and Math WGs
   static constexpr uint32_t LoadRegisterRequirement = 40;
   static constexpr uint32_t MmaRegisterRequirement = 232;
@@ -207,22 +205,23 @@ public:
     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<ProblemShape, ElementAccumulator>(
-      args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups);
-    workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
-
     void* epilogue_workspace = workspace_ptr + workspace_offset;
     workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
     workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
 
+    void* scheduler_workspace = workspace_ptr + workspace_offset;
+    workspace_offset += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
+      args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups);
+    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);
+      problem_shape_MNKL, TileShape{}, ClusterShape{}, hw_info, args.scheduler, scheduler_workspace, NumEpilogueSubTiles
+      );
 
     return {
       args.mode,
@@ -254,13 +253,12 @@ public:
     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, 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);
 
+    workspace_size += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
+      args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles);
+    workspace_size = round_nearest(workspace_size,  MinWorkspaceAlignment);
     return workspace_size;
   }
 
@@ -273,17 +271,17 @@ public:
     constexpr uint32_t NumEpilogueSubTiles = CollectiveEpilogue::get_store_pipe_increment(TileShape{});
     static constexpr uint32_t NumAccumulatorMtxs = 1;
 
-    status = TileScheduler::template initialize_workspace<ProblemShape, ElementAccumulator>(
-      args.scheduler, workspace_ptr + workspace_offset, stream, args.problem_shape, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles, NumAccumulatorMtxs, cuda_adapter);
-    workspace_offset += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
-      args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles);
+    status = CollectiveEpilogue::initialize_workspace(args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream, cuda_adapter);
+    workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
     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, cuda_adapter);
-    workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
+    status = TileScheduler::template initialize_workspace<ProblemShape, ElementAccumulator>(
+      args.scheduler, workspace_ptr + workspace_offset, stream, args.problem_shape, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles, NumAccumulatorMtxs, cuda_adapter);
+    workspace_offset += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
+      args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles);
     workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
     if (status != Status::kSuccess) {
       return status;
@@ -314,6 +312,7 @@ public:
   operator()(Params const& params, char* smem_buf) {
     using namespace cute;
     using X = Underscore;
+
 #if defined(__CUDA_ARCH_FEAT_SM90_ALL)
 #  define ENABLE_SM90_KERNEL_LEVEL 1
 #endif
@@ -487,7 +486,6 @@ public:
           // 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.load(
@@ -581,11 +579,10 @@ public:
         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)
+        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,

include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp

@@ -105,14 +105,24 @@ public:
   static_assert(!cute::is_same_v<TileScheduler_, StreamKScheduler>, "Ping-pong kernel does not currently support stream-K scheduler.");
   using TileSchedulerTag = TileScheduler_;
   using TileScheduler = typename detail::TileSchedulerSelector<
-    TileScheduler_, ArchTag, TileShape, ClusterShape>::Scheduler;
+                                          TileSchedulerTag, 
+                                          ArchTag, 
+                                          TileShape,
+                                          ClusterShape
+                                          >::Scheduler;
   using TileSchedulerArguments = typename TileScheduler::Arguments;
   using TileSchedulerParams = typename TileScheduler::Params;
-
+  // Warp specialization thread count per threadblock
+  static constexpr uint32_t NumMainloopLoadThreads = NumThreadsPerWarp;      // 1 warp
+  static constexpr uint32_t NumEpilogueLoadThreads = NumThreadsPerWarp;      // 1 warp for C
   static constexpr uint32_t NumLoadWarpGroups = 1;
   static constexpr uint32_t NumMmaWarpGroups = 2;
-  static constexpr uint32_t MaxThreadsPerBlock = CUTE_STATIC_V(size(TiledMma{})) + (NumMmaWarpGroups * NumThreadsPerWarpGroup);
+  static constexpr uint32_t NumMMAThreads = size(TiledMma{});                 // 4 warp 
+  static constexpr uint32_t MaxThreadsPerBlock = NumMMAThreads * NumMmaWarpGroups + (NumLoadWarpGroups * NumThreadsPerWarpGroup);
   static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
+  
+  static_assert(NumMMAThreads == 128, "Pingpong kernel must have TiledMMA operating using 128 threads.");
+  static_assert(MaxThreadsPerBlock == 384, "Pingpong kernel must have 384 threads in total.");
 
   /// Register requirement for Load and Math WGs
   static constexpr uint32_t LoadRegisterRequirement = 40;
@@ -142,7 +152,7 @@ public:
       alignas(16) MathWarpGroupOrderBarrierStorage math_wg_order;
       alignas(16) typename LoadWarpOrderBarrier::SharedStorage load_order;
     } pipelines;
-
+    
     struct TensorStorage : cute::aligned_struct<128, _1> {
       using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage;
       using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage;
@@ -208,16 +218,17 @@ public:
     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<ProblemShape, ElementAccumulator>(
-      args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups);
-    workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
-
     void* epilogue_workspace = workspace_ptr + workspace_offset;
     workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
     workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
 
+    void* scheduler_workspace = workspace_ptr + workspace_offset;
+    workspace_offset += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
+      args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups);
+    workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
+
     void* mainloop_workspace = nullptr;
+    constexpr uint32_t NumEpilogueSubTiles = CollectiveEpilogue::get_store_pipe_increment(TileShape{});
 
     return {
       args.mode,
@@ -225,7 +236,9 @@ 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)
+      TileScheduler::to_underlying_arguments(
+        problem_shape_MNKL, TileShape{}, ClusterShape{}, hw_info, args.scheduler, scheduler_workspace, NumEpilogueSubTiles
+      )
     };
   }
 
@@ -247,13 +260,14 @@ public:
   static size_t
   get_workspace_size(Arguments const& args) {
     size_t workspace_size = 0;
-    workspace_size += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
-      args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups);
-    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);
 
+    workspace_size += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
+      args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups);
+    workspace_size = round_nearest(workspace_size,  MinWorkspaceAlignment);
+
     return workspace_size;
   }
 
@@ -266,17 +280,17 @@ public:
     static constexpr uint32_t NumEpilogueSubTiles = 1;
     static constexpr uint32_t NumAccumulatorMtxs = 1;
 
-    status = TileScheduler::template initialize_workspace<ProblemShape, ElementAccumulator>(
-      args.scheduler, workspace_ptr + workspace_offset, stream, args.problem_shape, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles, NumAccumulatorMtxs, cuda_adapter);
-    workspace_offset += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
-      args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups);
+    status = CollectiveEpilogue::initialize_workspace(args.problem_shape, args.epilogue, workspace_ptr + workspace_offset, stream, cuda_adapter);
+    workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
     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, cuda_adapter);
-    workspace_offset += CollectiveEpilogue::get_workspace_size(args.problem_shape, args.epilogue);
+    status = TileScheduler::template initialize_workspace<ProblemShape, ElementAccumulator>(
+      args.scheduler, workspace_ptr + workspace_offset, stream, args.problem_shape, args.hw_info, NumMmaWarpGroups, NumEpilogueSubTiles, NumAccumulatorMtxs, cuda_adapter);
+    workspace_offset += TileScheduler::template get_workspace_size<ProblemShape, ElementAccumulator>(
+      args.scheduler, args.problem_shape, args.hw_info, NumMmaWarpGroups);
     workspace_offset = round_nearest(workspace_offset,  MinWorkspaceAlignment);
     if (status != Status::kSuccess) {
       return status;
@@ -308,9 +322,12 @@ public:
     using namespace cute;
     using X = Underscore;
 
+#if defined(__CUDA_ARCH_FEAT_SM90_ALL)
+#  define ENABLE_SM90_KERNEL_LEVEL 1
+#endif
 // Any Tensor Op MMA Atom in the WGMMA ISA is arch conditional to sm90a.
-#if ! defined(__CUDA_ARCH_FEAT_SM90_ALL)
-    printf("ERROR : Arch conditional MMA instruction used without targeting sm90a compute capability. Aborting.\n");
+#if ! defined(ENABLE_SM90_KERNEL_LEVEL)
+    printf("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
 #else
 
     // Preconditions
@@ -350,6 +367,7 @@ public:
       CollectiveEpilogue::prefetch_tma_descriptors(params.epilogue);
     }
 
+
     // Mainloop Load pipeline
     using MainloopPipeline = typename CollectiveMainloop::MainloopPipeline;
     typename MainloopPipeline::Params mainloop_pipeline_params;
@@ -450,8 +468,8 @@ public:
     auto d_tile_count = CollectiveEpilogue::get_store_pipe_increment(blk_shape);
 
     TileScheduler scheduler{params.scheduler};
-
     if (warp_group_role == WarpGroupRole::Consumer1) {
+
       // Advance 2nd Math WG to the next work tile for the startup
       scheduler.advance_to_next_work();
       // Advance 2nd Math WG pipeline states to the end of 1st Math WG
@@ -466,7 +484,7 @@ public:
 
     if (warp_group_role == WarpGroupRole::Producer) {
       cutlass::arch::warpgroup_reg_dealloc<LoadRegisterRequirement>();
-
+    
       // Mainloop Producer Warp
       if (producer_warp_role == ProducerWarpRole::Mainloop) {
         // Ensure that the prefetched kernel does not touch
@@ -546,6 +564,7 @@ public:
 
         // 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
 
@@ -564,7 +583,7 @@ public:
         return;
       }
       #endif
-
+      
       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));

include/cutlass/gemm/kernel/sm90_tile_scheduler.hpp

@@ -29,8 +29,8 @@
  *
  **************************************************************************************************/
 #pragma once
-#include "cutlass/gemm/kernel/static_tile_scheduler.hpp"
 
+#include "cutlass/gemm/kernel/static_tile_scheduler.hpp"
 
 namespace cutlass::gemm::kernel::detail {
 

include/cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp

@@ -337,12 +337,16 @@ public:
     uint64_t blk_per_grid_dim = divmod_cluster_shape_minor.divide(linear_idx - group_info.start_linear_idx);
     divmod_cluster_shape_major(cluster_id, cluster_major_offset, blk_per_grid_dim);
 
-    auto [cta_m_in_cluster, cta_n_in_cluster, _] = cute::block_id_in_cluster();
+    // With static schedulers, we launch grid such that all cluster are linear (1-D) order, i.e., 
+    // there can only be one cluster in the minor dimension. get_grid_shape() in scheduler params
+    // put cluster_shape.m/n() as the minor dimension based on raster order AlongN/M resp.
+    // Therefore, the offset of a CTA (inside a cluster) in the minor dimension can be directly be 
+    // inferred by the blockIdx along the minor dimension.
     if (raster_order == RasterOrder::AlongN) {
-      cluster_minor_offset = cta_m_in_cluster;
+      cluster_minor_offset = blockIdx.x;
     }
     else {
-      cluster_minor_offset = cta_n_in_cluster;
+      cluster_minor_offset = blockIdx.y;
     }
 
     uint64_t cluster_idx_minor, cluster_idx_major;

include/cutlass/gemm/kernel/sm90_tile_scheduler_stream_k.hpp

@@ -58,7 +58,9 @@ private:
   using UnderlyingArguments = typename UnderlyingScheduler::Arguments;
   using UnderlyingParams = typename UnderlyingScheduler::Params;
 
+  dim3 block_id_in_cluster_;
   uint64_t current_work_linear_idx_ = 0;
+  uint32_t unit_iter_start_ = 0;
 
 public:
 
@@ -240,25 +242,26 @@ public:
   CUTLASS_HOST_DEVICE
   PersistentTileSchedulerSm90StreamK() { };
 
-  CUTLASS_HOST_DEVICE
-  PersistentTileSchedulerSm90StreamK(Params const& params_) : scheduler_params(params_) {
+  CUTLASS_DEVICE
+  PersistentTileSchedulerSm90StreamK(Params const& params_) : scheduler_params(params_), block_id_in_cluster_(cute::block_id_in_cluster()) {
     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);
     }
+
   }
 
   CUTLASS_DEVICE
   WorkTileInfo
-  get_current_work() const {
-    return get_current_work_for_linear_idx(current_work_linear_idx_, scheduler_params);
+  get_current_work() {
+    return get_current_work_for_linear_idx(unit_iter_start_, current_work_linear_idx_, block_id_in_cluster_, scheduler_params);
   }
 
   CUTLASS_DEVICE
   static WorkTileInfo
-  get_current_work_for_linear_idx(uint64_t linear_idx, Params const& params) {
+  get_current_work_for_linear_idx(uint32_t &unit_iter_start, uint64_t linear_idx, dim3 block_id_in_cluster, Params const& params) {
     // The maximum number of work units is units_per_problem_ * splits_.
     // The multiplication by splits_ is used for handling split-K, in which
     // units_per_problem_ is equal to the total number of output tiles. To account
@@ -271,7 +274,7 @@ public:
     }
 
     WorkTileInfo work_tile_info;
-    assign_work(params, linear_idx, work_tile_info);
+    assign_work(params, linear_idx, block_id_in_cluster, work_tile_info, unit_iter_start);
     return work_tile_info;
   }
 
@@ -283,13 +286,15 @@ public:
   bool
   continue_current_work(WorkTileInfo& work_tile_info) const {
     return continue_current_work_for_linear_idx(
-      current_work_linear_idx_, work_tile_info, scheduler_params);
+      current_work_linear_idx_, unit_iter_start_, block_id_in_cluster_, work_tile_info, scheduler_params);
   }
 
   CUTLASS_DEVICE
   static bool
   continue_current_work_for_linear_idx(
     uint64_t linear_idx,
+    uint32_t unit_iter_start,
+    dim3 block_id_in_cluster,
     WorkTileInfo& work_tile_info,
     Params const& params) {
 
@@ -298,7 +303,7 @@ public:
     if (work_tile_info.k_tile_remaining == 0) {
       return false;
     }
-    assign_work(params, linear_idx, work_tile_info);
+    fast_assign_work(unit_iter_start, params, linear_idx, block_id_in_cluster, work_tile_info);
     return work_tile_info.is_valid();
   }
 
@@ -316,9 +321,11 @@ public:
       return false;
     }
     return not get_current_work_for_linear_idx(
+        unit_iter_start_,
         current_work_linear_idx_ + (
           uint64_t(gridDim.x) * uint64_t(gridDim.y) * uint64_t(gridDim.z) * uint64_t(advance_count)
           ),
+        block_id_in_cluster_,
         scheduler_params
     ).is_valid();
   }
@@ -420,22 +427,24 @@ public:
     uint64_t reduction_tile_idx = tile_idx;
     uint64_t num_peers = 0;
     uint64_t reduction_peer_offset = 0;
-    if (params.requires_separate_reduction()) {
+    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.
-      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 [first_peer_id, my_peer_id, last_peer_id] = tile_peer_range(params, tile_idx, work_tile_info);
+      auto peer_id_in_output_tile = my_peer_id - first_peer_id;
       num_peers = last_peer_id - first_peer_id + 1;
-      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{});
+      reduction_tile_idx = tile_idx * Params::max_peers_per_tile(params.sk_units_, params.sk_tiles_);
+      reduction_peer_offset = peer_id_in_output_tile * cute::size<0>(TileShape{}) * cute::size<1>(TileShape{}) * num_accumulator_mtxs;
     }
 
     // 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.
-    uint64_t reduction_offset =
-      (static_cast<uint64_t>(cute::size<0>(TileShape{})) * static_cast<uint64_t>(cute::size<1>(TileShape{})) * reduction_tile_idx * num_accumulator_mtxs) +
-      reduction_peer_offset +
+    uint64_t reduction_offset_base = (static_cast<uint64_t>(cute::size<0>(TileShape{})) * static_cast<uint64_t>(cute::size<1>(TileShape{})) * reduction_tile_idx * num_accumulator_mtxs) +
       (static_cast<uint64_t>(size(accumulators)) * barrier_idx * BarrierManager::ThreadCount);
+    uint64_t reduction_offset = reduction_offset_base + reduction_peer_offset;
 
     ElementAccumulator* group_reduction_workspace = reinterpret_cast<ElementAccumulator*>(params.reduction_workspace_) + reduction_offset;
 
@@ -457,7 +466,9 @@ public:
     if (params.divmod_splits_.divisor > 1) {
       reduction_tiles = params.units_per_problem_;
     }
-    else if (params.requires_separate_reduction()) {
+    else if (
+      params.requires_separate_reduction()
+      ) {
       reduction_tiles = params.sk_tiles_ * Params::max_peers_per_tile(params.sk_units_, params.sk_tiles_);
     }
     else {
@@ -470,29 +481,17 @@ public:
       reinterpret_cast<uint8_t*>(params.reduction_workspace_) + reduction_workspace_size);
 
     if (work_tile_info.is_reduction_unit()) {
-      plus<AccumulatorArrayT> add_fragments;
-      uint64_t peer_offset = size(accumulators) * num_barriers * BarrierManager::ThreadCount;
-
       // Wait until the peers collaborating on this output tile have all written
       // their accumulators to workspace.
       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 (uint64_t 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);
-      }
+      separate_reduction<FrgTensorC, BarrierManager>(accumulators, num_barriers, group_reduction_workspace, barrier_group_thread_idx, num_peers, num_accumulator_mtxs);
     }
     else if (!compute_epilogue(work_tile_info, params)) {
-      if (params.requires_separate_reduction() || work_tile_info.K_idx == 0) {
+      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);
@@ -513,12 +512,16 @@ public:
       BarrierManager::arrive_inc(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, increment);
     }
     else {
-      if (params.reduction_mode_ == ReductionMode::Deterministic) {
+      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
+        // Wait until the first split has stored its accumulators
         BarrierManager::wait_lt(barrier_idx, lock_workspace, barrier_group_thread_idx, lock_idx, 1);
       }
 
@@ -528,6 +531,36 @@ public:
     }
   }
 
+  template <class FrgTensorC, class BarrierManager>
+  CUTLASS_DEVICE
+  static void
+  separate_reduction(
+      FrgTensorC& accumulators,
+      uint32_t num_barriers,
+      typename FrgTensorC::value_type* reduction_workspace,
+      uint32_t thread_idx,
+      uint64_t num_peers,
+      uint32_t num_accumulator_mtxs) {
+    using AccumulatorArrayT = Array<typename FrgTensorC::value_type, size(FrgTensorC{})>;
+    using BlockStripedReduceT = BlockStripedReduce<BarrierManager::ThreadCount, AccumulatorArrayT>;
+
+    AccumulatorArrayT* accumulator_array = reinterpret_cast<AccumulatorArrayT*>(accumulators.data());
+
+    plus<AccumulatorArrayT> add_fragments;
+    uint64_t peer_offset = cute::size<0>(TileShape{}) * cute::size<1>(TileShape{}) * num_accumulator_mtxs;
+
+    for (uint64_t i = 0; i < num_peers; ++i) {
+      // Load peer fragment
+      AccumulatorArrayT addend_fragment;
+      auto peer_reduction_workspace = reinterpret_cast<AccumulatorArrayT*>(reduction_workspace + (i * peer_offset));
+
+      BlockStripedReduceT::load(addend_fragment, peer_reduction_workspace, thread_idx);
+
+      // Add peer fragment
+      *accumulator_array = add_fragments(*accumulator_array, addend_fragment);
+    }
+  }
+
   // Returns whether the block assigned this work should compute the epilogue for the corresponding
   // output tile. For the case of stream-K, this should only occur if the work is marked as the final split.
   CUTLASS_HOST_DEVICE
@@ -587,6 +620,7 @@ public:
       args.max_swizzle_size,
       args.raster_order,
       args.decomposition_mode,
+      args.reduction_mode,
       mma_warp_groups,
       sizeof_bits<BarrierType>::value,
       sizeof_bits<ElementAccumulator>::value,
@@ -627,6 +661,7 @@ public:
       args.max_swizzle_size,
       args.raster_order,
       args.decomposition_mode,
+      args.reduction_mode,
       mma_warp_groups,
       sizeof_bits<BarrierType>::value,
       sizeof_bits<ElementAccumulator>::value,
@@ -668,224 +703,235 @@ public:
     return get_current_work();
   }
 
-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.
+  // Given raster order and current work tile linear index, reset cta m and n index in the cluster.
   CUTLASS_DEVICE
-  static void
-  assign_work(
+  static dim3
+  get_current_work_cta_m_n_in_cluster(
+    Params const& params,
+    uint64_t linear_idx,
+    dim3 block_id_in_cluster) {
+    auto [cta_m_in_cluster_, cta_n_in_cluster_, _] = block_id_in_cluster;
+    uint64_t cta_m_in_cluster = static_cast<uint64_t>(cta_m_in_cluster_);
+    uint64_t cta_n_in_cluster = static_cast<uint64_t>(cta_n_in_cluster_);
+    return {static_cast<uint32_t>(cta_m_in_cluster), static_cast<uint32_t>(cta_n_in_cluster), _};
+  }
+
+private:
+
+  CUTLASS_DEVICE
+  static uint32_t
+  get_current_work_iter_start_possible_update_work_tile_k_remaining(
     Params const& params,
     uint64_t linear_idx,
     WorkTileInfo& work_tile_info) {
+    // In the CUTLASS 2.x implementation of stream K, stream-K work is assigned to each stream-K
+    // threadblock individually. For the most part, the set of K iterations corresponding to stream-K
+    // work was divided amongst stream-K threadblocks, and a threadblock determined which tile
+    // it would compute a (potentially-partial) output tile for based on the space of k iterations
+    // assigned to it. This often results in stream-K threadblocks processing tiles with different
+    // offsets in the K dimension from one another. This can reduce locality, but is lmitied to the
+    // (generally few) waves of threadblocks assigned to compute stream-K work.
+    //
+    // With the introduction of threadblock clusters, there is additional benefit to maintaining
+    // locality in the K dimension: shared portions of operands can be multicasted to threadblocks
+    // within a cluster. Thus, we would like to ensure that the assignment of stream-K work to
+    // threadblocks respects the ability to perform multicasting.
+    //
+    // To do so, we divide up the linearized stream-K units into clusters and share the same K
+    // offsets for work within clusters.
+    uint64_t cluster_linear_work_idx = params.div_cluster_size(linear_idx);
 
-    auto [cta_m_in_cluster_, cta_n_in_cluster_, _] = cute::block_id_in_cluster();
-    uint64_t cta_m_in_cluster = static_cast<uint64_t>(cta_m_in_cluster_);
-    uint64_t cta_n_in_cluster = static_cast<uint64_t>(cta_n_in_cluster_);
-    uint64_t output_tile_id = linear_idx;
-    if (linear_idx >= params.units_per_problem_ * params.divmod_splits_.divisor) {
-      // Separate-reduction work
-      auto cluster_size = params.get_cluster_size();
-      // Divide up the linearized separate reduction units into clusters
-      uint64_t 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;
+    uint64_t group_idx;
+    params.divmod_sk_groups_(cluster_linear_work_idx, group_idx, cluster_linear_work_idx);
 
-      work_tile_info.setup_separate_reduction(epi_subtile_idx);
+    // Determine whether we are in a "big group" that will process an additional
+    // stream-K cluster tile.
+    uint64_t sk_cluster_tiles = params.div_cluster_size(params.sk_tiles_);
+    uint64_t sk_cluster_tiles_in_group = params.divmod_sk_groups_.divide(sk_cluster_tiles);
+    if (group_idx < params.big_groups_) {
+      ++sk_cluster_tiles_in_group;
     }
-    else if (linear_idx >= params.sk_units_ && params.divmod_splits_.divisor == 1) {
-      // Data-parallel work
-      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;
+
+    // Determine whether we are in a "big unit" within the group, that will process
+    // an additional K chunk in the group.
+    uint64_t sk_tiles_in_group = sk_cluster_tiles_in_group * params.get_cluster_size();
+    uint64_t k_tiles_in_group = sk_tiles_in_group * params.divmod_tiles_per_output_tile_.divisor;
+    uint64_t k_tiles_per_unit_in_group = params.divmod_sk_units_per_group_.divide(k_tiles_in_group);
+    uint64_t 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);
+
+    bool is_split_k = params.divmod_splits_.divisor > 1;
+    uint64_t big_unit_cmp_lhs = is_split_k ? split : cluster_linear_work_idx;
+    uint64_t big_unit_cmp_rhs = is_split_k ? params.big_units_ : big_units_in_group;
+    uint64_t linear_idx_mult = is_split_k ? params.divmod_tiles_per_output_tile_.divisor : k_tiles_per_unit_in_group;
+    uint64_t k_tiles_per_split = is_split_k ? params.divmod_k_tiles_per_sk_unit_.divisor : 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) +
+                               (k_tiles_per_split * split);
+
+    // Adjust the starting position and number of k iterations for "big units," which
+    // 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_lhs;
+      ++k_tiles_in_my_split;
     }
     else {
-      // In the CUTLASS 2.x implementation of stream K, stream-K work is assigned to each stream-K
-      // threadblock individually. For the most part, the set of K iterations corresponding to stream-K
-      // work was divided amongst stream-K threadblocks, and a threadblock determined which tile
-      // it would compute a (potentially-partial) output tile for based on the space of k iterations
-      // assigned to it. This often results in stream-K threadblocks processing tiles with different
-      // offsets in the K dimension from one another. This can reduce locality, but is lmitied to the
-      // (generally few) waves of threadblocks assigned to compute stream-K work.
-      //
-      // With the introduction of threadblock clusters, there is additional benefit to maintaining
-      // locality in the K dimension: shared portions of operands can be multicasted to threadblocks
-      // within a cluster. Thus, we would like to ensure that the assignment of stream-K work to
-      // threadblocks respects the ability to perform multicasting.
-      //
-      // To do so, we divide up the linearized stream-K units into clusters and share the same K
-      // offsets for work within clusters.
-
-      uint64_t 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.
-      uint64_t sk_cluster_tiles = params.div_cluster_size(params.sk_tiles_);
-      uint64_t sk_cluster_tiles_in_group = params.divmod_sk_groups_.divide(sk_cluster_tiles);
-      if (group_idx < params.big_groups_) {
-        ++sk_cluster_tiles_in_group;
+      // Increment by one for each of the big clusters (since all big units precede this unit)
+      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;
       }
-
-      // Determine whether we are in a "big unit" within the group, that will process
-      // an additional K chunk in the group.
-      uint64_t sk_tiles_in_group = sk_cluster_tiles_in_group * params.get_cluster_size();
-      uint64_t k_tiles_in_group = sk_tiles_in_group * params.divmod_tiles_per_output_tile_.divisor;
-      uint64_t k_tiles_per_unit_in_group = params.divmod_sk_units_per_group_.divide(k_tiles_in_group);
-      uint64_t 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);
-
-      bool is_split_k = params.divmod_splits_.divisor > 1;
-      uint64_t big_unit_cmp_lhs = is_split_k ? split : cluster_linear_work_idx;
-      uint64_t big_unit_cmp_rhs = is_split_k ? params.big_units_ : big_units_in_group;
-      uint64_t linear_idx_mult = is_split_k ? params.divmod_tiles_per_output_tile_.divisor : k_tiles_per_unit_in_group;
-      uint64_t k_tiles_per_split = is_split_k ? params.divmod_k_tiles_per_sk_unit_.divisor : 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) +
-                                 (k_tiles_per_split * split);
-
-      // Adjust the starting position and number of k iterations for "big units," which
-      // 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_lhs;
-        ++k_tiles_in_my_split;
+      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;
       }
-      else {
-        // Increment by one for each of the big clusters (since all big units precede this unit)
-        unit_iter_start += big_unit_cmp_rhs;
+      else if (params.ktile_start_alignment_count_ == 2 && start_tile_k_tile % 2 != 0) {
+        // ktile for each SM start from even number
+        // If start from odd number ktile within the output tile
+        //    now start at the ktile one before my initial ktile start (take one ktile from prev sm)
+        // if end on odd number ktile within the output tile
+        //    now end at ktile that one before my ktile end (give one ktile to next sm)
+        unit_iter_start -= 1;
+        k_tiles_in_my_split += 1;
       }
+    }
+    if (work_tile_info.k_tile_count == 0) {
+      // This is a new unit
 
       if (!is_split_k) {
-        // Adjust the unit starting position and number of tiles to avoid
+        //
+        // Adjust the unit ending 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,
+        //
+
+        // 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.
-          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;
+          k_tiles_in_my_split -= end_tile_k_tile;
         }
-        else if (params.ktile_start_alignment_count == 2 && start_tile_k_tile % 2 != 0) {
+        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);
+        }
+        else if (params.ktile_start_alignment_count_ == 2 && end_tile_k_tile % 2 != 0) {
           // ktile for each SM start from even number
           // If start from odd number ktile within the output tile
           //    now start at the ktile one before my initial ktile start (take one ktile from prev sm)
-          // if end on odd number ktile within the output tile
+          // If end on odd number ktile within the output tile,
           //    now end at ktile that one before my ktile end (give one ktile to next sm)
-          unit_iter_start -= 1;
-          k_tiles_in_my_split += 1;
+          k_tiles_in_my_split -= 1;
         }
       }
 
-      if (work_tile_info.k_tile_count == 0) {
-        // This is a new unit
-
-        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);
-          }
-          else if (params.ktile_start_alignment_count == 2 && end_tile_k_tile % 2 != 0) {
-            // ktile for each SM start from even number
-            // If start from odd number ktile within the output tile
-            //    now start at the ktile one before my initial ktile start (take one ktile from prev sm)
-            // If end on odd number ktile within the output tile,
-            //    now end at ktile that one before my ktile end (give one ktile to next sm)
-            k_tiles_in_my_split -= 1;
-          }
-        }
-
-        work_tile_info.k_tile_remaining = k_tiles_in_my_split;
-      }
-
-      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
-      // tile, while the beginning (partial) tile is typically the ending of another
-      // output tile. Since ending portions of an output tile must reduce across
-      // 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.
-
-      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
-      output_tile_id *= params.get_cluster_size();
-
-      // The final linearized tile ID is in units of the cluster dimension over which we rasterize.
-      if (params.raster_order_ == RasterOrder::AlongN) {
-        output_tile_id += cta_n_in_cluster * params.divmod_cluster_shape_minor_.divisor;
-      }
-      else {
-        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(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(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 - output_tile_iter_start);
-      work_tile_info.k_tile_count = tile_iter_end - tile_iter_start;
+      work_tile_info.k_tile_remaining = k_tiles_in_my_split;
     }
+    return unit_iter_start;
+  }
+
+  // Update output tile index given existing remaining k tiles of current work tile.
+  CUTLASS_DEVICE
+  static uint64_t update_output_tile_id_and_work_tile_k(
+    Params const& params,
+    WorkTileInfo& work_tile_info,
+    uint64_t linear_idx,
+    uint32_t unit_iter_start,
+    uint64_t cta_m_in_cluster,
+    uint64_t cta_n_in_cluster) {
+    // we divide up the linearized stream-K units into clusters and share the same K
+    // offsets for work within clusters.
+    uint64_t cluster_linear_work_idx = params.div_cluster_size(linear_idx);
+
+    uint64_t unused, group_idx;
+    params.divmod_sk_groups_(unused, group_idx, cluster_linear_work_idx);
+
+    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
+    // tile, while the beginning (partial) tile is typically the ending of another
+    // output tile. Since ending portions of an output tile must reduce across
+    // 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.
+
+    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.
+    uint64_t 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
+    output_tile_id *= params.get_cluster_size();
+
+    // The final linearized tile ID is in units of the cluster dimension over which we rasterize.
+    if (params.raster_order_ == RasterOrder::AlongN) {
+      output_tile_id += cta_n_in_cluster * params.divmod_cluster_shape_minor_.divisor;
+    }
+    else {
+      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(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(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 - output_tile_iter_start);
+    work_tile_info.k_tile_count = tile_iter_end - tile_iter_start;
+
+    return output_tile_id;
+  }
+  // Given output tile index, update M, N, L index of current work tile info.
+  CUTLASS_DEVICE
+  static void
+  update_work_tile_m_n_l(
+    Params const& params,
+    uint32_t output_tile_id,
+    WorkTileInfo& work_tile_info,
+    uint64_t cta_m_in_cluster,
+    uint64_t cta_n_in_cluster) {
 
     uint64_t work_idx_l, remainder;
     params.divmod_batch_(work_idx_l, remainder, output_tile_id);
@@ -907,18 +953,81 @@ private:
     work_tile_info.L_idx = static_cast<int32_t>(work_idx_l);
   }
 
+  // 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.
+  CUTLASS_DEVICE
+  static void
+  assign_work(
+    Params const& params,
+    uint64_t linear_idx,
+    dim3 block_id_in_cluster,
+    WorkTileInfo& work_tile_info,
+    uint32_t &unit_iter_start) {
+
+    auto [cta_m_in_cluster, cta_n_in_cluster, _] =
+      get_current_work_cta_m_n_in_cluster(params, linear_idx, block_id_in_cluster);
+
+    uint64_t output_tile_id = linear_idx;
+    if (linear_idx >= params.units_per_problem_ * params.divmod_splits_.divisor) {
+      // Separate-reduction work
+      auto cluster_size = params.get_cluster_size();
+      // Divide up the linearized separate reduction units into clusters
+      uint64_t 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.divmod_splits_.divisor == 1) {
+      // Data-parallel work
+      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;
+    }
+    else {
+      unit_iter_start = get_current_work_iter_start_possible_update_work_tile_k_remaining(params, linear_idx, work_tile_info);
+      output_tile_id = update_output_tile_id_and_work_tile_k(params, work_tile_info,
+        linear_idx, unit_iter_start, cta_m_in_cluster, cta_n_in_cluster);
+    }
+    update_work_tile_m_n_l(params, output_tile_id, work_tile_info, cta_m_in_cluster, cta_n_in_cluster);
+  }
+
+  // The fast path to get current output tile index then update fields of work tile info
+  // when continuing current work tile is needed, since k tile starting index has precomputed
+  // in the first time fetching current work tile.
+  CUTLASS_DEVICE
+  static void
+  fast_assign_work(
+    uint32_t unit_iter_start,
+    Params const& params,
+    uint64_t linear_idx,
+    dim3 block_id_in_cluster,
+    WorkTileInfo& work_tile_info) {
+
+    auto [cta_m_in_cluster, cta_n_in_cluster, _] =
+      get_current_work_cta_m_n_in_cluster(params, linear_idx, block_id_in_cluster);
+
+    uint64_t output_tile_id = update_output_tile_id_and_work_tile_k(params, work_tile_info,
+      linear_idx, unit_iter_start, cta_m_in_cluster, cta_n_in_cluster);
+
+    update_work_tile_m_n_l(params, output_tile_id, work_tile_info, cta_m_in_cluster, cta_n_in_cluster);
+  }
+
   // 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) {
+  tile_peer_range(Params const& params, uint32_t tile_idx, WorkTileInfo const& work_tile_info) {
+    uint32_t cur_k_tile = static_cast<uint32_t>(work_tile_info.K_idx);
     uint32_t tile_idx_in_cluster_path = params.div_cluster_size(tile_idx);
     uint32_t start_k_tile = params.divmod_tiles_per_output_tile_.divisor * tile_idx_in_cluster_path;
     uint32_t end_k_tile = start_k_tile + params.divmod_tiles_per_output_tile_.divisor - 1;
     uint32_t big_unit_k_tiles = params.big_units_ * (params.divmod_k_tiles_per_sk_unit_.divisor + 1);
 
-    auto adjust_unit = [&](uint32_t k_tile, uint32_t unit_idx, uint32_t k_tiles_per_unit) {
-      uint32_t unit_k_start = unit_idx * k_tiles_per_unit;
-      uint32_t unit_k_end = unit_k_start + k_tiles_per_unit;
+    auto adjust_unit = [&](uint32_t k_tile, uint32_t unit_idx, uint32_t unit_k_start, uint32_t unit_k_end) {
       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,
@@ -943,17 +1052,22 @@ private:
       if (k_tile < big_unit_k_tiles) {
         // The tile is within the "big unit range"
         uint32_t unit_idx = params.divmod_k_tiles_per_sk_big_unit_.divide(k_tile);
-        return static_cast<uint64_t>(adjust_unit(k_tile, unit_idx, params.divmod_k_tiles_per_sk_big_unit_.divisor));
+        uint32_t unit_k_start = unit_idx * params.divmod_k_tiles_per_sk_big_unit_.divisor;
+        uint32_t unit_k_end = unit_k_start + params.divmod_k_tiles_per_sk_big_unit_.divisor;
+        return static_cast<uint64_t>(adjust_unit(k_tile, unit_idx, unit_k_start, unit_k_end));
       }
       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
-        uint32_t unit_idx = params.divmod_k_tiles_per_sk_unit_.divide(k_tile - big_unit_k_tiles) + params.big_units_;
-        return static_cast<uint64_t>(adjust_unit(k_tile, unit_idx, params.divmod_k_tiles_per_sk_unit_.divisor));
+        uint32_t unit_idx_after_big_units = params.divmod_k_tiles_per_sk_unit_.divide(k_tile - big_unit_k_tiles);
+        uint32_t unit_k_start = unit_idx_after_big_units * params.divmod_k_tiles_per_sk_unit_.divisor + (params.big_units_ * params.divmod_k_tiles_per_sk_big_unit_.divisor);
+        uint32_t unit_k_end = unit_k_start + params.divmod_k_tiles_per_sk_unit_.divisor;
+        uint32_t unit_idx = unit_idx_after_big_units + params.big_units_;
+        return static_cast<uint64_t>(adjust_unit(k_tile, unit_idx, unit_k_start, unit_k_end));
       }
     };
 
-    return cute::make_tuple(find_unit(start_k_tile), find_unit(cur_k_tile), find_unit(end_k_tile));
+    return cute::make_tuple(find_unit(start_k_tile), find_unit(start_k_tile + cur_k_tile), find_unit(end_k_tile));
   }
 };
 

include/cutlass/gemm/kernel/tile_scheduler.hpp

@@ -37,15 +37,11 @@
 
 #include "cutlass/arch/arch.h"
 #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 {
 
-////////////////////////////////////////////////////////////////////////////////
-
 //
 // Tags for specifying tile schedulers
 //
@@ -56,10 +52,12 @@ struct StreamKScheduler { };
 
 struct GroupScheduler { }; // Only used for Grouped GEMMs
 
+} // namespace cutlass::gemm
 ////////////////////////////////////////////////////////////////////////////////
 
-} // namespace cutlass::gemm
-
+#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::kernel::detail {

include/cutlass/gemm/kernel/tile_scheduler_params.h

@@ -50,6 +50,26 @@ namespace detail {
 
 ////////////////////////////////////////////////////////////////////////////////
 
+  CUTLASS_HOST_DEVICE
+  static uint32_t
+  get_max_cta_occupancy(
+    int max_sm_per_gpc,
+    GemmCoord cluster_shape,
+    int sm_count) {
+      // 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;
+      return cta_per_device;
+   }
 //
 // Parameters for SM90 tile schedulers
 //
@@ -247,20 +267,7 @@ struct PersistentTileSchedulerSm90Params {
       * 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);
-      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;
-
-      if (sm_count < cta_per_device) {
-        cta_per_device = sm_count;
-      }
+      cta_per_device = get_max_cta_occupancy(max_sm_per_gpc, cluster_shape, sm_count);
       if (raster_order == RasterOrder::AlongN) {
         launch_grid.y = possibly_truncate(
             cta_per_device       / cluster_shape.m(),
@@ -467,7 +474,7 @@ struct PersistentTileSchedulerSm90StreamKParams {
   static constexpr uint32_t max_sk_groups_ = 8u;
 
   // ktile start from even for each cta
-  uint32_t ktile_start_alignment_count { 1u };
+  uint32_t ktile_start_alignment_count_ { 1u };
 
   // Divides dividend by the cluster size
   CUTLASS_HOST_DEVICE
@@ -519,7 +526,7 @@ struct PersistentTileSchedulerSm90StreamKParams {
     ReductionMode reduction_mode,
     DecompositionMode decomposition_mode,
     void* workspace,
-    const uint32_t epilogue_subtile = 1
+    const uint32_t epilogue_subtile = 1u
   ) {
     dim3 problem_blocks = UnderlyingParams::get_tiled_cta_shape_mnl(
       problem_shape, tile_shape, cluster_shape);
@@ -559,6 +566,15 @@ struct PersistentTileSchedulerSm90StreamKParams {
     void* workspace,
     const uint32_t epilogue_subtile = 1
   ) {
+
+    #if !defined(__CUDACC_RTC__)
+    if (hw_info.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.");
+      hw_info.sm_count = KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
+    }
+    #endif // !defined(__CUDACC_RTC__) 
+
     UnderlyingParams underlying_params;
     underlying_params.initialize(
       problem_blocks,
@@ -568,115 +584,43 @@ struct PersistentTileSchedulerSm90StreamKParams {
       raster_order_option
     );
 
-    auto problem_blocks_l = problem_blocks.z;
+    // Set basic parameters that not affected by any heuristics in advance.
+    set_params_base(underlying_params, workspace);
 
-    auto problem_blocks_m = round_up(problem_blocks.x, (1 << underlying_params.log_swizzle_size_) * cluster_shape.m());
-    auto problem_blocks_n = round_up(problem_blocks.y, (1 << underlying_params.log_swizzle_size_) * cluster_shape.n());
-    uint64_t output_tiles = problem_blocks_m * problem_blocks_n * problem_blocks_l;
-
-    // Reduction workspace is at the beginning of the workspace. Lock workspace follows.
-    void* reduction_workspace = workspace;
-
-    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
-      if (splits > hw_info.sm_count) {
-        splits = hw_info.sm_count;
-      }
-
-      // Don't split by more than the K tile iterations
-      //
-      // splits is almost certainly nonnegative here (e.g., hw_info.sm_count,
-      // despite being an int, is a count), so it can safely be converted to unsigned
-      // in the comparison to avoid a signed-unsigned comparison warning-as-error.
-      if (static_cast<decltype(k_tiles_per_output_tile)>(splits) > k_tiles_per_output_tile) {
-        splits = k_tiles_per_output_tile;
-      }
-
-      // If splits == k_tiles_per_output_tiles, there will be one k_tile per cta
-      //   and this violate k_tile start from even requirements. Thus we need to
-      //   reduce the number of splits.
-      if (ktile_start_alignment_count > 1u &&
-           static_cast<decltype(k_tiles_per_output_tile)>(splits) == k_tiles_per_output_tile) { 
-        splits = k_tiles_per_output_tile / ktile_start_alignment_count;
-      }
-
-      set_params_basic(
-        underlying_params,
-        problem_blocks_m,
-        problem_blocks_n,
-        problem_blocks_l,
-        splits,
-        k_tiles_per_output_tile,
-        reduction_workspace,
-        reduction_mode
-      );
-      return;
-    }
-
-    // Calculate the maximum number of blocks from clusters of shape cluster_shape that we
-    // can fit within sm_count SMs.
-    dim3 grid = get_grid_shape(
+    // Call for internal streamk heuristic to setup streamk related params
+    stream_k_heuristic(
+      underlying_params,
       problem_blocks,
+      k_tiles_per_output_tile,
       cluster_shape,
       hw_info,
+      splits,
       max_swizzle,
-      raster_order_option
-    );
+      raster_order_option,
+      decomposition_mode,
+      reduction_mode,
+      epilogue_subtile
+    ); 
+  }
+  
+  // 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 calculate_groups(
+    UnderlyingParams underlying_params,
+    ReductionMode reduction_mode,
+    uint32_t problem_blocks_m,
+    uint32_t problem_blocks_n,
+    GemmCoord cluster_shape,
+    uint64_t cluster_size,
+    uint32_t sk_tiles,
+    uint64_t sk_cluster_tiles,
+    uint64_t sk_units,
+    uint32_t k_tiles_per_output_tile,
+    bool do_separate_reduction) {
 
-    uint64_t ctas_per_wave = grid.x * grid.y;
-    auto cluster_size = cluster_shape.m() * cluster_shape.n();
-    // 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,
-      cluster_size,
-      k_tiles_per_output_tile,
-      decomposition_mode
-    );
-    uint64_t dp_tiles = output_tiles - sk_tiles;
-
-    // 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
-    // (e.g., the callers into this scheduler will be persistent thread blocks).
-    // A work unit can encompass multiple output tiles worth of work (as will be the
-    // case for stream-K blocks).
-    // Since splitting is not required for data-parallel tiles, only one data-parallel unit
-    // is needed per data-parallel tile.
-    uint64_t dp_units = dp_tiles;
-
-    uint64_t ctas_per_sk_wave = ctas_per_wave;
-    uint64_t sk_units = get_num_sk_units(cluster_shape, ctas_per_sk_wave, sk_tiles, k_tiles_per_output_tile);
-
-    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();
@@ -691,14 +635,16 @@ struct PersistentTileSchedulerSm90StreamKParams {
     // 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) {
+    // Grouping is disabled when separate reduction is used because grouping is primarily an attempt
+    // to improve L2 locality, and L2-locality optimizations are unnecessary when the the kernel
+    // is a single wave (which is the case for separate reduction).
+    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) {
@@ -737,82 +683,281 @@ struct PersistentTileSchedulerSm90StreamKParams {
     if (groups == 1 && fallback_groups > 0) {
       groups = fallback_groups;
     }
+    return groups;
+  }
 
-    auto sk_units_per_group = sk_units / groups;
+  // Stream-K kernel use below function to set stream-K feature related parameters to choose
+  // optimal/customized decomposition mode.
+  void stream_k_heuristic(
+    UnderlyingParams underlying_params,
+    dim3 problem_blocks,
+    uint32_t k_tiles_per_output_tile,
+    GemmCoord cluster_shape,
+    KernelHardwareInfo hw_info,
+    int splits,
+    int max_swizzle,
+    RasterOrderOptions raster_order_option,
+    DecompositionMode decomposition_mode,
+    ReductionMode reduction_mode,
+    const uint32_t epilogue_subtile = 1
+    ) {
+    uint32_t groups = 0;
+    uint32_t sk_tiles = 0;
+    uint64_t sk_units = 0;
+    uint64_t cluster_size = 0;
+    uint64_t dp_units = 0;
+    uint64_t k_tiles_per_group = 0;
+    uint64_t k_tiles_per_sk_unit = 0;
+    uint64_t sk_big_groups = 0;
+    uint32_t sk_splits = 1;
+    // Self calculated optimal heuristic mode
+    DecompositionMode heuristic_mode =
+      select_decomposition_mode(
+        groups,
+        sk_tiles,
+        sk_units,
+        cluster_size,
+        dp_units,
+        k_tiles_per_group,
+        k_tiles_per_sk_unit,
+        sk_big_groups,
+        sk_splits,
+        underlying_params,
+        problem_blocks,
+        k_tiles_per_output_tile,
+        cluster_shape,
+        hw_info,
+        splits,
+        max_swizzle,
+        raster_order_option,
+        decomposition_mode,
+        reduction_mode,
+        epilogue_subtile
+      );
 
-    // 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_cluster_tiles_per_group = sk_cluster_tiles / groups;
-    uint64_t sk_tiles_per_group = sk_cluster_tiles_per_group * cluster_size;
+    // Given heuristic_mode returned from the heuristic() method, set params fields.
+    // Here, we decouple the params that have no relation with
+    // decomposition mode from the params that are decided within heuristic().
+    set_params(
+      heuristic_mode,
+      groups,
+      sk_tiles,
+      sk_units,
+      cluster_size,
+      dp_units,
+      k_tiles_per_group,
+      k_tiles_per_sk_unit,
+      sk_big_groups,
+      sk_splits,
+      underlying_params,
+      problem_blocks,
+      k_tiles_per_output_tile,
+      cluster_shape,
+      splits,
+      epilogue_subtile,
+      reduction_mode);
+  }
 
-    // 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_cluster_tiles % groups;
+  // Return the optimal decomposition result by heuristic.
+  DecompositionMode select_decomposition_mode(
+    uint32_t &groups,
+    uint32_t &sk_tiles,
+    uint64_t &sk_units,
+    uint64_t &cluster_size,
+    uint64_t &dp_units,
+    uint64_t &k_tiles_per_group,
+    uint64_t &k_tiles_per_sk_unit,
+    uint64_t &sk_big_groups,
+    uint32_t &sk_splits,
+    UnderlyingParams underlying_params,
+    dim3 problem_blocks,
+    uint32_t k_tiles_per_output_tile,
+    GemmCoord cluster_shape,
+    KernelHardwareInfo hw_info,
+    int splits,
+    int max_swizzle,
+    RasterOrderOptions raster_order_option,
+    DecompositionMode decomposition_mode,
+    ReductionMode reduction_mode,
+    uint32_t epilogue_subtile
+  ) {
 
-    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;
+    // Get block numbers in m, n and l dimensions
+    if (decomposition_mode == DecompositionMode::SplitK ||
+        (decomposition_mode == DecompositionMode::Heuristic && splits > 1)) {
+      // Short circuit to basic split-K decomposition
+      uint32_t adapted_splits = adjust_split_count(
+        splits, hw_info.sm_count, k_tiles_per_output_tile
+      );
+      sk_splits = adapted_splits;
+      return DecompositionMode::SplitK;
     }
-    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);
+    else {
+      // Calculate the maximum number of blocks from clusters of shape cluster_shape that we
+      // can fit within sm_count SMs.
+      // Get block numbers in m, n and l dimensions
+      auto problem_blocks_l = problem_blocks.z;
+      auto problem_blocks_m = round_up(problem_blocks.x, (1 << underlying_params.log_swizzle_size_) * cluster_shape.m());
+      auto problem_blocks_n = round_up(problem_blocks.y, (1 << underlying_params.log_swizzle_size_) * cluster_shape.n());
+      uint64_t output_tiles = problem_blocks_m * problem_blocks_n * problem_blocks_l;
+      dim3 grid = get_grid_shape(
+        problem_blocks,
+        cluster_shape,
+        hw_info,
+        max_swizzle,
+        raster_order_option
+      );
+      uint64_t ctas_per_wave = grid.x * grid.y;
+      cluster_size = cluster_shape.m() * cluster_shape.n();
+      // The number of output tiles to be computed in stream-K and data-parallel fashion, respectively.
+      sk_tiles = get_num_sk_tiles(
+        output_tiles,
+        ctas_per_wave,
+        cluster_size,
+        k_tiles_per_output_tile,
+        decomposition_mode
+      );
+      uint64_t dp_tiles = output_tiles - sk_tiles;
+      // 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
+      // (e.g., the callers into this scheduler will be persistent thread blocks).
+      // A work unit can encompass multiple output tiles worth of work (as will be the
+      // case for stream-K blocks).
+      // Since splitting is not required for data-parallel tiles, only one data-parallel unit
+      // is needed per data-parallel tile.
+      dp_units = dp_tiles;
+
+      uint64_t ctas_per_sk_wave = ctas_per_wave;
+      sk_units = get_num_sk_units(cluster_shape, ctas_per_sk_wave, sk_tiles, k_tiles_per_output_tile);
+
+      if (decomposition_mode == DecompositionMode::DataParallel ||
+          (decomposition_mode == DecompositionMode::Heuristic && sk_tiles == 0) ||
+          sk_units == 0) {
+        // Short circuit to basic data-parallel decomposition
+        return DecompositionMode::DataParallel;
+      }
+      else {
+        bool do_separate_reduction = should_perform_separate_reduction(
+          epilogue_subtile, sk_units, sk_tiles, dp_tiles, ctas_per_wave);
+        
+        uint64_t sk_cluster_tiles = sk_tiles / cluster_size;
+
+        groups = calculate_groups(underlying_params, reduction_mode, problem_blocks_m, problem_blocks_n, cluster_shape,
+          cluster_size, sk_tiles, sk_cluster_tiles, sk_units, k_tiles_per_output_tile, do_separate_reduction);
+
+        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_cluster_tiles_per_group = sk_cluster_tiles / groups;
+        uint64_t sk_tiles_per_group = sk_cluster_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.
+        sk_big_groups = sk_cluster_tiles % groups;
+
+        k_tiles_per_group = k_tiles_per_output_tile * sk_tiles_per_group;
+
+        // Number of k tiles computed per stream-K unit
+        k_tiles_per_sk_unit = k_tiles_per_group / sk_units_per_group;
+
+        DecompositionMode heuristic_mode;
+        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.
+          sk_splits = static_cast<uint32_t>(sk_units / sk_tiles);
+          heuristic_mode = DecompositionMode::SplitK;
+        }
+        else {
+          // Rest scenario is streamk
+          heuristic_mode = DecompositionMode::StreamK;
+        }
+        // Refresh heuristic_mode using analytical model before choosing streamk/separate_reduction decomposition,
+        // ideally it's to get the final decomposition more accuracy. Comment it as it is place holder at this moment.
+        #if 0
+        uint32_t total_waves = static_cast<uint32_t>((output_tiles + ctas_per_wave - 1) / ctas_per_wave);
+        analytical_model(heuristic_mode, k_tiles_per_output_tile, k_tiles_per_sk_unit,
+          sk_splits, epilogue_subtile, total_waves);
+        #endif
+        return heuristic_mode;
+      }
+    }
+  }
+
+  // Given decomposition mode output from heuristic, set all feilds of params.
+  void set_params(
+    DecompositionMode heuristic_mode,
+    uint32_t groups,
+    uint32_t sk_tiles,
+    uint64_t sk_units,
+    uint64_t cluster_size,
+    uint64_t dp_units,
+    uint64_t k_tiles_per_group,
+    uint64_t k_tiles_per_sk_unit,
+    uint64_t sk_big_groups,
+    uint32_t sk_splits,
+    UnderlyingParams underlying_params,
+    dim3 problem_blocks,
+    uint32_t k_tiles_per_output_tile,
+    GemmCoord cluster_shape,
+    uint32_t splits,
+    uint32_t epilogue_subtile,
+    ReductionMode reduction_mode) {
+    // The highest priority when customers set as splitk mode, may set
+    // with a adpated splits value rather than the original splits
+    // even it does not make sense
+    if (splits > 1 && heuristic_mode == DecompositionMode::SplitK) {
       set_params_basic(
         underlying_params,
-        problem_blocks_m,
-        problem_blocks_n,
-        problem_blocks_l,
-        sk_splits,
+        problem_blocks,
+        cluster_shape,
+        sk_splits, // split-k set by customers
         k_tiles_per_output_tile,
-        reduction_workspace,
         reduction_mode
       );
-      return;
     }
-    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
-    // is essentially a no-op.
-    divmod_clusters_mnl_ = FastDivmodU64(sk_units / cluster_size);
-    divmod_splits_ = FastDivmod(1);
-    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_;
-
-    // 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);
-    divmod_k_tiles_per_sk_unit_ = FastDivmod(static_cast<uint32_t>(k_tiles_per_sk_unit));
-    divmod_k_tiles_per_sk_big_unit_ = FastDivmod(static_cast<uint32_t>(k_tiles_per_sk_unit + 1));
-    reduction_mode_ = reduction_mode;
-    divmod_epilogue_subtile_ = FastDivmodU64(epilogue_subtile);
-    separate_reduction_units_ = reduction_units;
+    else if (heuristic_mode == DecompositionMode::DataParallel) {
+      set_params_basic(
+        underlying_params,
+        problem_blocks,
+        cluster_shape,
+        1, // fast path to fall back to the mode without any split scheme
+        k_tiles_per_output_tile,
+        reduction_mode
+      );
+    }
+    else if (heuristic_mode == DecompositionMode::SplitK) {
+      set_params_basic(
+        underlying_params,
+        problem_blocks,
+        cluster_shape,
+        sk_splits, // splits calculated by heuristic
+        k_tiles_per_output_tile,
+        reduction_mode
+      );
+    }
+    else {
+      // streamk
+      set_params_stream_k(
+        underlying_params,
+        k_tiles_per_output_tile,
+        groups,
+        sk_tiles,
+        sk_units,
+        cluster_size,
+        dp_units,
+        k_tiles_per_group,
+        k_tiles_per_sk_unit,
+        sk_big_groups,
+        reduction_mode,
+        1, /*epilogue_subtile*/
+        0  /*reduction_units*/
+      );
+    }
   }
 
   // Given the inputs, computes the physical grid we should launch.
@@ -897,7 +1042,6 @@ struct PersistentTileSchedulerSm90StreamKParams {
         // or if there is no work to be split.
         return 0;
       }
-
       //
       // The final wave is not full. Perform some stream-K work.
       //
@@ -971,11 +1115,13 @@ struct PersistentTileSchedulerSm90StreamKParams {
     int max_swizzle,
     RasterOrderOptions raster_order_option,
     DecompositionMode decomposition_mode,
+    ReductionMode reduction_mode,
     uint32_t mma_warp_groups,
     uint32_t barrier_bits,
     uint32_t accumulator_bits,
     uint32_t epilogue_subtile = 1,
-    uint32_t num_accumulator_mtxs = 1) {
+    uint32_t num_accumulator_mtxs = 1,
+    uint32_t ktile_start_alignment_count = 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());
@@ -989,12 +1135,6 @@ struct PersistentTileSchedulerSm90StreamKParams {
       barrier_workspace_size = 0;
       reduction_workspace_size = 0;
     }
-    else if (splits > 1 &&
-             (decomposition_mode == DecompositionMode::SplitK || decomposition_mode == DecompositionMode::Heuristic)) {
-      // 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, num_accumulator_mtxs);
-    }
     else {
       KernelHardwareInfo new_hw_info;
       new_hw_info.device_id = hw_info.device_id;
@@ -1025,20 +1165,42 @@ struct PersistentTileSchedulerSm90StreamKParams {
       uint64_t sk_units = get_num_sk_units(cluster_shape, ctas_per_sk_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);
+      if (decomposition_mode == DecompositionMode::SplitK ||
+         (decomposition_mode == DecompositionMode::Heuristic && splits > 1)) {
+        splits = adjust_split_count(
+          splits, new_hw_info.sm_count, k_tiles_per_output_tile
+        );
       }
 
-      // 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(reduction_tiles, tile_shape, accumulator_bits, num_accumulator_mtxs);
+      bool split_k_required = splits > 1 && (decomposition_mode == DecompositionMode::SplitK || decomposition_mode == DecompositionMode::Heuristic);
+      bool split_k_selected = decomposition_mode == DecompositionMode::Heuristic &&
+                              sk_units > sk_tiles &&
+                              sk_tiles != 0 &&
+                              sk_units % sk_tiles == 0;
+
+      if (split_k_required || split_k_selected) {
+        // 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, num_accumulator_mtxs);
+      }
+      else {
+        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(reduction_tiles, tile_shape, accumulator_bits, num_accumulator_mtxs);
+      }
     }
   }
   #endif // !defined(__CUDACC_RTC__)
@@ -1063,11 +1225,13 @@ struct PersistentTileSchedulerSm90StreamKParams {
     int max_swizzle,
     RasterOrderOptions raster_order_option,
     DecompositionMode decomposition_mode,
+    ReductionMode reduction_mode,
     uint32_t mma_warp_groups,
     uint32_t barrier_bits,
     uint32_t element_accumulator_bits,
     uint32_t epilogue_subtile,
-    uint32_t num_accumulator_mtxs) {
+    uint32_t num_accumulator_mtxs,
+    uint32_t ktile_start_alignment_count = 1) {
 
     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();
@@ -1082,11 +1246,13 @@ struct PersistentTileSchedulerSm90StreamKParams {
       max_swizzle,
       raster_order_option,
       decomposition_mode,
+      reduction_mode,
       mma_warp_groups,
       barrier_bits,
       element_accumulator_bits,
       epilogue_subtile,
-      num_accumulator_mtxs
+      num_accumulator_mtxs,
+      ktile_start_alignment_count
     );
   }
 
@@ -1104,11 +1270,13 @@ struct PersistentTileSchedulerSm90StreamKParams {
     int max_swizzle,
     RasterOrderOptions raster_order_option,
     DecompositionMode decomposition_mode,
+    ReductionMode reduction_mode,
     uint32_t mma_warp_groups,
     uint32_t barrier_bits,
     uint32_t element_accumulator_bits,
     uint32_t epilogue_subtile = 1,
-    uint32_t num_accumulator_mtxs = 1) {
+    uint32_t num_accumulator_mtxs = 1,
+    uint32_t ktile_start_alignment_count = 1) {
 
     size_t barrier_workspace_size = 0;
     size_t reduction_workspace_size = 0;
@@ -1126,11 +1294,13 @@ struct PersistentTileSchedulerSm90StreamKParams {
         max_swizzle,
         raster_order_option,
         decomposition_mode,
+        reduction_mode,
         mma_warp_groups,
         barrier_bits,
         element_accumulator_bits,
         epilogue_subtile,
-        num_accumulator_mtxs
+        num_accumulator_mtxs,
+        ktile_start_alignment_count
       );
     #endif
 
@@ -1151,11 +1321,13 @@ struct PersistentTileSchedulerSm90StreamKParams {
     int max_swizzle,
     RasterOrderOptions raster_order_option,
     DecompositionMode decomposition_mode,
+    ReductionMode reduction_mode,
     uint32_t mma_warp_groups,
     uint32_t barrier_bits,
     uint32_t element_accumulator_bits,
     uint32_t epilogue_subtile,
-    CudaHostAdapter* cuda_adapter = nullptr) {
+    CudaHostAdapter* cuda_adapter = nullptr,
+    uint32_t ktile_start_alignment_count = 1) {
 
     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();
@@ -1172,12 +1344,14 @@ struct PersistentTileSchedulerSm90StreamKParams {
       max_swizzle,
       raster_order_option,
       decomposition_mode,
+      reduction_mode,
       mma_warp_groups,
       barrier_bits,
       element_accumulator_bits,
       epilogue_subtile,
       1,
-      cuda_adapter
+      cuda_adapter,
+      ktile_start_alignment_count
     );
   }
 
@@ -1197,12 +1371,14 @@ struct PersistentTileSchedulerSm90StreamKParams {
     int max_swizzle,
     RasterOrderOptions raster_order_option,
     DecompositionMode decomposition_mode,
+    ReductionMode reduction_mode,
     uint32_t mma_warp_groups,
     uint32_t barrier_bits,
     uint32_t element_accumulator_bits,
     uint32_t epilogue_subtile = 1,
     uint32_t num_accumulator_mtxs = 1,
-    CudaHostAdapter* cuda_adapter = nullptr) {
+    CudaHostAdapter* cuda_adapter = nullptr,
+    uint32_t ktile_start_alignment_count = 1) {
 
     #if !defined(__CUDACC_RTC__)
       uint64_t barrier_workspace_size = 0;
@@ -1220,11 +1396,13 @@ struct PersistentTileSchedulerSm90StreamKParams {
         max_swizzle,
         raster_order_option,
         decomposition_mode,
+        reduction_mode,
         mma_warp_groups,
         barrier_bits,
         element_accumulator_bits,
         epilogue_subtile,
-        num_accumulator_mtxs
+        num_accumulator_mtxs,
+        ktile_start_alignment_count
       );
 
       if (barrier_workspace_size > 0) {
@@ -1242,31 +1420,41 @@ struct PersistentTileSchedulerSm90StreamKParams {
     return Status::kSuccess;
   }
 
+  // Set params for basic parameters, which will not affected by different decompositions.
+  void
+  set_params_base(UnderlyingParams const& underlying_params, void* reduction_workspace) {
+    divmod_cluster_shape_major_ = underlying_params.divmod_cluster_shape_major_;
+    divmod_cluster_shape_minor_ = underlying_params.divmod_cluster_shape_minor_;
+    divmod_cluster_blk_major_ = underlying_params.divmod_cluster_blk_major_;
+    log_swizzle_size_ = underlying_params.log_swizzle_size_;
+    raster_order_ = underlying_params.raster_order_;
+    reduction_workspace_ = reduction_workspace;
+  }
+
   void
   set_params_basic(
     UnderlyingParams const& underlying_params,
-    uint32_t blocks_m,
-    uint32_t blocks_n,
-    uint32_t blocks_l,
+    dim3 problem_blocks,
+    GemmCoord cluster_shape,
     uint32_t splits,
     uint32_t k_tiles_per_output_tile,
-    void* reduction_workspace,
     ReductionMode reduction_mode) {
 
-    divmod_cluster_shape_major_ = underlying_params.divmod_cluster_shape_major_;
-    divmod_cluster_shape_minor_ = underlying_params.divmod_cluster_shape_minor_;
+    auto blocks_l = problem_blocks.z;
+    auto blocks_m = round_up(problem_blocks.x,
+                             (1 << underlying_params.log_swizzle_size_) * cluster_shape.m());
+    auto blocks_n = round_up(problem_blocks.y,
+                             (1 << underlying_params.log_swizzle_size_) * cluster_shape.n());
+
     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;
+    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);
     divmod_splits_ = FastDivmod(splits);
-    divmod_cluster_blk_major_ = underlying_params.divmod_cluster_blk_major_;
-    log_swizzle_size_ = underlying_params.log_swizzle_size_;
     units_per_problem_ = blocks_m * blocks_n * blocks_l;
-    raster_order_ = underlying_params.raster_order_;
     big_units_ = k_tiles_per_output_tile % splits;
-    reduction_workspace_ = reduction_workspace;
     reduction_mode_ = reduction_mode;
     divmod_k_tiles_per_sk_unit_ = FastDivmod(k_tiles_per_output_tile / splits);
     divmod_k_tiles_per_sk_big_unit_ = FastDivmod(k_tiles_per_output_tile / splits + 1);
@@ -1278,6 +1466,55 @@ struct PersistentTileSchedulerSm90StreamKParams {
     separate_reduction_units_ = 0;
   }
 
+  // Set params for streamk(streamk, separate-reduction included) decomposition.
+  void
+  set_params_stream_k(
+    UnderlyingParams const& underlying_params,
+    uint32_t k_tiles_per_output_tile,
+    uint32_t groups,
+    uint32_t sk_tiles,
+    uint64_t sk_units,
+    uint64_t cluster_size,
+    uint64_t dp_units,
+    uint64_t k_tiles_per_group,
+    uint64_t k_tiles_per_sk_unit,
+    uint64_t sk_big_groups,
+    ReductionMode reduction_mode,
+    uint32_t epilogue_subtile,
+    uint32_t reduction_units) {
+    // stream-k and separate-reduction decompostions
+    divmod_batch_ = underlying_params.divmod_batch_;
+    divmod_tiles_per_output_tile_ = FastDivmod(k_tiles_per_output_tile);
+    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
+    // is essentially a no-op.
+    divmod_clusters_mnl_ = FastDivmodU64(sk_units / cluster_size);
+    divmod_splits_ = FastDivmod(1);
+    units_per_problem_ = static_cast<uint32_t>(dp_units + sk_units);
+
+    // Assign big_units_ assuming that group count == 1. This is unused by stream-K
+    // when group count > 1.
+    auto big_units_in_ctas = k_tiles_per_group % sk_units;
+
+    // Store big_units in terms of clusters. big_units_in_ctas is guaranteed to be divisible
+    // by cluster_size because both k_tiles_per_group and k_tiles_per_sk_unit must be a multiple
+    // of cluster_size.
+    auto big_units_in_clusters = big_units_in_ctas / cluster_size;
+    big_units_ = static_cast<uint32_t>(big_units_in_clusters);
+
+    big_groups_ = static_cast<uint32_t>(sk_big_groups);
+    sk_tiles_ = sk_tiles;
+    sk_units_ = static_cast<uint32_t>(sk_units);
+    divmod_k_tiles_per_sk_unit_ = FastDivmod(static_cast<uint32_t>(k_tiles_per_sk_unit));
+    divmod_k_tiles_per_sk_big_unit_ = FastDivmod(static_cast<uint32_t>(k_tiles_per_sk_unit + 1));
+    reduction_mode_ = reduction_mode;
+    divmod_epilogue_subtile_ = FastDivmodU64(epilogue_subtile);
+    separate_reduction_units_ = reduction_units;
+  }
+
   private:
   // Round up number of bytes to the nearest multiple of L2 cache line alignment
   CUTLASS_HOST_DEVICE
@@ -1286,8 +1523,31 @@ struct PersistentTileSchedulerSm90StreamKParams {
     constexpr size_t L2CacheLineSizeBytes = 128u;
     return (bytes + L2CacheLineSizeBytes - 1) / L2CacheLineSizeBytes * L2CacheLineSizeBytes;
   }
+
+  CUTLASS_HOST_DEVICE
+  static int adjust_split_count(
+      int splits,
+      int sm_count,
+      uint32_t k_tiles_per_output_tile
+      ) {
+    // Don't split by more than the available number of SMs
+    if (splits > sm_count) {
+      splits = sm_count;
+    }
+
+    // Don't split by more than the K tile iterations
+    if (static_cast<uint32_t>(splits) > k_tiles_per_output_tile) {
+      splits = k_tiles_per_output_tile;
+    }
+
+    // If k_tiles_per_output_tiles / splits == 1, there will be one k_tile per cta
+    //   and this violate k_tile start from even requirements. Thus we need to
+    //   reduce the number of splits.
+    return splits;
+  }
 };
 
+
 ////////////////////////////////////////////////////////////////////////////////
 
 // Parameters for SM90 persistent group scheduler (only used for Grouped Gemms)
@@ -1453,18 +1713,7 @@ struct PersistentTileSchedulerSm90GroupParams {
       // 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;
+      int cta_per_device = get_max_cta_occupancy(max_sm_per_gpc, cluster_shape, sm_count);
 
       if (raster_order == RasterOrder::AlongN) {
         launch_grid.y = possibly_truncate(

include/cutlass/gemm/thread/mma_sm50.h

@@ -147,7 +147,7 @@ struct MmaGeneric {
     CUTLASS_PRAGMA_UNROLL
     for (int k = 0; k < Shape::kK; ++k) {
       #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 860)
-      if (kMultipleOf2 && kAllFp32) {
+      if constexpr (kMultipleOf2 && kAllFp32) {
         //2x2 zigzag - m and n loops to increment by 2. Inner loop to process 4 multiply-adds in a 2x2 tile.
         CUTLASS_PRAGMA_UNROLL
         for (int n = 0; n < Shape::kN; n+=2) {
@@ -396,34 +396,36 @@ struct MmaGeneric<
     CUTLASS_PRAGMA_UNROLL
     for (int k = 0; k < Shape::kK; ++k) {
 
-      CUTLASS_PRAGMA_UNROLL
-      for (int n = 0; n < Shape::kN; ++n) {
-
+      {
         CUTLASS_PRAGMA_UNROLL
-        for (int m = 0; m < Shape::kM; ++m) {
+        for (int n = 0; n < Shape::kN; ++n) {
 
-          int m_serpentine = (n % 2) ? (Shape::kM - 1 - m) : m;
+          CUTLASS_PRAGMA_UNROLL
+          for (int m = 0; m < Shape::kM; ++m) {
 
-          MatrixCoord mn(m_serpentine, n);
-          MatrixCoord mk(m_serpentine, k);
-          MatrixCoord kn(k, n);
+            int m_serpentine = (n % 2) ? (Shape::kM - 1 - m) : m;
 
-          Array<ElementC, 1> d;
-          Array<ElementA, 1> a;
-          Array<ElementB, 1> b;
+            MatrixCoord mn(m_serpentine, n);
+            MatrixCoord mk(m_serpentine, k);
+            MatrixCoord kn(k, n);
 
-          d[0] = d_ref.at(mn);
-          a[0] = a_ref.at(mk);
-          b[0] = b_ref.at(kn);
+            Array<ElementC, 1> d;
+            Array<ElementA, 1> a;
+            Array<ElementB, 1> b;
 
-          if ((m == 0 && n) || m == Shape::kM - 1) {
-            mma_corner(d, a, b, d);
+            d[0] = d_ref.at(mn);
+            a[0] = a_ref.at(mk);
+            b[0] = b_ref.at(kn);
+
+            if ((m == 0 && n) || m == Shape::kM - 1) {
+              mma_corner(d, a, b, d);
+            }
+            else {
+              mma_column(d, a, b, d);
+            }
+
+            d_ref.at(mn) = d[0];
           }
-          else {
-            mma_column(d, a, b, d);
-          }
-
-          d_ref.at(mn) = d[0];
         }
       }
     }

include/cutlass/gemm/threadblock/ell_mma_multistage.h

@@ -243,12 +243,12 @@ public:
             if (is_offset_constant){
               auto ell_offset = ell_iter.get_offset_fast();
               is_valid = is_valid && (ell_offset >= 0);
-              gmem_ptr +=  ell_offset * sizeof(IteratorA::Element) / kSrcBytes;
+              gmem_ptr +=  ell_offset * sizeof(typename IteratorA::Element) / kSrcBytes;
             } else {
               int k_offset = iterator_A.get_k();
               auto ell_offset = ell_iter.get_offset(k_offset);
               is_valid = is_valid && (ell_offset >= 0);
-              gmem_ptr += (ell_offset * sizeof(IteratorA::Element)) / kSrcBytes;
+              gmem_ptr += (ell_offset * sizeof(typename IteratorA::Element)) / kSrcBytes;
             }
           }
 
@@ -287,12 +287,12 @@ public:
             if (is_offset_constant){
               auto ell_offset = ell_iter.get_offset_fast();
               is_valid = is_valid && (ell_offset >= 0);
-              gmem_ptr += ell_offset * sizeof(IteratorB::Element) / kSrcBytes;
+              gmem_ptr += ell_offset * sizeof(typename IteratorB::Element) / kSrcBytes;
             } else {
               int k_offset = iterator_B.get_k();
               auto ell_offset = ell_iter.get_offset(k_offset);
               is_valid = is_valid && (ell_offset >= 0);
-              gmem_ptr += ( ell_offset * sizeof(IteratorB::Element)) / kSrcBytes;
+              gmem_ptr += ( ell_offset * sizeof(typename IteratorB::Element)) / kSrcBytes;
             }
           }
 
@@ -359,12 +359,12 @@ public:
             if (is_offset_constant){
               auto ell_offset = ell_iterator.get_offset_fast();
               is_valid = is_valid && (ell_offset >= 0);
-              gmem_ptr +=  ell_offset * sizeof(IteratorA::Element) / kSrcBytes;
+              gmem_ptr +=  ell_offset * sizeof(typename IteratorA::Element) / kSrcBytes;
             } else {
               int k_offset = iterator_A.get_k();
               auto ell_offset = ell_iterator.get_offset(k_offset);
               is_valid = is_valid && (ell_offset >= 0);
-              gmem_ptr += (ell_offset * sizeof(IteratorA::Element)) / kSrcBytes;
+              gmem_ptr += (ell_offset * sizeof(typename IteratorA::Element)) / kSrcBytes;
             }
           }
 
@@ -401,12 +401,12 @@ public:
             if (is_offset_constant){
               auto ell_offset = ell_iterator.get_offset_fast();
               is_valid = is_valid && (ell_offset >= 0);
-              gmem_ptr += ell_offset * sizeof(IteratorB::Element) / kSrcBytes;
+              gmem_ptr += ell_offset * sizeof(typename IteratorB::Element) / kSrcBytes;
             } else {
               int k_offset = iterator_B.get_k();
               auto ell_offset = ell_iterator.get_offset(k_offset);
               is_valid = is_valid && (ell_offset >= 0);
-              gmem_ptr += ( ell_offset * sizeof(IteratorB::Element)) / kSrcBytes;
+              gmem_ptr += ( ell_offset * sizeof(typename IteratorB::Element)) / kSrcBytes;
             }
           }
 

include/cutlass/integer_subbyte.h

@@ -93,7 +93,7 @@ struct integer_subbyte {
       [[maybe_unused]] constexpr int lower_bound = -(1 << (Bits - 1));
       [[maybe_unused]] constexpr int upper_bound = (1 << (Bits - 1)) - 1;
       assert(value >= lower_bound);
-      assert(value < upper_bound);
+      assert(value <= upper_bound);
     }
     else {
       [[maybe_unused]] constexpr unsigned upper_bound = 1u << Bits;
@@ -112,7 +112,7 @@ struct integer_subbyte {
       [[maybe_unused]] constexpr int lower_bound = -(1 << (Bits - 1));
       [[maybe_unused]] constexpr int upper_bound = (1 << (Bits - 1)) - 1;
       assert(value >= lower_bound);
-      assert(value < upper_bound);
+      assert(value <= upper_bound);
     }
     else {
       [[maybe_unused]] constexpr unsigned upper_bound = 1u << Bits;
@@ -120,6 +120,10 @@ struct integer_subbyte {
     }
   }
 
+  CUTLASS_HOST_DEVICE explicit
+  integer_subbyte(uint8_t value)
+    : integer_subbyte(static_cast<unsigned>(value)) {}
+
   // Convert to the "external" integer type (int or unsigned)
   CUTLASS_HOST_DEVICE
   operator xint_t() const {

include/cutlass/kernel_launch.h

@@ -37,6 +37,7 @@
 #include <cuda_runtime_api.h>
 #include "cutlass/cutlass.h"
 #include "cutlass/trace.h"
+#include "cutlass/device_kernel.h" // cutlass::device_kernel
 
 namespace cutlass {
 

include/cutlass/layout/permute.h

@@ -38,11 +38,8 @@
     computation lies in operator() with private member variables  {col_permute_, row_permute_ and stride_} as new addresses after permute op.
 */
 #pragma once
-#if defined(__CUDACC_RTC__)
+
 #include <cuda/std/cassert>
-#else
-#include "assert.h"
-#endif
 #include "cutlass/cutlass.h"
 #include "cutlass/fast_math.h"
 #include "cutlass/layout/pitch_linear.h"

include/cutlass/layout/tensor.h

@@ -39,11 +39,9 @@
     defined in cutlass/tensor_ref.h.
 */
 #pragma once
-#if defined(__CUDACC_RTC__)
+
 #include <cuda/std/cassert>
-#else
-#include "assert.h"
-#endif
+
 #include "cutlass/cutlass.h"
 #include "cutlass/fast_math.h"
 #include "cutlass/layout/pitch_linear.h"

include/cutlass/layout/tensor_op_multiplicand_sm70.h

@@ -37,6 +37,7 @@
 #include "cutlass/cutlass.h"
 #include "cutlass/coord.h"
 #include "cutlass/layout/pitch_linear.h"
+#include "cutlass/matrix_coord.h" // cutlass::MatrixCoord
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 

include/cutlass/numeric_conversion.h

@@ -95,7 +95,6 @@ struct NumericConverter {
 //
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-#if defined(__CUDA_ARCH__)
 template <>
 struct NumericConverter<int32_t, float, FloatRoundStyle::round_to_nearest> {
 
@@ -103,50 +102,17 @@ struct NumericConverter<int32_t, float, FloatRoundStyle::round_to_nearest> {
   using source_type = float;
   static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
 
-  CUTLASS_DEVICE
+  CUTLASS_HOST_DEVICE
   static result_type convert(source_type const & s) {
+    #if __CUDA_ARCH__
     return __float2int_rn(s);
-  }
-
-  CUTLASS_DEVICE
-  result_type operator()(source_type const &s) const {
-    return convert(s);
-  }
-};
-
-template <>
-struct NumericConverter<int32_t, float, FloatRoundStyle::round_toward_zero> {
-
-  using result_type = int32_t;
-  using source_type = float;
-  static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero;
-
-  CUTLASS_DEVICE
-  static result_type convert(source_type const & s) {
-
-    return __float2int_rz(s);
-  }
-
-  CUTLASS_DEVICE
-  result_type operator()(source_type const &s) const {
-    return convert(s);
-  }
-};
-
-#elif !defined(__CUDACC_RTC__)
-
-template <>
-struct NumericConverter<int32_t, float, FloatRoundStyle::round_to_nearest> {
-
-  using result_type = int32_t;
-  using source_type = float;
-  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
-
-  static result_type convert(source_type const & s) {
+    #elif !defined(__CUDACC_RTC__)
     std::fesetround(FE_TONEAREST);
-    return (result_type)std::nearbyint(s);
+    return static_cast<result_type>(std::nearbyint(s));
+    #endif
   }
 
+  CUTLASS_HOST_DEVICE
   result_type operator()(source_type const &s) const {
     return convert(s);
   }
@@ -159,16 +125,21 @@ struct NumericConverter<int32_t, float, FloatRoundStyle::round_toward_zero> {
   using source_type = float;
   static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero;
 
+  CUTLASS_HOST_DEVICE
   static result_type convert(source_type const & s) {
+    #if __CUDA_ARCH__
+    return __float2int_rz(s);
+    #elif !defined(__CUDACC_RTC__)
     std::fesetround(FE_TOWARDZERO);
     return (result_type)std::nearbyint(s);
+    #endif
   }
 
+  CUTLASS_HOST_DEVICE
   result_type operator()(source_type const &s) const {
     return convert(s);
   }
 };
-#endif
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 //
@@ -176,7 +147,6 @@ struct NumericConverter<int32_t, float, FloatRoundStyle::round_toward_zero> {
 //
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-#if defined(__CUDA_ARCH__)
 template <>
 struct NumericConverter<int8_t, float, FloatRoundStyle::round_to_nearest> {
 
@@ -184,13 +154,21 @@ struct NumericConverter<int8_t, float, FloatRoundStyle::round_to_nearest> {
   using source_type = float;
   static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
 
-  CUTLASS_DEVICE
+  CUTLASS_HOST_DEVICE
   static result_type convert(source_type const & s) {
-
+    #if defined(__CUDA_ARCH__)
     int32_t intermediate;
     asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(intermediate) : "f"(s));
-
     return static_cast<result_type>(intermediate);
+    #elif !defined(__CUDACC_RTC__)
+    std::fesetround(FE_TONEAREST);
+    int32_t intermediate = (int32_t)std::nearbyint(s);
+    // Low-end saturation
+    intermediate = std::max(intermediate, (int32_t)std::numeric_limits<int8_t>::lowest());
+    // High-end saturation
+    intermediate = std::min(intermediate, (int32_t)std::numeric_limits<int8_t>::max());
+    return static_cast<result_type>(intermediate);
+    #endif
   }
 
   CUTLASS_HOST_DEVICE
@@ -206,16 +184,24 @@ struct NumericConverter<int8_t, float, FloatRoundStyle::round_toward_zero> {
   using source_type = float;
   static FloatRoundStyle const round_style =  FloatRoundStyle::round_toward_zero;
 
-  CUTLASS_DEVICE
+  CUTLASS_HOST_DEVICE
   static result_type convert(source_type const & s) {
-
+    #if defined(__CUDA_ARCH__)
     int32_t intermediate;
     asm volatile("cvt.rzi.sat.s8.f32 %0, %1;" : "=r"(intermediate) : "f"(s));
-
     return static_cast<result_type>(intermediate);
+    #elif !defined(__CUDACC_RTC__)
+    std::fesetround(FE_TOWARDZERO);
+    int32_t intermediate = (int32_t)std::nearbyint(s);
+    // Low-end saturation
+    intermediate = std::max(intermediate, (int32_t)std::numeric_limits<int8_t>::lowest());
+    // High-end saturation
+    intermediate = std::min(intermediate, (int32_t)std::numeric_limits<int8_t>::max());
+    return static_cast<result_type>(intermediate);
+    #endif 
   }
 
-  CUTLASS_DEVICE
+  CUTLASS_HOST_DEVICE
   result_type operator()(source_type const &s) const {
     return convert(s);
   }
@@ -228,13 +214,21 @@ struct NumericConverter<uint8_t, float, FloatRoundStyle::round_to_nearest> {
   using source_type = float;
   static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
 
-  CUTLASS_DEVICE
+  CUTLASS_HOST_DEVICE
   static result_type convert(source_type const & s) {
-
+    #if defined(__CUDA_ARCH__)
     int32_t intermediate;
     asm volatile("cvt.rni.sat.u8.f32 %0, %1;" : "=r"(intermediate) : "f"(s));
-
     return static_cast<result_type>(intermediate);
+    #elif !defined(__CUDACC_RTC__)
+    std::fesetround(FE_TONEAREST);
+    int32_t intermediate = (int32_t)std::nearbyint(s);
+    // Low-end saturation
+    intermediate = std::max(intermediate, (int32_t)std::numeric_limits<uint8_t>::lowest());
+    // High-end saturation
+    intermediate = std::min(intermediate, (int32_t)std::numeric_limits<uint8_t>::max());
+    return static_cast<result_type>(intermediate);
+    #endif
   }
 
   CUTLASS_HOST_DEVICE
@@ -250,125 +244,29 @@ struct NumericConverter<uint8_t, float, FloatRoundStyle::round_toward_zero> {
   using source_type = float;
   static FloatRoundStyle const round_style =  FloatRoundStyle::round_toward_zero;
 
-  CUTLASS_DEVICE
+  CUTLASS_HOST_DEVICE
   static result_type convert(source_type const & s) {
-
+    #if __CUDA_ARCH__
     int32_t intermediate;
     asm volatile("cvt.rzi.sat.u8.f32 %0, %1;" : "=r"(intermediate) : "f"(s));
-
     return static_cast<result_type>(intermediate);
-  }
-
-  CUTLASS_DEVICE
-  result_type operator()(source_type const &s) const {
-    return convert(s);
-  }
-};
-
-#elif !defined(__CUDACC_RTC__)
-
-template <>
-struct NumericConverter<int8_t, float, FloatRoundStyle::round_to_nearest> {
-
-  using result_type = int8_t;
-  using source_type = float;
-  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
-
-  static result_type convert(source_type const & s) {
-    std::fesetround(FE_TONEAREST);
-    int32_t intermediate = (int32_t)std::nearbyint(s);
-
-    // Low-end saturation
-    intermediate = std::max(intermediate, (int32_t)std::numeric_limits<int8_t>::lowest());
-
-    // High-end saturation
-    intermediate = std::min(intermediate, (int32_t)std::numeric_limits<int8_t>::max());
-
-    return static_cast<result_type>(intermediate);
-  }
-
-  result_type operator()(source_type const &s) const {
-    return convert(s);
-  }
-};
-
-template <>
-struct NumericConverter<int8_t, float, FloatRoundStyle::round_toward_zero> {
-
-  using result_type = int8_t;
-  using source_type = float;
-  static FloatRoundStyle const round_style =  FloatRoundStyle::round_toward_zero;
-
-  static result_type convert(source_type const & s) {
+    #elif !defined(__CUDACC_RTC__)
     std::fesetround(FE_TOWARDZERO);
     int32_t intermediate = (int32_t)std::nearbyint(s);
-
-    // Low-end saturation
-    intermediate = std::max(intermediate, (int32_t)std::numeric_limits<int8_t>::lowest());
-
-    // High-end saturation
-    intermediate = std::min(intermediate, (int32_t)std::numeric_limits<int8_t>::max());
-
-    return static_cast<result_type>(intermediate);
-  }
-
-  result_type operator()(source_type const &s) const {
-    return convert(s);
-  }
-};
-
-template <>
-struct NumericConverter<uint8_t, float, FloatRoundStyle::round_to_nearest> {
-
-  using result_type = uint8_t;
-  using source_type = float;
-  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
-
-  static result_type convert(source_type const & s) {
-    std::fesetround(FE_TONEAREST);
-    int32_t intermediate = (int32_t)std::nearbyint(s);
-
     // Low-end saturation
     intermediate = std::max(intermediate, (int32_t)std::numeric_limits<uint8_t>::lowest());
-
     // High-end saturation
     intermediate = std::min(intermediate, (int32_t)std::numeric_limits<uint8_t>::max());
-
     return static_cast<result_type>(intermediate);
+    #endif
   }
 
+  CUTLASS_HOST_DEVICE
   result_type operator()(source_type const &s) const {
     return convert(s);
   }
 };
 
-template <>
-struct NumericConverter<uint8_t, float, FloatRoundStyle::round_toward_zero> {
-
-  using result_type = uint8_t;
-  using source_type = float;
-  static FloatRoundStyle const round_style =  FloatRoundStyle::round_toward_zero;
-
-  static result_type convert(source_type const & s) {
-    std::fesetround(FE_TOWARDZERO);
-    int32_t intermediate = (int32_t)std::nearbyint(s);
-
-    // Low-end saturation
-    intermediate = std::max(intermediate, (int32_t)std::numeric_limits<uint8_t>::lowest());
-
-    // High-end saturation
-    intermediate = std::min(intermediate, (int32_t)std::numeric_limits<uint8_t>::max());
-
-    return static_cast<result_type>(intermediate);
-  }
-
-  result_type operator()(source_type const &s) const {
-    return convert(s);
-  }
-};
-
-#endif
-
 /////////////////////////////////////////////////////////////////////////////////////////////////
 //
 // Partial specializations for float => integer_subbyte
@@ -3281,88 +3179,88 @@ namespace detail {
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-#if defined(__CUDA_ARCH__)
-/// Partial specialization for Array<int8_t, 8> <= Array<int4b_t, 8>
-template <
-  FloatRoundStyle Round
->
-struct NumericArrayConverter<int8_t, int4b_t, 8, Round> {
-
-  using result_type = Array<int8_t, 8>;
-  using source_type = Array<int4b_t, 8>;
-  static FloatRoundStyle const round_style = Round;
-
-  CUTLASS_DEVICE
-  static result_type convert(source_type const & source) {
-
-    unsigned const& storage = reinterpret_cast<unsigned const &>(source);
-    unsigned out[2];
-
-    asm volatile(
-        "{\n"
-        "  .reg .u32 tmp0, tmp1, tmp2;\n"
-        "  shl.b32 tmp0, %2, 4;\n"                // tmp0 = x1x2x3x4x5x6x7__
-        "  and.b32 tmp0, tmp0, 0xf0f0f0f0;\n"     // tmp0 = x1__x3__x5__x7__
-        "  prmt.b32 tmp1, tmp0, tmp0, 0xba98;\n"  // tmp1 = s1s3s5s7
-        "  and.b32 tmp1, tmp1, 0xf0f0f0f0;\n"     // tmp1 = s1__s3__s5__s7__
-        "  shr.u32 tmp0, tmp0, 4;\n"              // tmp0 = __x1__x3__x5__x7
-        "  or.b32 tmp2, tmp0, tmp1;\n"            // tmp2 = y1y3y5y7
-        "  and.b32 tmp0, %2, 0xf0f0f0f0;\n"       // tmp0 = x0__x2__x4__x6__
-        "  prmt.b32 tmp1, tmp0, tmp0, 0xba98;\n"  // tmp1 = s0s2s4s6
-        "  and.b32 tmp1, tmp1, 0xf0f0f0f0;\n"     // tmp1 = s0__s2__s4__s6__
-        "  shr.u32 tmp0, tmp0, 4;\n"              // tmp0 = __x0__x2__x4__x6
-        "  or.b32 tmp0, tmp0, tmp1;\n"            // tmp0 = y0y2y4y6
-        "  prmt.b32 %0, tmp2, tmp0, 0x5140;\n"    // %0 = y0y1y2y3
-        "  prmt.b32 %1, tmp2, tmp0, 0x7362;\n"    // %1 = y4y5y6y7
-        "}\n"
-        : "=r"(out[0]), "=r"(out[1])
-        : "r"(storage));
-
-    return reinterpret_cast<result_type const &>(out);
-  }
-
-  CUTLASS_DEVICE
-  result_type operator()(source_type const &s) const {
-    return convert(s);
-  }
-};
-
 /// Partial specialization for Array<int8_t> <= Array<int4b_t>
 template <
   int N,
   FloatRoundStyle Round
 >
 struct NumericArrayConverter<int8_t, int4b_t, N, Round> {
-  static_assert(!(N % 8), "N must be multiple of 8.");
+
+  static_assert(N % 8 == 0, "N must be a multiple of 8");
 
   using result_type = Array<int8_t, N>;
   using source_type = Array<int4b_t, N>;
   static FloatRoundStyle const round_style = Round;
 
-  CUTLASS_DEVICE
+  CUTLASS_HOST_DEVICE
   static result_type convert(source_type const & source) {
+   
+    #if defined(__CUDA_ARCH__)
 
-    NumericArrayConverter<int8_t, int4b_t, 8, Round> convert_vector_;
+    if constexpr ( N == 8 ) {
+      
+      unsigned const& storage = reinterpret_cast<unsigned const &>(source);
+      unsigned out[2];
 
-    result_type result;
+      asm volatile(
+          "{\n"
+          "  .reg .u32 tmp0, tmp1, tmp2;\n"
+          "  shl.b32 tmp0, %2, 4;\n"                // tmp0 = x1x2x3x4x5x6x7__
+          "  and.b32 tmp0, tmp0, 0xf0f0f0f0;\n"     // tmp0 = x1__x3__x5__x7__
+          "  prmt.b32 tmp1, tmp0, tmp0, 0xba98;\n"  // tmp1 = s1s3s5s7
+          "  and.b32 tmp1, tmp1, 0xf0f0f0f0;\n"     // tmp1 = s1__s3__s5__s7__
+          "  shr.u32 tmp0, tmp0, 4;\n"              // tmp0 = __x1__x3__x5__x7
+          "  or.b32 tmp2, tmp0, tmp1;\n"            // tmp2 = y1y3y5y7
+          "  and.b32 tmp0, %2, 0xf0f0f0f0;\n"       // tmp0 = x0__x2__x4__x6__
+          "  prmt.b32 tmp1, tmp0, tmp0, 0xba98;\n"  // tmp1 = s0s2s4s6
+          "  and.b32 tmp1, tmp1, 0xf0f0f0f0;\n"     // tmp1 = s0__s2__s4__s6__
+          "  shr.u32 tmp0, tmp0, 4;\n"              // tmp0 = __x0__x2__x4__x6
+          "  or.b32 tmp0, tmp0, tmp1;\n"            // tmp0 = y0y2y4y6
+          "  prmt.b32 %0, tmp2, tmp0, 0x5140;\n"    // %0 = y0y1y2y3
+          "  prmt.b32 %1, tmp2, tmp0, 0x7362;\n"    // %1 = y4y5y6y7
+          "}\n"
+          : "=r"(out[0]), "=r"(out[1])
+          : "r"(storage));
 
-    Array<int8_t, 8> *result_ptr = reinterpret_cast<Array<int8_t, 8> *>(&result);
-    Array<int4b_t, 8> const *source_ptr = reinterpret_cast<Array<int4b_t, 8> const *>(&source);
-
-    CUTLASS_PRAGMA_UNROLL
-    for (int i = 0; i < N / 8; ++i) {
-      result_ptr[i] = convert_vector_(source_ptr[i]);
+      return reinterpret_cast<result_type const &>(out);
+      
+    } else {
+      
+      NumericArrayConverter<int8_t, int4b_t, 8, Round> convert_vector_;
+      
+      result_type result;
+      
+      Array<int8_t, 8> *result_ptr = reinterpret_cast<Array<int8_t, 8> *>(&result);
+      Array<int4b_t, 8> const *source_ptr = reinterpret_cast<Array<int4b_t, 8> const *>(&source);
+      
+      CUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < N / 8; ++i) {
+        result_ptr[i] = convert_vector_(source_ptr[i]);
+      }
+      
+      return result;
     }
-
+    
+    #else
+    
+    result_type result;
+    NumericConverter<int8_t, int4b_t, Round> convert_;
+    
+    CUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = convert_(source[i]);
+    }
+    
     return result;
+    
+    #endif // __CUDA_ARCH__
   }
 
-  CUTLASS_DEVICE
+  CUTLASS_HOST_DEVICE
   result_type operator()(source_type const &s) const {
     return convert(s);
   }
 };
-#endif // defined(__CUDA_ARCH__)
 
 /// Partial specialization for Array<cutlass::float_e4m3_t, N> <= Array<cutlass::int4b_t, N>
 template <FloatRoundStyle Round, int N>

include/cutlass/numeric_size.h

@@ -68,6 +68,15 @@ bits_to_bytes(T bits) {
   return (R(bits) + R(7)) / R(8);
 }
 
+/// Returns the number of bits required to hold a specified number of bytes
+template <class R = int, class T>
+CUTLASS_HOST_DEVICE
+constexpr
+R
+bytes_to_bits(T bytes) {
+  return R(bytes) * R(8);
+}
+
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
 template <class T>

include/cutlass/numeric_types.h

@@ -34,8 +34,6 @@
 */
 #pragma once
 
-#include "cutlass/cutlass.h"
-#include "cutlass/platform/platform.h"
 #include "cutlass/numeric_size.h"
 
 /////////////////////////////////////////////////////////////////////////////////////////////////

include/cutlass/pipeline/sm90_pipeline.hpp

@@ -51,6 +51,65 @@ namespace cutlass {
 
 using namespace cute;
 
+namespace detail {
+
+// Helper function for DEBUG checks
+template<class ThreadCategory>
+CUTLASS_DEVICE
+bool pipeline_is_producer(ThreadCategory role) {
+  return (role == ThreadCategory::Producer || role == ThreadCategory::ProducerConsumer);
+}
+
+template<class ThreadCategory>
+CUTLASS_DEVICE
+void pipeline_check_is_producer(ThreadCategory role) {
+  #ifndef NDEBUG
+  if (!pipeline_is_producer(role)) {
+    asm volatile ("brkpt;\n" ::);
+  }
+  #endif
+}
+
+template<class ThreadCategory>
+CUTLASS_DEVICE
+bool pipeline_is_consumer(ThreadCategory role) {
+  return (role == ThreadCategory::Consumer || role == ThreadCategory::ProducerConsumer);
+}
+
+template<class ThreadCategory>
+CUTLASS_DEVICE
+void pipeline_check_is_consumer(ThreadCategory role) {
+  #ifndef NDEBUG
+  if (!pipeline_is_consumer(role)) {
+    asm volatile ("brkpt;\n" ::);
+  }
+  #endif
+}
+
+CUTLASS_DEVICE
+cute::tuple<bool, uint32_t> spread_arrivals_to_warp(int thread_idx_in_warp) {
+  constexpr uint32_t MaxClusterSize = 16;
+  bool is_signaling_thread = (thread_idx_in_warp % (32 / MaxClusterSize)) == 0;
+  auto layout = Layout<Shape<_4,_4>,Stride<_4, _1>>{};
+  uint32_t thread_row = thread_idx_in_warp / 8;
+  uint32_t thread_col = (thread_idx_in_warp % 8) / 2;
+  uint32_t dst_blockid = layout(thread_row, thread_col);
+  return cute::make_tuple(is_signaling_thread, dst_blockid);
+}
+
+CUTLASS_DEVICE
+cute::tuple<bool, uint32_t> spread_arrivals_to_warpgroup(int thread_idx_in_warpgroup, int warp_idx) {
+  constexpr uint32_t MaxClusterSize = 16;
+  bool is_signaling_thread = (thread_idx_in_warpgroup % (NumThreadsPerWarpGroup / MaxClusterSize)) == 0;
+  auto layout = cute::composition(Swizzle<2,0,-2>{},
+                                  Layout<Shape<_4,_4>,Stride<_4,_1>>{});
+  uint32_t thread_row = warp_idx % 4;
+  uint32_t thread_col = (thread_idx_in_warpgroup / 8) % 4;
+  uint32_t dst_blockid = layout(thread_row, thread_col);
+  return cute::make_tuple(is_signaling_thread, dst_blockid);
+}
+} // namespace detail
+
 enum class BarrierStatus : uint32_t {
   WaitAgain = 0u,
   WaitDone  = 1u,
@@ -210,7 +269,7 @@ PipelineState<Pipeline::Stages> make_producer_start_state() {
 // Currently, it is optional to elect a leader for the Consumers
 template <int Stages_>
 class PipelineTmaAsync {
-public :
+public:
   using FullBarrier = cutlass::arch::ClusterTransactionBarrier;
   using EmptyBarrier = cutlass::arch::ClusterBarrier;
   using ProducerBarrierType = FullBarrier::ValueType;
@@ -237,68 +296,92 @@ public :
     uint32_t num_consumers = 0;
   };
 
-  // Constructor
-  template<class ClusterShape>
+  template <class ClusterShape>
+  static
   CUTLASS_DEVICE
-  PipelineTmaAsync(SharedStorage& storage, Params params, ClusterShape cluster_shape)
+  void
+  init_barriers(SharedStorage& storage, Params params, ClusterShape cluster_shape) {
+    int warp_idx = canonical_warp_idx_sync();
+    bool is_initializing_warp = (warp_idx == 0);
+    if (is_initializing_warp) {
+      // Barrier FULL and EMPTY init
+      constexpr int producer_arv_cnt = 1;
+      uint32_t const num_consumer_warpgroups_per_cluster = params.num_consumers / NumThreadsPerWarpGroup;
+      uint32_t multicast_consumer_arrival_count = params.num_consumers; // If cluster_size is 1
+      if (cute::size(cluster_shape) > 1) {
+        multicast_consumer_arrival_count = (cute::size<0>(cluster_shape) + cute::size<1>(cluster_shape) - 1) *
+              num_consumer_warpgroups_per_cluster;
+      }
+
+      cutlass::arch::detail::initialize_barrier_array_pair_aligned<decltype(storage.full_barrier_), decltype(storage.empty_barrier_), Stages>(
+          storage.full_barrier_, storage.empty_barrier_, producer_arv_cnt, multicast_consumer_arrival_count);
+    }
+  }
+
+  template<class ClusterShape, class InitBarriers, class InitMasks>
+  CUTLASS_DEVICE
+  PipelineTmaAsync(SharedStorage& storage, Params params, ClusterShape cluster_shape, InitBarriers = {}, InitMasks = {})
       : params_(params)
       , full_barrier_ptr_(&storage.full_barrier_[0])
       , empty_barrier_ptr_(&storage.empty_barrier_[0]) {
 
     int warp_idx = canonical_warp_idx_sync();
+    int thread_idx = threadIdx.x;
     int lane_predicate = cute::elect_one_sync();
 
-    if (warp_idx == 0 && lane_predicate == 1) {
-      // Barrier FULL init
-      for (int i = 0; i < Stages; ++i) {
-        full_barrier_ptr_[i].init(1);
+    static_assert(cute::is_same_v<InitBarriers, cute::true_type> || cute::is_same_v<InitBarriers, cute::false_type>);
+    static_assert(cute::is_same_v<InitMasks, cute::true_type> || cute::is_same_v<InitMasks, cute::false_type>);
+    if constexpr (cute::is_same_v<InitBarriers, cute::true_type>) {
+      init_barriers(storage, params_, cluster_shape);
+    }
+
+    if constexpr (cute::is_same_v<InitMasks, cute::true_type>) {
+      // Logic to optimally schedule Empty Arrives
+      // Goal : To divide SYNCS Empty Arrival duty equally amongst the Warp-Group (128 threads)
+      dim3 block_id = cute::block_id_in_cluster();
+      auto cluster_size = cute::size(cluster_shape);
+
+      if (cluster_size == 1) {
+        is_signaling_thread_ = true;
+        dst_blockid_ = 0;
       }
-      uint32_t const num_consumer_warpgroups_per_cluster = params_.num_consumers / NumThreadsPerWarpGroup;
-      uint32_t const multicast_consumer_arrival_count = (cute::size<0>(cluster_shape) + cute::size<1>(cluster_shape) - 1) *
-          num_consumer_warpgroups_per_cluster;
-      // Barrier EMPTY init
-      for (int i = 0; i < Stages; ++i) {
-        empty_barrier_ptr_[i].init(multicast_consumer_arrival_count);
+      else {
+        // STEP 1 : Use Cute Layout function to generate an optimal dst block-id (0-15)
+        if (params_.num_consumers % NumThreadsPerWarpGroup == 0) {
+          auto [is_signaling_thread, dst_blockid] = detail::spread_arrivals_to_warpgroup(thread_idx % NumThreadsPerWarpGroup, warp_idx);
+          is_signaling_thread_ = is_signaling_thread;
+          dst_blockid_ = dst_blockid;
+        }
+        else if (params_.num_consumers == 32) {
+          auto [is_signaling_thread, dst_blockid] = detail::spread_arrivals_to_warp(thread_idx % 32);
+          is_signaling_thread_ = is_signaling_thread;
+          dst_blockid_ = dst_blockid;
+        }
+        else {
+          is_signaling_thread_ = 0;
+          #ifndef NDEBUG
+            asm volatile ("brkpt;\n" ::);
+          #endif
+        }
+
+        // STEP 2: Find if this dst block-id needs an arrival for this problem
+        is_signaling_thread_ &= dst_blockid_ < cluster_size;
+        is_signaling_thread_ &= is_same_row_or_col(dst_blockid_, block_id, cluster_shape);
       }
     }
-    cutlass::arch::fence_barrier_init();
-
-    // Logic to optimally schedule Empty Arrives
-    // Goal : To divide SYNCS Empty Arrival duty equally amongst the Warp-Group (128 threads)
-    dim3 block_id = cute::block_id_in_cluster();
-    auto cluster_size = cute::size(cluster_shape);
-    static constexpr int MaxClusterSize = 16;
-
-    // STEP 1 : Use Cute Layout function to generate an optimal dst block-id (0-15)
-    if (params_.num_consumers % NumThreadsPerWarpGroup == 0) {
-      int thread_idx = threadIdx.x % NumThreadsPerWarpGroup;
-      is_signalling_thread_ = (thread_idx % (NumThreadsPerWarpGroup / MaxClusterSize)) == 0;
-      auto layout = cute::composition(Swizzle<2,0,-2>{},
-                                      Layout<Shape<_4,_4>,Stride<_4,_1>>{});
-      uint32_t thread_row = warp_idx % 4;
-      uint32_t thread_col = (thread_idx / 8) % 4;
-      dst_blockid_ = layout(thread_row, thread_col);
-    }
-    else if (params_.num_consumers == 32) {
-      int thread_idx = threadIdx.x % 32;
-      is_signalling_thread_ = (thread_idx % (32 / MaxClusterSize)) == 0;
-      auto layout = Layout<Shape<_4,_4>,Stride<_4, _1>>{};
-      uint32_t thread_row = thread_idx / 8;
-      uint32_t thread_col = (thread_idx % 8) / 2;
-      dst_blockid_ = layout(thread_row, thread_col);
-    }
-    else {
-      is_signalling_thread_ = 0;
-      #ifndef NDEBUG
-        asm volatile ("brkpt;\n" ::);
-      #endif
-    }
-
-    // STEP 2: Find if this dst block-id needs an arrival for this problem
-    is_signalling_thread_ &= dst_blockid_ < cluster_size;
-    is_signalling_thread_ &= is_same_row_or_col(dst_blockid_, block_id, cluster_shape);
   }
 
+  // Constructor
+  template<class ClusterShape>
+  CUTLASS_DEVICE
+  PipelineTmaAsync(SharedStorage& storage, Params params, ClusterShape cluster_shape)
+      : PipelineTmaAsync(storage, params, cluster_shape, cute::true_type{}, cute::true_type{}) { }
+  
+  template<class ClusterShape, class InitBarriers>
+  CUTLASS_DEVICE
+  PipelineTmaAsync(SharedStorage& storage, Params params, ClusterShape cluster_shape, InitBarriers = {})
+      : PipelineTmaAsync(storage, params, cluster_shape, InitBarriers{}, cute::true_type{}) { }
+
   template <class ClusterShape>
   CUTLASS_DEVICE
   bool is_same_row_or_col(int dst_block_id, dim3 block_id, ClusterShape cluster_shape) {
@@ -347,6 +430,7 @@ public :
   // This should be called once before kernel exits.
   CUTLASS_DEVICE
   void producer_tail(PipelineState state) {
+    detail::pipeline_check_is_producer(params_.role);
     for (int count = 0; count < Stages; ++count) {
       empty_barrier_ptr_[state.index()].wait(state.phase());
       ++state;
@@ -386,15 +470,16 @@ public :
     consumer_release(state.index());
   }
 
-private :
+private:
   uint32_t dst_blockid_ = 0;
-  uint32_t is_signalling_thread_ = 0;
+  uint32_t is_signaling_thread_ = 0;
   FullBarrier *full_barrier_ptr_ = nullptr;
   EmptyBarrier *empty_barrier_ptr_ = nullptr;
   Params params_;
 
   CUTLASS_DEVICE
   ProducerToken producer_try_acquire(uint32_t stage, uint32_t phase, uint32_t skip_wait) {
+    detail::pipeline_check_is_producer(params_.role);
     if (skip_wait) {
       return {BarrierStatus::WaitDone};
     }
@@ -404,6 +489,7 @@ private :
 
   CUTLASS_DEVICE
   void producer_acquire(uint32_t stage, uint32_t phase, ProducerToken barrier_token) {
+    detail::pipeline_check_is_producer(params_.role);
     if (barrier_token != BarrierStatus::WaitDone) {
       empty_barrier_ptr_[stage].wait(phase);
     }
@@ -454,6 +540,7 @@ private :
 
   CUTLASS_DEVICE
   ConsumerToken consumer_try_wait(uint32_t stage, uint32_t phase, uint32_t skip_wait) {
+    detail::pipeline_check_is_consumer(params_.role);
     if (skip_wait) {
       return {BarrierStatus::WaitDone};
     }
@@ -463,6 +550,7 @@ private :
 
   CUTLASS_DEVICE
   ConsumerToken consumer_test_wait(uint32_t stage, uint32_t phase, uint32_t skip_wait) {
+    detail::pipeline_check_is_consumer(params_.role);
     if (skip_wait) {
       return {BarrierStatus::WaitDone};
     }
@@ -473,12 +561,14 @@ private :
   // Wait for producer to commit transactions (done by TMA)
   CUTLASS_DEVICE
   void consumer_wait(uint32_t stage, uint32_t phase) {
+    detail::pipeline_check_is_consumer(params_.role);
     full_barrier_ptr_[stage].wait(phase);
   }
 
   // Wait for producer to commit transactions (done by TMA)
   CUTLASS_DEVICE
   void consumer_wait(uint32_t stage, uint32_t phase, ConsumerToken barrier_token) {
+    detail::pipeline_check_is_consumer(params_.role);
     if (barrier_token == BarrierStatus::WaitAgain) {
       full_barrier_ptr_[stage].wait(phase);
     }
@@ -488,7 +578,8 @@ private :
   // Ensures all blocks in the Same Row and Column get notifed.
   CUTLASS_DEVICE
   void consumer_release(uint32_t stage, uint32_t skip = false) {
-    empty_barrier_ptr_[stage].arrive(dst_blockid_, is_signalling_thread_ & (!skip));
+    detail::pipeline_check_is_consumer(params_.role);
+    empty_barrier_ptr_[stage].arrive(dst_blockid_, is_signaling_thread_ & (!skip));
     #ifndef NDEBUG
     if (params_.role == ThreadCategory::Producer || params_.role == ThreadCategory::NonParticipant) {
       asm volatile ("brkpt;\n" ::);
@@ -625,7 +716,7 @@ private:
 ///////////////////////////////////////////////////////////////////////////////////////////////////
 template <int Stages_>
 class PipelineTransactionAsync {
-public :
+public:
   using FullBarrier = cutlass::arch::ClusterTransactionBarrier;
   using EmptyBarrier = cutlass::arch::ClusterBarrier;
   using ProducerBarrierType = FullBarrier::ValueType;
@@ -653,26 +744,45 @@ public :
     uint32_t dst_blockid = cute::block_rank_in_cluster();
   };
 
-  // Constructor
+  static
   CUTLASS_DEVICE
-  PipelineTransactionAsync(SharedStorage& storage, Params const& params)
+  void
+  init_barriers(SharedStorage& storage, Params const& params) {
+    FullBarrier *full_barrier_ptr = storage.full_barrier_.data();
+    EmptyBarrier *empty_barrier_ptr = storage.empty_barrier_.data();
+    int warp_idx = canonical_warp_idx_sync();
+    bool is_initializing_warp = (warp_idx == 0);
+    if (is_initializing_warp) {
+      // Barrier FULL and EMPTY init
+      cutlass::arch::detail::initialize_barrier_array_pair_aligned<decltype(full_barrier_ptr), decltype(empty_barrier_ptr), Stages>(
+          full_barrier_ptr, empty_barrier_ptr, params.producer_arv_count, params.consumer_arv_count);
+    }
+  }
+
+  // Constructor
+  template<class InitBarriers>
+  CUTLASS_DEVICE
+  PipelineTransactionAsync(SharedStorage& storage, Params const& params, InitBarriers = cute::true_type{})
     : params_(params)
     , full_barrier_ptr_(storage.full_barrier_.data())
     , empty_barrier_ptr_(storage.empty_barrier_.data()) {
+
     int warp_idx = canonical_warp_idx_sync();
     int lane_predicate = cute::elect_one_sync();
 
-    // Barrier FULL, EMPTY init
-    // Init is done only by thread 0 of the block
-    if (warp_idx == 0 && lane_predicate) {
-      for (int i = 0; i < Stages; ++i) {
-        full_barrier_ptr_[i].init(params.producer_arv_count);
-        empty_barrier_ptr_[i].init(params.consumer_arv_count);
-      }
+    static_assert(cute::is_same_v<InitBarriers, cute::true_type> || cute::is_same_v<InitBarriers, cute::false_type>);
+
+    if constexpr (cute::is_same_v<InitBarriers, cute::true_type>) {
+      init_barriers(storage, params);
     }
-    cutlass::arch::fence_barrier_init();
+
   }
 
+  // Constructor
+  CUTLASS_DEVICE
+  PipelineTransactionAsync(SharedStorage& storage, Params const& params) :
+    PipelineTransactionAsync(storage, params, cute::true_type{}) { }
+
   ////////////////////
   // Producer APIs
   ////////////////////
@@ -758,6 +868,7 @@ private:
 
   CUTLASS_DEVICE
   ProducerToken producer_try_acquire(uint32_t stage, uint32_t phase, uint32_t skip_wait) {
+    detail::pipeline_check_is_producer(params_.role);
     if (skip_wait) {
       return {BarrierStatus::WaitDone};
     }
@@ -767,6 +878,7 @@ private:
 
   CUTLASS_DEVICE
   void producer_acquire(uint32_t stage, uint32_t phase, ProducerToken barrier_token) {
+    detail::pipeline_check_is_producer(params_.role);
     if (barrier_token == BarrierStatus::WaitAgain) {
       empty_barrier_ptr_[stage].wait(phase);
     }
@@ -775,11 +887,13 @@ private:
   // Perform an expect-tx operation on the stage's full barrier. Must be called by 1 thread
   CUTLASS_DEVICE
   void producer_expect_transaction(uint32_t stage) {
+    detail::pipeline_check_is_producer(params_.role);
     full_barrier_ptr_[stage].expect_transaction(params_.transaction_bytes);
   }
 
   CUTLASS_DEVICE
   void producer_commit(uint32_t stage) {
+    detail::pipeline_check_is_producer(params_.role);
     full_barrier_ptr_[stage].arrive(params_.dst_blockid);
   }
 
@@ -790,6 +904,7 @@ private:
 
   CUTLASS_DEVICE
   ConsumerToken consumer_try_wait(uint32_t stage, uint32_t phase, uint32_t skip_wait) {
+    detail::pipeline_check_is_consumer(params_.role);
     if (skip_wait) {
       return {BarrierStatus::WaitDone};
     }
@@ -799,6 +914,7 @@ private:
 
   CUTLASS_DEVICE
   ConsumerToken consumer_test_wait(uint32_t stage, uint32_t phase, uint32_t skip_wait) {
+    detail::pipeline_check_is_consumer(params_.role);
     if (skip_wait) {
       return {BarrierStatus::WaitDone};
     }
@@ -808,6 +924,7 @@ private:
 
   CUTLASS_DEVICE
   void consumer_wait(uint32_t stage, uint32_t phase, ConsumerToken barrier_token) {
+    detail::pipeline_check_is_consumer(params_.role);
     if (barrier_token == BarrierStatus::WaitAgain) {
       full_barrier_ptr_[stage].wait(phase);
     }
@@ -815,6 +932,7 @@ private:
 
   CUTLASS_DEVICE
   void consumer_release(uint32_t stage, uint32_t skip = false) {
+    detail::pipeline_check_is_consumer(params_.role);
     empty_barrier_ptr_[stage].arrive(params_.dst_blockid, (not skip));
   }
 };
@@ -841,7 +959,7 @@ namespace PipelineDetail {
 
 template <int Stages_>
 class PipelineAsync {
-public :
+public:
   static constexpr uint32_t Stages = Stages_;
   using SharedStorage = PipelineDetail::PipelineAsyncSharedStorage<Stages>;
   using FullBarrier = typename SharedStorage::FullBarrier;
@@ -864,33 +982,46 @@ public :
     uint32_t dst_blockid = cute::block_rank_in_cluster();
   };
 
-  // Default assumption when only storage is passed is :
-  // => single producer, single consumer & they are in the same block (within the Cluster)
+  static
   CUTLASS_DEVICE
-  PipelineAsync(SharedStorage& storage)
-    : PipelineAsync(storage, {}) {}
+  void
+  init_barriers(SharedStorage& storage, Params params) {
+    int warp_idx = canonical_warp_idx_sync();
+    bool is_initializing_warp = (warp_idx == 0);
+    if (is_initializing_warp) {
+      // Barrier FULL and EMPTY init
+      cutlass::arch::detail::initialize_barrier_array_pair_aligned<decltype(storage.full_barrier_), decltype(storage.empty_barrier_), Stages>(
+          storage.full_barrier_, storage.empty_barrier_, params.producer_arv_count, params.consumer_arv_count);
+    }
+  }
+
+  template<class InitBarriers>
+  CUTLASS_DEVICE
+  PipelineAsync(
+    SharedStorage& storage,
+    Params const& params,
+    InitBarriers = {}) :
+      params_(params),
+      full_barrier_ptr_(&storage.full_barrier_[0]),
+      empty_barrier_ptr_(&storage.empty_barrier_[0]) {
+
+    static_assert(cute::is_same_v<InitBarriers, cute::true_type> || cute::is_same_v<InitBarriers, cute::false_type>);
+    if constexpr (cute::is_same_v<InitBarriers, cute::true_type>) {
+      init_barriers(storage, params_);
+    }
+  }
 
   CUTLASS_DEVICE
   PipelineAsync(
     SharedStorage& storage,
     Params const& params) :
-      params_(params),
-      full_barrier_ptr_(&storage.full_barrier_[0]),
-      empty_barrier_ptr_(&storage.empty_barrier_[0]) {
+      PipelineAsync(storage, params, cute::true_type{}) { }
 
-    int warp_idx = canonical_warp_idx_sync();
-    int lane_predicate = cute::elect_one_sync();
-
-    // Barrier FULL, EMPTY init
-    // Init is done only by thread 0 of the block
-    if (warp_idx == 0 && lane_predicate == 1) {
-      for (int i = 0; i < Stages; ++i) {
-        full_barrier_ptr_[i].init(params.producer_arv_count);
-        empty_barrier_ptr_[i].init(params.consumer_arv_count);
-      }
-    }
-    cutlass::arch::fence_barrier_init();
-  }
+  // Default assumption when only storage is passed is :
+  // => single producer, single consumer & they are in the same block (within the Cluster)
+  CUTLASS_DEVICE
+  PipelineAsync(SharedStorage& storage)
+    : PipelineAsync(storage, {}, cute::true_type{}) {}
 
   ////////////////////
   // Producer APIs
@@ -983,6 +1114,7 @@ private:
 
   CUTLASS_DEVICE
   ProducerToken producer_try_acquire(uint32_t stage, uint32_t phase, uint32_t skip_wait) {
+    detail::pipeline_check_is_producer(params_.role);
     if (skip_wait) {
       return {BarrierStatus::WaitDone};
     }
@@ -992,6 +1124,7 @@ private:
 
   CUTLASS_DEVICE
   void producer_acquire(uint32_t stage, uint32_t phase, ProducerToken barrier_token) {
+    detail::pipeline_check_is_producer(params_.role);
     if (barrier_token == BarrierStatus::WaitAgain) {
       empty_barrier_ptr_[stage].wait(phase);
     }
@@ -999,11 +1132,13 @@ private:
 
   CUTLASS_DEVICE
   void producer_commit(uint32_t stage) {
+    detail::pipeline_check_is_producer(params_.role);
     full_barrier_ptr_[stage].arrive();
   }
 
   CUTLASS_DEVICE
   ConsumerToken consumer_try_wait(uint32_t stage, uint32_t phase, uint32_t skip_wait) {
+    detail::pipeline_check_is_consumer(params_.role);
     if (skip_wait) {
       return {BarrierStatus::WaitDone};
     }
@@ -1013,6 +1148,7 @@ private:
 
   CUTLASS_DEVICE
   ConsumerToken consumer_test_wait(uint32_t stage, uint32_t phase, uint32_t skip_wait) {
+    detail::pipeline_check_is_consumer(params_.role);
     if (skip_wait) {
       return {BarrierStatus::WaitDone};
     }
@@ -1022,6 +1158,7 @@ private:
 
   CUTLASS_DEVICE
   void consumer_wait(uint32_t stage, uint32_t phase) {
+    detail::pipeline_check_is_consumer(params_.role);
     bool done = full_barrier_ptr_[stage].test_wait(phase);
     if (!done) {
       full_barrier_ptr_[stage].wait(phase);
@@ -1030,6 +1167,7 @@ private:
 
   CUTLASS_DEVICE
   void consumer_wait(uint32_t stage, uint32_t phase, ConsumerToken barrier_token) {
+    detail::pipeline_check_is_consumer(params_.role);
     if (barrier_token == BarrierStatus::WaitAgain) {
       full_barrier_ptr_[stage].wait(phase);
     }
@@ -1037,6 +1175,7 @@ private:
 
   CUTLASS_DEVICE
   void consumer_release(uint32_t stage) {
+    detail::pipeline_check_is_consumer(params_.role);
     empty_barrier_ptr_[stage].arrive(params_.dst_blockid);
   }
 };
@@ -1075,7 +1214,7 @@ public:
     uint32_t group_size;
   };
 
-private :
+private:
   // In future this Params object can be replaced easily with a CG object
   Params params_;
   Barrier *barrier_ptr_;
@@ -1110,7 +1249,6 @@ public:
         }
       }
     }
-    cutlass::arch::fence_barrier_init();
   }
 
   // Wait on a stage to be unlocked

include/cutlass/platform/platform.h

@@ -106,7 +106,11 @@
 #include <cuda/std/cstdint>
 #include <cuda/std/limits>
 #else
-#include <stdint.h>
+#include <type_traits>
+#include <utility>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
 #endif
 
 #if !defined(__CUDACC_RTC__)
@@ -134,6 +138,10 @@
 #define CUTLASS_OS_WINDOWS
 #endif
 
+#if defined(__clang__) && defined(__CUDA__)
+#define CUTLASS_CLANG_CUDA 1
+#endif
+
 /******************************************************************************
  * Macros
  ******************************************************************************/
@@ -298,30 +306,13 @@ namespace platform {
 
 #if defined(__CUDACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1500))
 
-/// std::integral_constant
-template <typename value_t, value_t V>
-struct integral_constant;
-
-/// std::integral_constant
-template <typename value_t, value_t V>
-struct integral_constant {
-  static const value_t value = V;
-
-  typedef value_t value_type;
-  typedef integral_constant<value_t, V> type;
-
-  CUTLASS_HOST_DEVICE operator value_type() const { return value; }
-
-  CUTLASS_HOST_DEVICE const value_type operator()() const { return value; }
-};
-
 #else
 
-using std::integral_constant;
 using std::pair;
 
 #endif
 
+using CUTLASS_STL_NAMESPACE::integral_constant;
 using CUTLASS_STL_NAMESPACE::bool_constant;
 using CUTLASS_STL_NAMESPACE::true_type;
 using CUTLASS_STL_NAMESPACE::false_type;

include/cutlass/predicate_vector.h

@@ -35,15 +35,14 @@
 #pragma once
 
 #if defined(__CUDACC_RTC__)
-#include <cuda/std/cassert>
 #include <cuda/std/cstdint>
 #else
-#include <assert.h>
-#include <stdint.h>
+#include <cstdint>
 #endif
 
-#include "cutlass/cutlass.h"
+#include <cuda/std/cassert>
 
+#include "cutlass/cutlass.h"
 #include "cutlass/platform/platform.h"
 
 namespace cutlass {

include/cutlass/real.h

@@ -35,6 +35,8 @@
 
 #pragma once
 
+#include <cutlass/detail/helper_macros.hpp> // CUTLASS_DEVICE
+
 namespace cutlass {
 
 /// Used to determine the real-valued underlying type of a numeric type T.

include/cutlass/reduction/thread/reduction_operators.h

@@ -172,7 +172,7 @@ struct ReduceArrayOperation<logical_and<uint1b_t>, uint1b_t, N> {
       item = (item || !bits);
     }
 
-    return uint1b_t(!item);
+    return uint1b_t{!item};
   }
 };
 
@@ -195,7 +195,7 @@ struct ReduceArrayOperation<logical_or<uint1b_t>, uint1b_t, N> {
       item = (item || bits);
     }
 
-    return uint1b_t(item);
+    return uint1b_t{item};
   }
 };