diff --git a/tests/fmha_v3_stage_c_example7.py b/tests/fmha_v3_stage_c_example7.py index b72e8318..afe27334 100644 --- a/tests/fmha_v3_stage_c_example7.py +++ b/tests/fmha_v3_stage_c_example7.py @@ -369,44 +369,36 @@ class FmhaV3StageCMulti: # Wait for MMA's PV[N-1] to commit before reading O. final_o_bar.arrive_and_wait() - # === Reference-style scaled epilogue (mirrors CUTLASS FMHA correction_epilog) === - # Pattern: TMEM → reg (paired load atom) → scale in reg → FP32→BF16 in reg - # → SMEM (paired store atom) → TMA SMEM→GMEM. No TMEM round-trip. + # === O normalization via TMEM load → scale → TMEM store === + # Matches CUTLASS reference's correction_rescale pattern. + # Uses Ld32x32bOp / St32x32bOp with the SAME Repetition so the + # register tile shapes match (paired atoms). - corr_tile_size = 16 # matches the reference + corr_tile_size = 16 - # Sub-tile the O C-fragment for column-wise iteration + # Sub-tile the O C-fragment tOtO_i = cute.logical_divide(tOtO0, cute.make_layout((128, corr_tile_size))) cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) tOcO = pv_thr.partition_C(cO) tOcO_i = cute.logical_divide(tOcO, cute.make_layout((128, corr_tile_size))) - tOsO = pv_thr.partition_C(sC[None, None, 0]) - tOsO_i = cute.logical_divide(tOsO, cute.make_layout((128, corr_tile_size))) - # Paired atoms via sm100_utils (same as CUTLASS reference) - epi_subtile = (self.epi_tile[0], corr_tile_size) - tmem_load_op = utils.sm100.get_tmem_load_op( - self.pv_mma_tiler, - self.c_layout, - self.o_dtype, + # TMEM load + store atoms (paired — same Repetition) + tmem_load_o_atom = cute.make_copy_atom( + tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype, - epi_subtile, - use_2cta_instrs=False, ) - tiled_tmem_load_o = tcgen05.make_tmem_copy( - tmem_load_op, tOtO_i[(None, None), 0] - ) - thr_tmem_load_o = tiled_tmem_load_o.get_slice(sfw_idx) - smem_store_op = utils.sm100.get_smem_store_op( - self.c_layout, self.o_dtype, self.acc_dtype, tiled_tmem_load_o - ) - tiled_smem_store_o = cute.make_tiled_copy_D( - smem_store_op, tiled_tmem_load_o + tmem_store_o_atom = cute.make_copy_atom( + tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), + self.acc_dtype, ) + tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tOtO_i[(None, None), 0]) + tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_i[(None, None), 0]) + thr_load_o = tiled_tmem_load_o.get_slice(sfw_idx) + thr_store_o = tiled_tmem_store_o.get_slice(sfw_idx) - tTMEM_LOADtO = thr_tmem_load_o.partition_S(tOtO_i[(None, None), None]) - tTMEM_LOADsO = thr_tmem_load_o.partition_D(tOsO_i[(None, None), None]) - tTMEM_LOADcO = thr_tmem_load_o.partition_D(tOcO_i[(None, None), None]) + tTMEM_LOADtO = thr_load_o.partition_S(tOtO_i[(None, None), None]) + tTMEM_LOADcO = thr_load_o.partition_D(tOcO_i[(None, None), None]) + tTMEM_STOREtO = thr_store_o.partition_D(tOtO_i[(None, None), None]) # Scale = 1/row_sum inv_row_sum = Float32(1.0) / row_sum @@ -414,7 +406,9 @@ class FmhaV3StageCMulti: n_corr = self.pv_mma_tiler[1] // corr_tile_size for i in range(n_corr): tTMEM_LOADtO_i = tTMEM_LOADtO[None, 0, 0, i] - tTMEM_LOADsO_i = tTMEM_LOADsO[None, 0, 0, i] + tTMEM_STOREtO_i = cute.make_tensor( + tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout + ) tTMrO = cute.make_rmem_tensor( tTMEM_LOADcO[None, 0, 0, i].shape, self.acc_dtype ) @@ -424,47 +418,25 @@ class FmhaV3StageCMulti: for j in range(cute.size(tTMrO), vectorize=True): tTMrO[j] = tTMrO[j] * inv_row_sum - # FP32 → BF16 in registers - tSMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype) - o_vec = tTMrO.load() - tSMrO.store(o_vec.to(self.o_dtype)) + # Write back to TMEM + cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STOREtO_i) - # Registers → SMEM via paired atom - cute.copy(tiled_smem_store_o, tSMrO, tTMEM_LOADsO_i) + cute.arch.fence_view_async_tmem_store() - cute.arch.fence_view_async_tmem_load() - # Async-proxy fence so the TMA store sees the SMEM writes. - cute.arch.fence_proxy("async.shared", space="cta") - # Use NamedBarrier to sync softmax warps with TMA store warp - epi_sync_bar = pipeline.NamedBarrier( - barrier_id=self.epilog_sync_bar_id, - num_threads=32 * len(self.epilogue_warp_id), - ) - epi_sync_bar.arrive_and_wait() - - # TMA SMEM -> GMEM. One warp issues the copy; the rest waited at - # the named barrier above. (Match epilogue_tma_store's behavior - # with all-thread arrive.) - if warp_idx == self.epilogue_warp_id[0]: - # Partition sC and gC for TMA. - tOsO_tma, tOgO_tma = cpasync.tma_partition( - tma_c, - 0, - cute.make_layout(1), - cute.group_modes(sC, 0, 2), - cute.group_modes(tCgC, 0, 2), - ) - # tOgO_tma still has the trailing tile/batch coords; we want - # the first (and only) tile here. - cute.copy(tma_c, tOsO_tma[None, 0], tOgO_tma[None, 0, 0, 0]) - cute.arch.cp_async_bulk_commit_group() - cute.arch.cp_async_bulk_wait_group(0, read=True) - - # Release the acc pipe so MMA's producer_tail can complete. + # Standard epilogue: TMEM → SMEM → GMEM via TMA store. + # O in TMEM is now scaled by 1/row_sum. + tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) acc_cons_st = pipeline.make_pipeline_state( pipeline.PipelineUserType.Consumer, self.num_acc_stage ) - acc_pipe.consumer_release(acc_cons_st) + c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)) + c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp) + acc_cons_st = utils.gemm.sm100.epilogue_tma_store( + self, tidx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile, + 0, const_expr(lambda x: x), (0, 0, 0), + acc_cons_st, acc_pipe, c_pipe, + ) + c_pipe.producer_tail() tmem.relinquish_alloc_permit() tmem.free(tmem_ptr)