v4.4 release update v2. (#2999)

include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized.hpp

@@ -832,6 +832,12 @@ public:
     mainloop_pipeline.init_masks(cluster_shape, block_id_in_cluster);
     accumulator_pipeline.init_masks(cluster_shape, block_id_in_cluster);
 
+    // Ensure that the prefetched kernel does not touch
+    // unflushed global memory prior to this instruction.
+    // For the static grouped scheduler, the problem shapes
+    // might be produced by a previous kernel in global memory.
+    cutlass::arch::wait_on_dependent_grids();
+
     // TileID scheduler
     TileScheduler scheduler(
       (!IsTensorMapUpdateAsync || is_participant.sched || is_participant.tensor_map_updater)
@@ -842,12 +848,6 @@ public:
     );
 
     auto work_tile_info = [&] () {
-      // Ensure that the prefetched kernel does not touch
-      // unflushed global memory prior to this instruction.
-      // For the static grouped scheduler, the problem shapes
-      // might be produced by a previous kernel in global memory.
-      cutlass::arch::wait_on_dependent_grids();
-
       if constexpr (IsTensorMapUpdateAsync) {
         return scheduler.initial_work_tile_info(cluster_shape, [] (typename TileScheduler::CLCResponse response) {
           CLCResponseWithAdditionalInformation response_with_additional_info = response;

include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_input_transform.hpp

@@ -652,12 +652,12 @@ public:
     // Allocate accumulators
     auto acc_shape = collective_mainloop.partition_accumulator_shape();
 
-    // TileID scheduler
-    TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, block_id_in_cluster);
-
     // Ensure memory ops in this kernel are not done prior to completion of dependent grids.
     cutlass::arch::wait_on_dependent_grids();
 
+    // TileID scheduler
+    TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, block_id_in_cluster);
+
     typename TileScheduler::WorkTileInfo work_tile_info = scheduler.initial_work_tile_info(cluster_shape);
     auto cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
 

include/cutlass/gemm/kernel/sm100_gemm_array_tma_warpspecialized_mma_transform.hpp

@@ -735,12 +735,12 @@ public:
     Tensor accumulators = cutlass::detail::make_sm100_accumulator<AccumulatorPipelineStageCount, IsOverlappingAccum>(
         tiled_mma, acc_shape, EpilogueTile{});
 
-    // TileID scheduler
-    TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, block_id_in_cluster);
-
     // Ensure memory ops in this kernel are not done prior to completion of dependent grids.
     cutlass::arch::wait_on_dependent_grids();
 
+    // TileID scheduler
+    TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, block_id_in_cluster);
+
     typename TileScheduler::WorkTileInfo work_tile_info = scheduler.initial_work_tile_info(cluster_shape);
     auto cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
 

include/cutlass/gemm/kernel/sm100_tile_scheduler_group.hpp

@@ -116,12 +116,15 @@ public:
 
   CUTLASS_DEVICE
   PersistentTileSchedulerSm100Group() { }
-
+  
+  // Note: constructing this tile scheduler can touch global memory that was
+  // written to by the prior kernel.
   CUTLASS_DEVICE
   PersistentTileSchedulerSm100Group(CLCResponse* clc_response_ptr, Params const& params)
     : scheduler_params(params),
       scheduler_sm90(params.params_sm90_, clc_response_ptr) { }
-
+  // Note: constructing this tile scheduler can touch global memory that was
+  // written to by the prior kernel.
   CUTLASS_DEVICE
   PersistentTileSchedulerSm100Group(CLCResponse* clc_response_ptr, Params const& params, dim3 /* block_id_in_cluster */)
     : scheduler_params(params),

include/cutlass/gemm/kernel/sm103_blockscaled_gemm_array_tma_warpspecialized.hpp

@@ -752,12 +752,12 @@ public:
     Tensor accumulators = cutlass::detail::make_sm100_accumulator<AccumulatorPipelineStageCount, IsOverlappingAccum>(
         tiled_mma, acc_shape, EpilogueTile{});
 
-    // TileID scheduler
-    TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, block_id_in_cluster);
-
     // Ensure memory ops in this kernel are not done prior to completion of dependent grids.
     cutlass::arch::wait_on_dependent_grids();
 
