diff --git a/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp b/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp index 9e326d2b..987dd622 100644 --- a/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp +++ b/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp @@ -608,9 +608,10 @@ public: // Get pipeline stage increments from tensor shapes auto k_tile_count = size<3>(gA_mkl); - #ifdef CUTLASS_ENABLE_GDC_FOR_SM90 + // Ensure that the kernel does not touch + // unflushed global memory prior to this instruction cutlass::arch::wait_on_dependent_grids(); - #endif + if (warp_group_role == WarpGroupRole::Producer) { cutlass::arch::warpgroup_reg_dealloc(); diff --git a/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp b/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp index be70bdff..bd979c7b 100644 --- a/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp +++ b/include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_pingpong.hpp @@ -646,9 +646,10 @@ public: // Get pipeline stage increments from tensor shapes auto k_tile_count = size<3>(gA_mkl); - #ifdef CUTLASS_ENABLE_GDC_FOR_SM90 + // Ensure that the kernel does not touch + // unflushed global memory prior to this instruction cutlass::arch::wait_on_dependent_grids(); - #endif + if (warp_group_role == WarpGroupRole::Producer) { cutlass::arch::warpgroup_reg_dealloc();