CUTLASS 3.4.0 (#1286)

* CUTLASS 3.4.0

* Update CHANGELOG.md

---------

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

test/unit/gemm/device/default_gemm_configuration.hpp

@@ -169,7 +169,7 @@ struct DefaultGemmConfigurationToCutlass3Types<
   using TiledMma = TiledMMA<
       MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
       Layout<Shape<_2,_2,_1>>,  // 2x2x1 thread group
-      Layout<Shape<_1,_2,_1>>>; // 1x2x1 value group for 16x16x16 MMA and LDSM
+      Tile<_32,_32,_16>>;       // 32x32x16 MMA for LDSM, 1x2x1 value group
 
   // A
   static constexpr int kAlignmentA = 8;
@@ -301,7 +301,7 @@ struct DefaultGemmConfigurationToCutlass3Types<
   using TiledMma = TiledMMA<
       MMA_Atom<SM80_16x8x8_F32TF32TF32F32_TN>,
       Layout<Shape<_2,_2,_1>, Stride<_2, _1, _1>>, // 2x2x1 thread group
-      Layout<Shape<_1,_2,_1>>>;                    // 1x2x1 value group for 16x16x8 and LDSM
+      Tile<_32,_32,_8>>;                           // 32x32x8 MMA for LDSM, 1x2x1 value group
 
   // A
   static constexpr int kAlignmentA = 4;
@@ -352,7 +352,7 @@ struct DefaultGemmConfigurationToCutlass3Types<
   using TiledMma = TiledMMA<
       MMA_Atom<SM80_16x8x32_S32S8S8S32_TN>,
       Layout<Shape<_2,_2,_1>>,   // 2x2x1 thread group
-      Layout<Shape<_1,_2,_1>>>;  // 1x2x1 value group for 16x16x32 and LDSM
+      Tile<_32,_32,_32>>;        // 16x16x32 MMA for LDSM, 1x2x1 value group
 
   // A (M,K)  K-major
   using SmemLayoutAtomA = decltype(
@@ -798,9 +798,9 @@ struct DefaultGemmConfigurationToCutlass3Types<
   using DispatchPolicy = MainloopSm80CpAsync<3>;
   using TiledMma = TiledMMA<
       MMA_Atom<UniversalFMA<ElementAccumulator, ElementA, ElementB, ElementC>>,
-      Layout<Shape<_16, _16, _1>>,
-      Layout<Shape< _2,  _2, _1>>,
-      Tile<Layout<_2,_16>,Layout<_2,_16>,Underscore>>;
+      Layout<Shape<_16, _16, _1>>,                            // 16x16x1 thread group
+      Tile<Layout<Shape<_16,_2>,Stride<_2,_1>>,               // 32x32x1 MMA with perm for load vectorization
+           Layout<Shape<_16,_2>,Stride<_2,_1>>,Underscore>>;
 
   // A (M,K)  M-major
   using SmemLayoutAtomA = Layout<Shape<_128,_16>>;
@@ -920,9 +920,8 @@ struct DefaultGemmConfigurationToCutlass3Types<
   using DispatchPolicy = MainloopSm80CpAsync<3>;
   using TiledMma = TiledMMA<
       MMA_Atom<UniversalFMA<ElementAccumulator, ElementA, ElementB, ElementC>>,
-      Layout<Shape<_16, _16, _1>>,
-      Layout<Shape< _2,  _1, _1>>,
-      Tile<Layout<_2,_16>,Underscore,Underscore>>;
+      Layout<Shape<_16, _16, _1>>,                                      // 16x16x1 thread group
+      Tile<Layout<Shape<_16,_2>,Stride<_2,_1>>,Underscore,Underscore>>; // 32x16x1 MMA with perm for load vectorization
 
   // A (M,K)  M-major
   using SmemLayoutAtomA = Layout<Shape<_128,_16>>;
@@ -982,9 +981,8 @@ struct DefaultGemmConfigurationToCutlass3Types<
   using DispatchPolicy = MainloopSm80CpAsync<3>;
   using TiledMma = TiledMMA<
       MMA_Atom<UniversalFMA<ElementAccumulator, ElementA, ElementB, ElementC>>,
-      Layout<Shape<_16, _16, _1>>,
-      Layout<Shape< _1,  _2, _1>>,
-      Tile<Underscore,Layout<_2,_16>,Underscore>>;
+      Layout<Shape<_16, _16, _1>>,                                      // 16x16x1 thread group
+      Tile<Underscore,Layout<Shape<_16,_2>,Stride<_2,_1>>,Underscore>>; // 16x32x1 MMA with perm for load vectorization
 
   // A (M,K)  K-major
   using SmemLayoutAtomA = Layout<Shape <_128,          _16>,
@@ -1041,8 +1039,9 @@ struct DefaultGemmConfigurationToCutlass3Types<
   using TiledMma = TiledMMA<
       MMA_Atom<SM80_8x8x4_F64F64F64F64_TN>,            // Atom
       Layout<Shape<_2,_2,_1>>,                         // Atom layout
-      Layout<Shape<_2,_2,_1>>,                         // Val layout
-      Tile<Layout<_2,_16>,Layout<_2,_16>,Underscore>>; // Mode permutations
+      Tile<Layout<Shape<_16,_2>,Stride<_2,_1>>,        // 32x32x4 MMA with perm for load vectorization
+           Layout<Shape<_16,_2>,Stride<_2,_1>>,
+           Underscore>>;
 
   // A  (M,K)  K-Major
   using SmemLayoutAtomA = decltype(
@@ -1119,8 +1118,9 @@ struct DefaultGemmConfigurationToCutlass3Types<
   using TiledMma = TiledMMA<
       MMA_Atom<SM80_8x8x4_F64F64F64F64_TN>,            // Atom
       Layout<Shape<_2,_2,_1>>,                         // Atom layout
-      Layout<Shape<_2,_2,_1>>,                         // Val layout
-      Tile<Layout<_2,_16>,Layout<_2,_16>,Underscore>>; // Mode permutations
+      Tile<Layout<Shape<_16,_2>,Stride<_2,_1>>,        // 32x32x4 MMA with perm for load vectorization
+           Layout<Shape<_16,_2>,Stride<_2,_1>>,
+           Underscore>>;
 
   // A  (M,K)  M-Major
   using SmemLayoutAtomA = decltype(
@@ -1183,8 +1183,9 @@ struct DefaultGemmConfigurationToCutlass3Types<
   using TiledMma = TiledMMA<
       MMA_Atom<SM80_8x8x4_F64F64F64F64_TN>,            // Atom
       Layout<Shape<_2,_2,_1>>,                         // Atom layout
-      Layout<Shape<_2,_2,_1>>,                         // Val layout
-      Tile<Layout<_2,_16>,Layout<_2,_16>,Underscore>>; // Mode permutations
+      Tile<Layout<Shape<_16,_2>,Stride<_2,_1>>,        // 32x32x4 MMA with perm for load vectorization
+           Layout<Shape<_16,_2>,Stride<_2,_1>>,
+           Underscore>>;
 
   // A  (M,K)  M-Major
   using SmemLayoutAtomA = decltype(
@@ -1247,8 +1248,9 @@ struct DefaultGemmConfigurationToCutlass3Types<
   using TiledMma = TiledMMA<
       MMA_Atom<SM80_8x8x4_F64F64F64F64_TN>,            // Atom
       Layout<Shape<_2,_2,_1>>,                         // Atom layout
-      Layout<Shape<_2,_2,_1>>,                         // Val layout
-      Tile<Layout<_2,_16>,Layout<_2,_16>,Underscore>>; // Mode permutations
+      Tile<Layout<Shape<_16,_2>,Stride<_2,_1>>,        // 32x32x4 MMA with perm for load vectorization
+           Layout<Shape<_16,_2>,Stride<_2,_1>>,
+           Underscore>>;
 
   // A  (M,K)  K-Major
   using SmemLayoutAtomA = decltype(

test/unit/gemm/device/gemm_testbed_3x.hpp

@@ -58,6 +58,7 @@
 #include "cutlass/fast_math.h"
 #include "cutlass/platform/platform.h"
 #include "cutlass/epilogue/fusion/operations.hpp"
+#include "cutlass/gemm/kernel/tile_scheduler_params.h"
 
 #include "cute/int_tuple.hpp"
 #include "cute/layout.hpp"
@@ -192,6 +193,7 @@ struct TestbedImpl {
   using ActivationFunctor = ActivationFunctor_<ElementCompute>;
 
   using RasterOrderOptions = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90::RasterOrderOptions;
+  using DecompositionMode = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90StreamKParams::DecompositionMode;
 
   static_assert(cute::rank(StrideC{}) == 3, "StrideCD must be rank-3: [M, N, L]");
   static_assert(cute::rank(StrideD{}) == 3, "StrideCD must be rank-3: [M, N, L]");
@@ -248,6 +250,10 @@ struct TestbedImpl {
 
   // Used to force multi-wave tests for persistent kernel schedules
   constexpr static int MaxSmCount = 16;
+
+  cutlass::ComplexTransform TransformA = Gemm::kTransformA;
+  cutlass::ComplexTransform TransformB = Gemm::kTransformB;
+
   //
   // Methods
   //
@@ -462,7 +468,7 @@ struct TestbedImpl {
     auto Vbeta = cute::make_tensor(static_cast<ElementCompute*>(nullptr),
         cute::make_layout(cute::make_shape(M, cute::_1{})));
 
-    cutlass::reference::host::GettMainloopParams<ElementAccumulator, decltype(A), decltype(B)> mainloop_params{A, B};
+    cutlass::reference::host::GettMainloopParams<ElementAccumulator, decltype(A), decltype(B)> mainloop_params{A, B, TransformA, TransformB};
 
     cutlass::reference::host::GettEpilogueParams<
         ElementScalar,
@@ -523,6 +529,9 @@ struct TestbedImpl {
     Gemm& gemm_op,
     typename Gemm::Arguments& arguments,
     cutlass::device_memory::allocation<uint8_t>& workspace) {
+    int M = cute::size<0>(problem_size);
+    int N = cute::size<1>(problem_size);
+    int K = cute::size<2>(problem_size);
     int L = 1;
     if constexpr(cute::rank(ProblemShapeType{}) == 4) {
       L = cute::size<3>(problem_size);
@@ -561,7 +570,8 @@ struct TestbedImpl {
     detail::Iterations iterations = detail::Iterations{},
     RasterOrderOptions raster_order = RasterOrderOptions::Heuristic,
     detail::MaxSwizzleSize max_swizzle = detail::MaxSwizzleSize{},
-    detail::Splits splits = detail::Splits{})
+    detail::Splits splits = detail::Splits{},
+    DecompositionMode decomposition_mode = DecompositionMode::Heuristic)
   {
     // Fail test if insufficient CUDA device
     if (!sufficient()) {
@@ -586,14 +596,6 @@ struct TestbedImpl {
       hw_info.sm_count = this->sm_count;
     }
 
-    typename Gemm::GemmKernel::TileScheduler::Arguments scheduler_args;
-    if constexpr (std::is_same_v<typename Gemm::GemmKernel::TileSchedulerTag, cutlass::gemm::StreamKScheduler>) {
-      scheduler_args = { static_cast<int>(splits), static_cast<int>(max_swizzle), raster_order };
-    }
-    else {
-      scheduler_args = { static_cast<int>(max_swizzle), raster_order };
-    }
-
     // DefaultEpilogue
     auto arguments = typename Gemm::Arguments {
       cutlass::gemm::GemmUniversalMode::kGemm,
@@ -606,10 +608,20 @@ struct TestbedImpl {
         {alpha, beta},
         tensor_C.device_data(), stride_c, tensor_D.device_data(), stride_d
       },
-      hw_info,
-      scheduler_args
+      hw_info
     };
 
+    if constexpr (std::is_same_v<typename Gemm::GemmKernel::TileSchedulerTag, cutlass::gemm::StreamKScheduler>) {
+      arguments.scheduler.splits = static_cast<int>(splits);
+      arguments.scheduler.max_swizzle_size = static_cast<int>(max_swizzle);
+      arguments.scheduler.raster_order = raster_order;
+      arguments.scheduler.decomposition_mode = decomposition_mode;
+
+    } else {
+      arguments.scheduler.max_swizzle_size = static_cast<int>(max_swizzle);
+      arguments.scheduler.raster_order = raster_order;
+    }
+
     Gemm gemm_op;
 
     size_t workspace_size = Gemm::get_workspace_size(arguments);
@@ -683,6 +695,7 @@ struct Testbed3x {
   using LayoutTagD = typename TestBedImpl::LayoutTagD;
 
   using RasterOrderOptions = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90::RasterOrderOptions;
+  using DecompositionMode = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90StreamKParams::DecompositionMode;
 
   // Detail Implementation
   TestBedImpl impl_;
@@ -723,11 +736,12 @@ struct Testbed3x {
     RasterOrderOptions raster_order = RasterOrderOptions::Heuristic,
     detail::MaxSwizzleSize max_swizzle = detail::MaxSwizzleSize{},
     detail::Splits splits = detail::Splits{},
+    DecompositionMode decomposition_mode = DecompositionMode::Heuristic,
     bool profiling = false,
     detail::Iterations iterations = detail::Iterations{})
   {
     return impl_.run(
-        problem_size, alpha, beta, profiling, iterations, raster_order, max_swizzle, splits
+        problem_size, alpha, beta, profiling, iterations, raster_order, max_swizzle, splits, decomposition_mode
         );
   }
 };
@@ -768,6 +782,7 @@ struct Testbed3xFusionOperation {
   static_assert(cute::is_base_of_v<cutlass::epilogue::fusion::FusionOperation, FusionOp>);
 
   using RasterOrderOptions = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90::RasterOrderOptions;
+  using DecompositionMode = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90StreamKParams::DecompositionMode;
 
   // fusion types are potentially void if the fusion is not supported
   // helper so we don't try to construct HostTensor with void type
@@ -818,6 +833,7 @@ struct Testbed3xFusionOperation {
   cutlass::HostTensor<ElementAmax, LayoutTagScalar> abs_max_D;
   cutlass::HostTensor<ElementAux , LayoutTagAux   > tensor_Aux;
   cutlass::gemm::TagToStrideC_t<   LayoutTagAux   > stride_Aux;
+
   // References
   cutlass::HostTensor<ElementBias, LayoutTagVector> reference_dbias;
   cutlass::HostTensor<ElementAux , LayoutTagAux   > reference_Aux;
@@ -977,7 +993,6 @@ struct Testbed3xFusionOperation {
         cutlass::reference::host::TensorFill(reference_abs_max_Aux.host_view(), ElementAmax(0));
       }
     }
-
   }
 
   template <
@@ -1219,6 +1234,7 @@ struct Testbed3xFusionOperation {
     RasterOrderOptions raster_order = RasterOrderOptions::Heuristic,
     detail::MaxSwizzleSize max_swizzle = detail::MaxSwizzleSize{},
     detail::Splits splits = detail::Splits{},
+    DecompositionMode decomposition_mode = DecompositionMode::Heuristic,
     bool profiling = false,
     detail::Iterations iterations = detail::Iterations{})
   {
@@ -1234,7 +1250,7 @@ struct Testbed3xFusionOperation {
     typename Gemm::Arguments arguments;
     cutlass::KernelHardwareInfo hw_info;
     cudaDeviceProp prop;
-    
+
     hw_info.device_id = 0;
     if (not profiling) {
       impl_.sm_count = std::min(impl_.MaxSmCount, cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id));
@@ -1251,11 +1267,6 @@ struct Testbed3xFusionOperation {
     /// A/B/C/D Tensor
     initialize(problem_size, alpha_, beta_);
 
-    typename Gemm::GemmKernel::TileScheduler::Arguments scheduler_args;
-    if constexpr (std::is_same_v<typename Gemm::GemmKernel::TileSchedulerTag, cutlass::gemm::StreamKScheduler>) {
-      scheduler_args = { static_cast<int>(splits) };
-    }
-
     arguments = typename Gemm::Arguments{
       cutlass::gemm::GemmUniversalMode::kGemm,
       problem_size,
@@ -1270,10 +1281,19 @@ struct Testbed3xFusionOperation {
         impl_.tensor_D.device_data(),
         impl_.stride_d
       }, // Epilogue arguments end
-      hw_info,
-      scheduler_args
+      hw_info
     };
-    
+
+    if constexpr (std::is_same_v<typename Gemm::GemmKernel::TileSchedulerTag, cutlass::gemm::StreamKScheduler>) {
+      arguments.scheduler.splits = static_cast<int>(splits);
+      arguments.scheduler.max_swizzle_size = static_cast<int>(max_swizzle);
+      arguments.scheduler.raster_order = raster_order;
+      arguments.scheduler.decomposition_mode = decomposition_mode;
+    } else {
+      arguments.scheduler.max_swizzle_size = static_cast<int>(max_swizzle);
+      arguments.scheduler.raster_order = raster_order;
+    }
+
     auto coord_0 = cutlass::make_Coord(0);
     if constexpr (IsLegacy) {
       arguments.epilogue.thread = {
@@ -1313,12 +1333,18 @@ struct Testbed3xFusionOperation {
       }
 
       // example of how to set kernel activation arguments
+      // see ActivationFunctor::Arguments in activation.h for definition
+      // if Arguments doesn't exist then fusion_args.activation is empty
       if constexpr (cute::is_same_v<ActivationFunctor, cutlass::epilogue::thread::ScaledGELU_taylor<ElementCompute>>) {
-        // see ActivationFunctor::Arguments in activation.h for definition
-        // if Arguments doesn't exist then fusion_args.activation is empty
         fusion_args.activation.scale = ElementCompute(1);
       }
 
+      // Treat Clamp as ReLU
+      if constexpr (cute::is_same_v<ActivationFunctor, cutlass::epilogue::thread::Clamp<ElementCompute>>) {
+        fusion_args.activation.lower_bound = 0;
+        fusion_args.activation.upper_bound = std::numeric_limits<ElementCompute>::max();
+      }
+
       if constexpr (IsAbsMaxEnabledD) {
         fusion_args.amax_D_ptr = abs_max_D.device_data();
       }
@@ -1381,7 +1407,6 @@ struct Testbed3xFusionOperation {
         std::cout << "Error : Failed : with alpha: " << float(alpha_) << ", beta: " << float(beta_)
                   << "\n";
       }
-
       return passed;
     }
   }
@@ -1413,13 +1438,21 @@ bool TestAll(double alpha = 1.0, double beta = 0.0, Testbed testbed = {}) {
 
   std::vector<int> problem_size_k = {max_alignment, TileShapeK * (Stages + 1) - max_alignment};
 
+  using DecompositionMode = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90StreamKParams::DecompositionMode;
+  std::vector<DecompositionMode> decomposition_modes = {DecompositionMode::Heuristic};
   std::vector<int> problem_splits = {1};
-  if constexpr (std::is_same_v<typename Gemm::GemmKernel::TileSchedulerTag, cutlass::gemm::StreamKScheduler>) {
+  static constexpr bool UsesStreamKScheduler = std::is_same_v<typename Gemm::GemmKernel::TileSchedulerTag, cutlass::gemm::StreamKScheduler>;
+  if constexpr (UsesStreamKScheduler) {
     problem_splits.push_back(2);
     problem_splits.push_back(3);
 
-    // As many splits as there are maximum k tiles
-    problem_splits.push_back(Stages + 1);
+    decomposition_modes.push_back(DecompositionMode::DataParallel);
+    decomposition_modes.push_back(DecompositionMode::SplitK);
+    decomposition_modes.push_back(DecompositionMode::StreamK);
+
+    // Use larger K sizes for stream-K tests
+    static constexpr int min_tiles_per_sk_unit = cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90StreamKParams::min_iters_per_sk_unit_;
+    problem_size_k = {TileShapeK * min_tiles_per_sk_unit, TileShapeK * 3 * min_tiles_per_sk_unit - max_alignment};
   }
 
   using RasterOrderOptions = typename cutlass::gemm::kernel::detail::PersistentTileSchedulerSm90::RasterOrderOptions;
@@ -1433,33 +1466,53 @@ bool TestAll(double alpha = 1.0, double beta = 0.0, Testbed testbed = {}) {
       for (int k : problem_size_k) {
         for (auto raster_order : raster_orders) {
           for (int max_swizzle_size : max_swizzle_sizes) {
-            for (int splits : problem_splits) {
-              ProblemShapeType problem_size;
-              if constexpr (cute::rank(ProblemShapeType{}) == 4) {
-                problem_size = ProblemShapeType{m, n, k, /* l */ 1};
-              }
-              else {
-                problem_size = ProblemShapeType{m, n, k};
-              }
+            for (DecompositionMode decomp_mode : decomposition_modes) {
 
-              passed = testbed.run(
-                problem_size,
-                cutlass::from_real<ElementScalar>(alpha),
-                cutlass::from_real<ElementScalar>(beta),
-                raster_order,
-                detail::MaxSwizzleSize(max_swizzle_size),
-                detail::Splits(splits)
-              );
+              std::vector<int> problem_splits = {1};
+              if (UsesStreamKScheduler && (decomp_mode == DecompositionMode::Heuristic || decomp_mode == DecompositionMode::SplitK)) {
+                auto max_splits = (k + TileShapeK - 1) / TileShapeK;
+                if (max_splits > 2) {
+                  problem_splits.push_back(2);
+                }
+                if (max_splits > 3) {
+                  problem_splits.push_back(3);
+                }
 
-              if (!passed) {
-                return false;
+                problem_splits.push_back(max_splits);
+
+                // Test the case in which we ask for more splits than there are K tiles in the GEMM. In this
+                // case, split-K will fall back to a splitting factor of `max_splits`.
+                problem_splits.push_back(max_splits + 1);
               }
-            }
-          }
-        }
-      }
-    }
-  }
+              for (int splits : problem_splits) {
+                ProblemShapeType problem_size;
+                if constexpr (cute::rank(ProblemShapeType{}) == 4) {
+                  problem_size = ProblemShapeType{m, n, k, /* l */ 1};
+                }
+                else {
+                  problem_size = ProblemShapeType{m, n, k};
+                }
+
+                passed = testbed.run(
+                  problem_size,
+                  cutlass::from_real<ElementScalar>(alpha),
+                  cutlass::from_real<ElementScalar>(beta),
+                  raster_order,
+                  detail::MaxSwizzleSize(max_swizzle_size),
+                  detail::Splits(splits),
+                  decomp_mode
+                );
+
+                if (!passed) {
+                  return false;
+                }
+              } // splits
+            } // decomposition_mode
+          } // max_swizzle_size
+        } // raster_order
+      } // k
+    } // n
+  } // m
 
   // if we do support batched GEMM, just run one test on it to save on test time
   if constexpr (cute::rank(ProblemShapeType{}) == 4) {

test/unit/gemm/device/gemm_testbed_3x_evt.hpp

@@ -382,7 +382,7 @@ public:
   HostAuxLoad(){}
   template<typename ProblemShapeType>
   HostAuxLoad(ProblemShapeType problem_size, TestBedImpl impl, bool check_relative_equality=false)
-    : Base(check_relative_equality), impl_(impl){
+    : Base(check_relative_equality), impl_(impl) {
     auto problem_shape_NMKL = cute::append<4>(problem_size, 1);
     auto [_M, _N, K, _L] = problem_shape_NMKL;
     auto aux_coord = cutlass::make_Coord(_M * _L, _N);

test/unit/gemm/device/sm90_evt_operations.hpp

@@ -267,7 +267,7 @@ template<
 using Sm90LinCombAuxLoad =
   Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementOutput, ElementCompute, RoundStyle>, // beta * C + (alpha * acc + bias)
     Sm90ScalarBroadcast<ElementScalar>, // beta
-    Sm90SrcFetch, // C
+    Sm90SrcFetch<ElementOutput>, // C
     Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementCompute, ElementCompute, RoundStyle>, // alpha * acc + bias
       Sm90ScalarBroadcast<ElementScalar>, // alpha
       Sm90AccFetch, // acc
@@ -295,7 +295,7 @@ template<
 using Sm90LinCombEVTDAG =
   Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementOutput, ElementCompute, RoundStyle>, // beta * C + (alpha * acc + aux)
     Sm90ScalarBroadcast<ElementScalar>, // beta
-    Sm90SrcFetch, // C
+    Sm90SrcFetch<ElementOutput>, // C
     Sm90TopologicalVisitor<
       ElementCompute,
       cute::tuple<
@@ -349,7 +349,7 @@ using Sm90LinCombDAGEVT =
       Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementCompute, ElementCompute, RoundStyle>,
         Sm90ScalarBroadcast<ElementScalar>,
         Sm90AccFetch,
-        Sm90SrcFetch
+        Sm90SrcFetch<ElementOutput>
       >
     >,
     Sm90ColBroadcast<0, typename EpilogueDescriptor::TileShape, ElementBias>,
@@ -371,7 +371,7 @@ template<
 using Sm90LinCombPerColumnBias =
   Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementOutput, ElementCompute, RoundStyle>, // beta * C + (alpha * acc + bias)
     Sm90ScalarBroadcast<ElementScalar>, // beta
-    Sm90SrcFetch, // C
+    Sm90SrcFetch<ElementOutput>, // C
     Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementCompute, ElementCompute, RoundStyle>, // alpha * acc + bias
       Sm90ScalarBroadcast<ElementScalar>, // alpha
       Sm90AccFetch, // acc
@@ -403,7 +403,7 @@ using Sm90LinCombPerColumnReduce =
   Sm90EVT<Sm90RowReduction<RegReduceFn, GmemReduceFn, 0, CtaTileShapeMNK, ElementReduce, ElementCompute, RoundStyle>, // per column reduce
     Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementOutput, ElementCompute, RoundStyle>, // beta * C + alpha * acc
       Sm90ScalarBroadcast<ElementScalar>, // beta
-      Sm90SrcFetch, // C
+      Sm90SrcFetch<ElementOutput>, // C
       Sm90EVT<Sm90Compute<multiplies, ElementCompute, ElementCompute, RoundStyle>, // alpha * acc
         Sm90ScalarBroadcast<ElementScalar>, // alpha
         Sm90AccFetch // acc
@@ -428,7 +428,7 @@ using Sm90LinCombPerRowReduce =
   Sm90EVT<Sm90ColReduction<RegReduceFn, GmemReduceFn, 0, CtaTileShapeMNK, ElementReduce, ElementCompute, RoundStyle>, // per column reduce
     Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementOutput, ElementCompute, RoundStyle>, // beta * C + alpha * acc
       Sm90ScalarBroadcast<ElementScalar>, // beta
-      Sm90SrcFetch, // C
+      Sm90SrcFetch<ElementOutput>, // C
       Sm90EVT<Sm90Compute<multiplies, ElementCompute, ElementCompute, RoundStyle>, // alpha * acc
         Sm90ScalarBroadcast<ElementScalar>, // alpha
         Sm90AccFetch // acc
@@ -452,7 +452,7 @@ using Sm90LinCombScalarReduce =
   Sm90EVT<Sm90ScalarReduction<RegReduceFn, GmemReduceFn, ElementReduce, ElementCompute, RoundStyle>, // per column reduce
     Sm90EVT<Sm90Compute<homogeneous_multiply_add, ElementOutput, ElementCompute, RoundStyle>, // beta * C + alpha * acc
       Sm90ScalarBroadcast<ElementScalar>, // beta
-      Sm90SrcFetch, // C
+      Sm90SrcFetch<ElementOutput>, // C
       Sm90EVT<Sm90Compute<multiplies, ElementCompute, ElementCompute, RoundStyle>, // alpha * acc
         Sm90ScalarBroadcast<ElementScalar>, // alpha
         Sm90AccFetch // acc

test/unit/gemm/device/sm90_gemm_f16_f16_f16_tensor_op_f32_cluster_warpspecialized_cooperative_bias_elementwise.cu

@@ -389,7 +389,7 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_cooperative_epilogue, 25
 
   using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative;
   using FusionOperation = cutlass::epilogue::fusion::LinCombPerRowBiasEltActAux<
-      LayoutC, cutlass::epilogue::thread::ReLu, cutlass::half_t, float, cutlass::half_t, float>;
+      LayoutC, cutlass::epilogue::thread::ReLu, cutlass::half_t, float, cutlass::half_t, float, void>;
 
   using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
       cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
@@ -434,7 +434,7 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_cooperative_epilogue, 25
 
   using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative;
   using FusionOperation = cutlass::epilogue::fusion::LinCombPerRowBiasEltActAux<
-      LayoutC, cutlass::epilogue::thread::ReLu, cutlass::half_t, float, cutlass::half_t, cutlass::half_t>;
+      LayoutC, cutlass::epilogue::thread::ReLu, cutlass::half_t, float, cutlass::half_t, cutlass::half_t, void>;
 
   using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
       cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
@@ -480,7 +480,7 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_cooperative_epilogue, 25
   using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative;
   // ReLU with uint1b_t aux will compute dReLU/dZ as the aux output, i.e. Aux(i) = (Z(i) >= 0) ? 1 : 0
   using FusionOperation = cutlass::epilogue::fusion::LinCombPerRowBiasEltActAux<
-      LayoutC, cutlass::epilogue::thread::ReLU, cutlass::half_t, float, cutlass::uint1b_t, int8_t>;
+      LayoutC, cutlass::epilogue::thread::ReLU, cutlass::half_t, float, cutlass::uint1b_t, int8_t, void>;
 
   using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
       cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
@@ -525,7 +525,7 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_cooperative_epilogue, 25
 
   using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative;
   using FusionOperation = cutlass::epilogue::fusion::LinCombDeEltActDePerRowBias<
-      LayoutC, cutlass::epilogue::thread::dReLU, cutlass::half_t, float, cutlass::uint1b_t, float>;
+      LayoutC, cutlass::epilogue::thread::dReLU, cutlass::half_t, float, cutlass::uint1b_t, float, void>;
 
   using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
       cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
@@ -570,7 +570,7 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_cooperative_epilogue, 25
 
   using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecializedCooperative;
   using FusionOperation = cutlass::epilogue::fusion::LinCombDeEltAct<
-      LayoutC, cutlass::epilogue::thread::dGELU, cutlass::half_t, float, cutlass::half_t>;
+      LayoutC, cutlass::epilogue::thread::dGELU, cutlass::half_t, float, cutlass::half_t, void>;
 
   using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
       cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,

test/unit/gemm/device/sm90_gemm_f16_f16_f16_tensor_op_f32_cluster_warpspecialized_pingpong_bias_elementwise.cu

@@ -335,7 +335,7 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_persistent_epilogue, 128
 
   using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized;
   using FusionOperation = cutlass::epilogue::fusion::LinCombPerRowBiasEltActAux<
-      LayoutC, cutlass::epilogue::thread::ReLu, cutlass::half_t, float, cutlass::half_t, float>;
+      LayoutC, cutlass::epilogue::thread::ReLu, cutlass::half_t, float, cutlass::half_t, float, void>;
 
   using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
       cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
@@ -380,7 +380,7 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_persistent_epilogue, 128
 
   using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized;
   using FusionOperation = cutlass::epilogue::fusion::LinCombPerRowBiasEltActAux<
-      LayoutC, cutlass::epilogue::thread::ReLu, cutlass::half_t, float, cutlass::half_t, cutlass::half_t>;
+      LayoutC, cutlass::epilogue::thread::ReLu, cutlass::half_t, float, cutlass::half_t, cutlass::half_t, void>;
 
   using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
       cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
@@ -425,7 +425,7 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_persistent_epilogue, 128
 
   using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized;
   using FusionOperation = cutlass::epilogue::fusion::LinCombPerRowBiasEltActAux<
-      LayoutC, cutlass::epilogue::thread::ReLu, cutlass::half_t, float, cutlass::half_t, int8_t>;
+      LayoutC, cutlass::epilogue::thread::ReLu, cutlass::half_t, float, cutlass::half_t, int8_t, void>;
 
   using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
       cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
@@ -470,7 +470,7 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_pingpong_epilogue, 128x1
 
   using EpilogueSchedule = cutlass::epilogue::TmaWarpSpecialized;
   using FusionOperation = cutlass::epilogue::fusion::LinCombDeEltActDePerRowBias<
-      LayoutC, cutlass::epilogue::thread::dReLU, cutlass::half_t, float, cutlass::uint1b_t, float>;
+      LayoutC, cutlass::epilogue::thread::dReLU, cutlass::half_t, float, cutlass::uint1b_t, float, void>;
 
   using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
       cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,

test/unit/gemm/device/sm90_gemm_f16_f16_f16_tensor_op_f32_cooperative_stream_k.cu

@@ -95,7 +95,8 @@ TEST(SM90_Device_Gemm_f16t_f16t_f32n_tensor_op_gmma_f32_cooperative_stream_k, 12
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 TEST(SM90_Device_Gemm_f16t_f16t_f32n_tensor_op_gmma_f32_cooperative_stream_k, 256x128x64_1x2x1) {
@@ -136,7 +137,8 @@ TEST(SM90_Device_Gemm_f16t_f16t_f32n_tensor_op_gmma_f32_cooperative_stream_k, 25
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -178,7 +180,8 @@ TEST(SM90_Device_Gemm_f16t_f16t_f32n_tensor_op_gmma_f32_cooperative_stream_k, 12
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -218,7 +221,8 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32n_tensor_op_gmma_f32_cooperative_stream_k, 25
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -258,7 +262,8 @@ TEST(SM90_Device_Gemm_f16n_f16t_f32n_tensor_op_gmma_f32_cooperative_stream_k, 12
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -298,7 +303,8 @@ TEST(SM90_Device_Gemm_f16n_f16n_f32n_tensor_op_gmma_f32_cooperative_stream_k, 25
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -341,7 +347,8 @@ TEST(SM90_Device_Gemm_f16t_f16t_f32n_tensor_op_gmma_f32_cooperative_stream_k, 12
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -381,7 +388,8 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32n_tensor_op_gmma_f32_cooperative_stream_k, 12
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -421,7 +429,8 @@ TEST(SM90_Device_Gemm_f16n_f16t_f32n_tensor_op_gmma_f32_cooperative_stream_k, 12
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -461,7 +470,8 @@ TEST(SM90_Device_Gemm_f16n_f16n_f32n_tensor_op_gmma_f32_cooperative_stream_k, 12
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 
@@ -505,7 +515,8 @@ TEST(SM90_Device_Gemm_f16t_f16t_f32n_tensor_op_gmma_f32_cooperative_stream_k, 12
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -545,7 +556,8 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32n_tensor_op_gmma_f32_cooperative_stream_k, 12
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -585,7 +597,8 @@ TEST(SM90_Device_Gemm_f16n_f16t_f32n_tensor_op_gmma_f32_cooperative_stream_k, 12
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -625,7 +638,8 @@ TEST(SM90_Device_Gemm_f16n_f16n_f32n_tensor_op_gmma_f32_cooperative_stream_k, 12
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 
@@ -669,7 +683,8 @@ TEST(SM90_Device_Gemm_f16t_f16t_f32n_tensor_op_gmma_f32_cooperative_stream_k, 25
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -709,7 +724,8 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32n_tensor_op_gmma_f32_cooperative_stream_k, 25
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -749,7 +765,8 @@ TEST(SM90_Device_Gemm_f16n_f16t_f32n_tensor_op_gmma_f32_cooperative_stream_k, 25
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -789,7 +806,8 @@ TEST(SM90_Device_Gemm_f16n_f16n_f32n_tensor_op_gmma_f32_cooperative_stream_k, 25
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 TEST(SM90_Device_Gemm_f16t_f16n_f16n_tensor_op_gmma_f32_cooperative_stream_k_epilogue, 256x128x64_2x2x1) {
@@ -827,7 +845,8 @@ TEST(SM90_Device_Gemm_f16t_f16n_f16n_tensor_op_gmma_f32_cooperative_stream_k_epi
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 TEST(SM90_Device_Gemm_f16t_f16n_f16t_tensor_op_gmma_f32_cooperative_stream_k_epilogue, 256x128x64_2x2x1) {
@@ -865,7 +884,8 @@ TEST(SM90_Device_Gemm_f16t_f16n_f16t_tensor_op_gmma_f32_cooperative_stream_k_epi
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 TEST(SM90_Device_Gemm_f16t_f16n_f32n_tensor_op_gmma_f32_cooperative_stream_k_epilogue, 128x128x64_2x2x1) {
@@ -903,7 +923,8 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32n_tensor_op_gmma_f32_cooperative_stream_k_epi
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_cooperative_stream_k_epilogue, 128x128x64_2x2x1) {
@@ -941,7 +962,8 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_cooperative_stream_k_epi
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAll<Gemm>(1.0, 1.0));
 }
 
 TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_cooperative_stream_k_epilogue, 256x128x64_2x2x1_BiasF32_ReLU) {
@@ -985,8 +1007,8 @@ TEST(SM90_Device_Gemm_f16t_f16n_f32t_tensor_op_gmma_f32_cooperative_stream_k_epi
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
 
-  bool passed = test::gemm::device::TestAllBiasElementwise<Gemm>();
-  EXPECT_TRUE(passed);
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 1.0));
 }
 
 #endif // defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)

test/unit/gemm/device/sm90_gemm_f8_f8_f32_tensor_op_f32_cooperative_stream_k.cu

@@ -99,7 +99,8 @@ TEST(SM90_Device_Gemm_e4m3t_e4m3n_f32n_tensor_op_gmma_f32_cooperative_stream_k,
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 1.0));
 }
 
 TEST(SM90_Device_Gemm_e4m3t_e4m3n_f32n_tensor_op_gmma_f32_cooperative_stream_k, 256x128x128_1x1x1) {
@@ -146,7 +147,8 @@ TEST(SM90_Device_Gemm_e4m3t_e4m3n_f32n_tensor_op_gmma_f32_cooperative_stream_k,
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 1.0));
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -197,7 +199,8 @@ TEST(SM90_Device_Gemm_e4m3t_e4m3n_f32n_tensor_op_gmma_f32_cooperative_stream_k,
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 1.0));
 }
 
 TEST(SM90_Device_Gemm_e4m3t_e4m3n_f32n_tensor_op_gmma_f32_cooperative_stream_k, 256x128x128_1x2x1) {
@@ -244,7 +247,8 @@ TEST(SM90_Device_Gemm_e4m3t_e4m3n_f32n_tensor_op_gmma_f32_cooperative_stream_k,
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 1.0));
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -295,7 +299,8 @@ TEST(SM90_Device_Gemm_e4m3t_e4m3n_f32n_tensor_op_gmma_f32_cooperative_stream_k,
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 1.0));
 }
 
 TEST(SM90_Device_Gemm_e4m3t_e4m3n_f32n_tensor_op_gmma_f32_cooperative_stream_k, 256x128x128_1x4x1) {
@@ -342,7 +347,8 @@ TEST(SM90_Device_Gemm_e4m3t_e4m3n_f32n_tensor_op_gmma_f32_cooperative_stream_k,
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 1.0));
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -393,7 +399,8 @@ TEST(SM90_Device_Gemm_e4m3t_e4m3n_f32n_tensor_op_gmma_f32_cooperative_stream_k,
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 1.0));
 }
 
 TEST(SM90_Device_Gemm_e4m3t_e4m3n_f32n_tensor_op_gmma_f32_cooperative_stream_k, 256x128x128_4x1x1) {
@@ -440,7 +447,8 @@ TEST(SM90_Device_Gemm_e4m3t_e4m3n_f32n_tensor_op_gmma_f32_cooperative_stream_k,
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 1.0));
 }
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -491,7 +499,8 @@ TEST(SM90_Device_Gemm_e4m3t_e4m3n_f32n_tensor_op_gmma_f32_cooperative_stream_k,
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 1.0));
 }
 
 TEST(SM90_Device_Gemm_e4m3t_e4m3n_f32n_tensor_op_gmma_f32_cooperative_stream_k, 256x128x128_2x4x1_fp8_fast_accum) {
@@ -538,7 +547,8 @@ TEST(SM90_Device_Gemm_e4m3t_e4m3n_f32n_tensor_op_gmma_f32_cooperative_stream_k,
   >;
 
   using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
-  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>());
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 0.0));
+  EXPECT_TRUE(test::gemm::device::TestAllBiasElementwise<Gemm>(1.0, 1.0));
 }
 
 #endif // defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)

test/unit/gemm/device/sm90_gemm_stream_k_scheduler.cu

@@ -101,8 +101,34 @@ test_scheduler(
   cutlass::KernelHardwareInfo hw_info{0, sm_count};
   auto params = Scheduler::to_underlying_arguments(problem_shape_mnkl, tile_shape, cluster_shape, hw_info, {splits}, nullptr);
 
+  typename Scheduler::Arguments args{};
+
+  // Set up the grid for the problem
+  dim3 grid = Scheduler::get_grid_shape(problem_shape_mnkl, tile_shape, cluster_shape, hw_info, args);
+
+  auto print_info = [&]() {
+    std::cout << "Failed with problem size "
+      << size<0>(problem_shape_mnkl) << "x"
+      << size<1>(problem_shape_mnkl) << "x"
+      << size<2>(problem_shape_mnkl) << "x"
+      << size<3>(problem_shape_mnkl)
+      << " and grid size " << grid.x << "x"
+      << grid.y << "x" << grid.z
+      << " splits=" << params.splits_
+      << " k_iter=" << params.divmod_tiles_per_output_tile_.divisor
+      << " big_units_=" << params.big_units_
+      << " big_groups_=" << params.big_groups_
+      << " sk_tiles=" << params.sk_tiles_
+      << " sk_units=" << params.sk_units_
+      << " k_tiles_per_sk_unit=" << params.k_tiles_per_sk_unit_
+      << " units_per_problem=" << params.units_per_problem_
+      << " groups=" << params.divmod_sk_groups_.divisor << std::endl;
+  };
+
   // If we expect the schedule to be data-parallel only, ensure that no stream-K tiles are launched.
   if (expect_data_parallel && params.sk_tiles_ != 0) {
+    print_info();
+    std::cout << "Expected stream-K to select a data-parallel decomposition." << std::endl;
     return false;
   }
 
@@ -114,15 +140,11 @@ test_scheduler(
   // Initialize counters to zero
   cudaError_t err = cudaMemset((void*)visit_counters.get(), 0, sizeof(int) * total_counters);
   if (err != cudaSuccess) {
-    std::cerr << __FILE__ << ":" << __LINE__ << " cudaMemset failed with error: " << cudaGetErrorString(err) << std::endl;
+    print_info();
+    std::cout << __FILE__ << ":" << __LINE__ << " cudaMemset failed with error: " << cudaGetErrorString(err) << std::endl;
     return false;
   }
 
-  typename Scheduler::Arguments args{};
-
-  // Set up the grid for the problem
-  dim3 grid = Scheduler::get_grid_shape(problem_shape_mnkl, tile_shape, cluster_shape, hw_info, args);
-
   // Set up cluster and cluster launch. This is needed even for this simple kernel because
   // the SM90 scheduler needs to be able to query the CTA id within a cluster, which requires
   // explicitly launching with clusters.
@@ -161,7 +183,8 @@ test_scheduler(
   err = cudaLaunchKernelExC(&launch_config, kernel, kernel_params);
 
   if (err != cudaSuccess) {
-    std::cerr << __FILE__ << ":" << __LINE__
+    print_info();
+    std::cout << __FILE__ << ":" << __LINE__
               << " cudaLaunchKernelExC failed with error: "
               << cudaGetErrorString(err) << std::endl;
     return false;
@@ -169,7 +192,8 @@ test_scheduler(
 
   err = cudaDeviceSynchronize();
   if (err != cudaSuccess) {
-    std::cerr << __FILE__ << ":" << __LINE__
+    print_info();
+    std::cout << __FILE__ << ":" << __LINE__
               << " scheduler kernel failed with error: "
               << cudaGetErrorString(err) << std::endl;
     return false;
@@ -181,20 +205,7 @@ test_scheduler(
 
   for (size_t i = 0; i < host_visit_counts.size(); ++i) {
     if (host_visit_counts[i] != 1) {
-      std::cout << "Failed with problem size "
-                << size<0>(problem_shape_mnkl) << "x"
-                << size<1>(problem_shape_mnkl) << "x"
-                << size<2>(problem_shape_mnkl) << "x"
-                << size<3>(problem_shape_mnkl)
-                << " and grid size " << grid.x << "x"
-                << grid.y << "x" << grid.z
-                << " splits=" << params.splits_
-                << " k_iter=" << params.divmod_tiles_per_output_tile_.divisor
-                << " big_units=" << params.big_units_
-                << " sk_tiles=" << params.sk_tiles_
-                << " sk_units=" << params.sk_units_
-                << " k_tiles_per_sk_unit=" << params.k_tiles_per_sk_unit_
-                << " units_per_problem=" << params.units_per_problem_ << std::endl;
+      print_info();
       std::cout << "Error at idx: " << i << ". Got count " << host_visit_counts[i] << std::endl;
       return false;
     }
@@ -301,7 +312,7 @@ TEST(SM90_Device_Gemm_stream_k_scheduler, 256x128x64_2x1x1) {
   // Test various data-parallel cases
   EXPECT_TRUE(test_data_parallel(/*blocks_m=*/ 4, /*blocks_n=*/ 4, tile_shape, cluster_shape, /*sm_count=*/ 16));
   EXPECT_TRUE(test_data_parallel(/*blocks_m=*/16, /*blocks_n=*/ 4, tile_shape, cluster_shape, /*sm_count=*/ 64));
-  EXPECT_TRUE(test_data_parallel(/*blocks_m=*/ 4, /*blocks_n=*/27, tile_shape, cluster_shape, /*sm_count=*/108));
+  EXPECT_TRUE(test_data_parallel(/*blocks_m=*/ 8, /*blocks_n=*/27, tile_shape, cluster_shape, /*sm_count=*/108));
 
   // Test various stream-K cases
   EXPECT_TRUE(test_stream_k(tile_shape, cluster_shape, /*sm_count=*/ 16));