+    // TileID scheduler
+    TileScheduler scheduler(&shared_storage.clc_response[0], params.scheduler, block_id_in_cluster);
+
     typename TileScheduler::WorkTileInfo work_tile_info = scheduler.initial_work_tile_info(cluster_shape);
     auto cta_coord_mnkl = scheduler.work_tile_to_cta_coord(work_tile_info);
 

include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp

@@ -454,16 +454,6 @@ public:
     // Kernel level shared memory storage
     SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
 
-    auto scheduler = [&] () {
-      // Group scheduler requires a different constructor that takes a response ptr
-      if constexpr (cute::is_same_v<SchedulerTag, GroupScheduler>) {
-        return TileScheduler{params.scheduler, shared_storage.scheduler_response};
-      }
-      else {
-        return TileScheduler{params.scheduler};
-      }
-    } ();
-
     // In a warp specialized kernel, collectives expose data movement and compute operations separately
     CollectiveMainloop collective_mainloop;
     CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
@@ -585,6 +575,16 @@ public:
     // Ensure memory ops in this kernel are not done prior to completion of dependent grids.
     cutlass::arch::wait_on_dependent_grids();
 
+    auto scheduler = [&] () {
+      // Group scheduler requires a different constructor that takes a response ptr
+      if constexpr (cute::is_same_v<SchedulerTag, GroupScheduler>) {
+        return TileScheduler{params.scheduler, shared_storage.scheduler_response};
+      }
+      else {
+        return TileScheduler{params.scheduler};
+      }
+    } ();
+
     auto work_tile_info = scheduler.initial_work_tile_info(ClusterShape{});
 
     if (not work_tile_info.is_valid()) {

include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp

@@ -463,16 +463,6 @@ public:
     // Kernel level shared memory storage
     SharedStorage& shared_storage = *reinterpret_cast<SharedStorage*>(smem_buf);
 
-    auto scheduler = [&] () {
-      // Group scheduler requires a different constructor that takes a response ptr
-      if constexpr (cute::is_same_v<SchedulerTag, GroupScheduler>) {
-        return TileScheduler{params.scheduler, shared_storage.scheduler_response};
-      }
-      else {
-        return TileScheduler{params.scheduler};
-      }
-    } ();
-
     // In a warp specialized kernel, collectives expose data movement and compute operations separately
     CollectiveMainloop collective_mainloop;
     CollectiveEpilogue collective_epilogue(params.epilogue, shared_storage.tensors.epilogue);
@@ -600,6 +590,16 @@ public:
     // Ensure memory ops in this kernel are not done prior to completion of dependent grids.
     cutlass::arch::wait_on_dependent_grids();
 
+    auto scheduler = [&] () {
+      // Group scheduler requires a different constructor that takes a response ptr
+      if constexpr (cute::is_same_v<SchedulerTag, GroupScheduler>) {
+        return TileScheduler{params.scheduler, shared_storage.scheduler_response};
+      }
+      else {
+        return TileScheduler{params.scheduler};
+      }
+    } ();
+
     auto work_tile_info = scheduler.initial_work_tile_info(ClusterShape{});
 
     if (not work_tile_info.is_valid()) {

include/cutlass/gemm/kernel/sm90_tile_scheduler_group.hpp

@@ -246,6 +246,8 @@ public:
 
   PersistentTileSchedulerSm90Group() = default;
 
+  // Note: constructing this tile scheduler can touch global memory that was
+  // written to by the prior kernel.
   CUTLASS_DEVICE explicit PersistentTileSchedulerSm90Group(Params const& params_, SchedulerResponse* response_ptr) : scheduler_params(params_), response_ptr_(response_ptr) {
     // MSVC requires protecting use of CUDA-specific nonstandard syntax,
     // like blockIdx and gridDim, with __CUDA_ARCH__.