diff --git a/python/cudnn/gemm/cutedsl/grouped/quant/moe_blockscaled_grouped_gemm_quant_rubin.py b/python/cudnn/gemm/cutedsl/grouped/quant/moe_blockscaled_grouped_gemm_quant_rubin.py index 10862358f..75c1cdd87 100644 --- a/python/cudnn/gemm/cutedsl/grouped/quant/moe_blockscaled_grouped_gemm_quant_rubin.py +++ b/python/cudnn/gemm/cutedsl/grouped/quant/moe_blockscaled_grouped_gemm_quant_rubin.py @@ -2254,6 +2254,9 @@ def kernel( if reverse_subtile: real_subtile_idx = self.cta_tile_shape_mnk[1] // self.epi_tile_n_required - 1 - subtile_idx + tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)] + cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) + # C1 fix: fence + early release for overlapping_accum if cutlass.const_expr(self.overlapping_accum): if subtile_idx == self.iter_acc_early_release_in_epilogue: @@ -2262,9 +2265,6 @@ def kernel( acc_pipeline.consumer_release(acc_consumer_state) acc_consumer_state.advance() - tTR_tAcc_mn = tTR_tAcc[(None, None, None, real_subtile_idx)] - cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) - # For breuse, update mProb based on which M half this subtile belongs to. # With transform, subtiles interleave M groups: even = bkeep, odd = breuse. if cutlass.const_expr(self.enable_breuse and self.has_prob):