diff --git a/deep_gemm/include/deep_gemm/fp8_gemm.cuh b/deep_gemm/include/deep_gemm/fp8_gemm.cuh index 0d9d3f7..5fb5b66 100644 --- a/deep_gemm/include/deep_gemm/fp8_gemm.cuh +++ b/deep_gemm/include/deep_gemm/fp8_gemm.cuh @@ -292,7 +292,7 @@ fp8_gemm_kernel(__nv_bfloat16* gmem_d, float* scales_b, int* grouped_layout, if constexpr (kNumTMAMulticast == 1) { lane_idx == 0 ? empty_barriers[s]->arrive() : void(); } else { - auto target_cta_idx = scheduler.is_tma_multicast_valid(m_block_idx) ? lane_idx : cute::block_rank_in_cluster(); + auto target_cta_idx = scheduler.num_blocks_in_group != 1 ? lane_idx : cute::block_rank_in_cluster(); lane_idx < kNumTMAMulticast ? empty_barriers[s]->arrive(target_cta_idx) : void(); } }; diff --git a/deep_gemm/include/deep_gemm/scheduler.cuh b/deep_gemm/include/deep_gemm/scheduler.cuh index ebe65a6..6a595b8 100644 --- a/deep_gemm/include/deep_gemm/scheduler.cuh +++ b/deep_gemm/include/deep_gemm/scheduler.cuh @@ -45,7 +45,6 @@ struct Scheduler { } __device__ __forceinline__ bool is_tma_multicast_valid(const uint32_t& m_block_idx) { - return true; if (num_blocks_in_group == 1) return false; if constexpr (kGemmType == GemmType::Normal or kGemmType == GemmType::GroupedMasked) {