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 c69defbf..6e3ce1e0 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 @@ -614,6 +614,9 @@ public: if (producer_warp_role == ProducerWarpRole::Scheduler) { // GroupScheduler requires a producer warp to iterate over the group infos and push // the work tile infos to the downstream pipelines. + #ifdef CUTLASS_ENABLE_GDC_FOR_SM90 + cutlass::arch::wait_on_dependent_grids(); + #endif if constexpr (cute::is_same_v) { do { auto [next_work_tile_info, increment_pipe] = scheduler.advance_to_next_work(tile_scheduler_pipeline, tile_scheduler_pipe_producer_state); @@ -627,6 +630,9 @@ public: } // Mainloop Producer Warp else if (producer_warp_role == ProducerWarpRole::Mainloop) { + #ifdef CUTLASS_ENABLE_GDC_FOR_SM90 + cutlass::arch::wait_on_dependent_grids(); + #endif int32_t curr_batch = idx2crd(work_tile_info.L_idx, shape<4>(gB_nkl)); // Usually just returns work_tile_info.L_idx; int32_t const mock_l_coord = 0; int32_t const sm_idx = blockIdx.x + (blockIdx.y * gridDim.x); @@ -791,6 +797,9 @@ public: } // Mainloop Auxiliary Load Producer Warp End // Epilogue Producer Warp else if (producer_warp_role == ProducerWarpRole::Epilogue && collective_epilogue.is_producer_load_needed()) { + #ifdef CUTLASS_ENABLE_GDC_FOR_SM90 + cutlass::arch::wait_on_dependent_grids(); + #endif int32_t const sm_idx = blockIdx.x + (blockIdx.y * gridDim.x); int32_t const sm_count = params.hw_info.sm_count; @@ -994,6 +1003,11 @@ public: // Get next work tile auto [next_work_tile_info, increment_pipe] = scheduler.fetch_next_work(work_tile_info, tile_scheduler_pipeline, tile_scheduler_pipe_consumer_state); + #ifdef CUTLASS_ENABLE_GDC_FOR_SM90 + if (!next_work_tile_info.is_valid()) { + cutlass::arch::launch_dependent_grids(); + } + #endif work_tile_info = next_work_tile_info; if (increment_pipe) { ++tile_scheduler_pipe_consumer_state;