From 3c134f7e902324dbde00388b4cdac3f3fbbf5e6e Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 01:18:56 +0000 Subject: [PATCH] fix: replace TMEM round-trip normalize with CUTLASS correction_epilog pattern --- tests/unit/test_fmha_v3_stage_c.py | 106 ++++++++++++++++++++--------- 1 file changed, 72 insertions(+), 34 deletions(-) diff --git a/tests/unit/test_fmha_v3_stage_c.py b/tests/unit/test_fmha_v3_stage_c.py index 8dde4630..45767c22 100644 --- a/tests/unit/test_fmha_v3_stage_c.py +++ b/tests/unit/test_fmha_v3_stage_c.py @@ -399,48 +399,86 @@ class FmhaV3StageCMulti: # Wait for MMA's PV[N-1] to commit before reading O. final_o_bar.arrive_and_wait() - # === Final O normalization: O *= 1/row_sum === - # DIAG: use 1.0 instead of 1/row_sum to test raw PV output - inv_row_sum = Float32(1.0) + # === Final O normalization + epilogue: CUTLASS correction_epilog pattern === + # ONE-WAY trip: TMEM → reg (normalize + FP32→BF16) → SMEM → TMA → GMEM + # NO TMEM round-trip. Hand-constructed atoms corrupt data on round-trip. + inv_row_sum = Float32(1.0) / row_sum - tTMrO = cute.make_rmem_tensor( - (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype + # Build paired atoms for TMEM load → SMEM store + epi_corr_tile_size = 32 * 8 // self.o_dtype.width # 16 for BF16 + epi_subtile = (self.epi_tile[0], epi_corr_tile_size) + + tOsO = pv_thr.partition_C(sC) + cO_epi = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) + tOcO_epi = pv_thr.partition_C(cO_epi) + + tOtO_epi = cute.logical_divide(tOtO0, cute.make_layout((128, epi_corr_tile_size))) + tOsO_epi = cute.logical_divide(tOsO, cute.make_layout((128, epi_corr_tile_size))) + tOcO_epi = cute.logical_divide(tOcO_epi, cute.make_layout((128, epi_corr_tile_size))) + + tmem_load_epi_atom = utils.sm100.get_tmem_load_op( + self.pv_mma_tiler, + self.c_layout, + self.o_dtype, + self.acc_dtype, + epi_subtile, + use_2cta_instrs=False, + ) + tiled_tmem_load_epi = tcgen05.make_tmem_copy( + tmem_load_epi_atom, tOtO_epi[(None, None), 0] + ) + thr_tmem_load_epi = tiled_tmem_load_epi.get_slice(sfw_idx) + + smem_store_epi_atom = utils.sm100.get_smem_store_op( + self.c_layout, + self.o_dtype, + self.acc_dtype, + tiled_tmem_load_epi, + ) + tiled_smem_store_epi = cute.make_tiled_copy_D( + smem_store_epi_atom, tiled_tmem_load_epi ) - for i in range(n_corr_tiles): - tTMrO_i_ = tTMrO[None, i] - tTMrO_i_layout = cute.composition( - tTMrO_i_.layout, cute.make_layout(tTMrO.shape[0]) - ) - tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout) - tTMEM_LOADtO_i = cute.make_tensor( - tTMEM_LOADtO.iterator + i * corr_tile_size, tTMEM_LOADtO.layout - ) - tTMEM_STOREtO_i = cute.make_tensor( - tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout - ) + tTMEM_LOADtO_epi = thr_tmem_load_epi.partition_S(tOtO_epi[(None, None), None]) + tTMEM_LOADsO_epi = thr_tmem_load_epi.partition_D(tOsO_epi[(None, None), None]) + tTMEM_LOADcO_epi = thr_tmem_load_epi.partition_D(tOcO_epi[(None, None), None]) - cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO_i) - for j in cutlass.range(cute.size(tTMrO_i), vectorize=True): - tTMrO_i[j] = tTMrO_i[j] * inv_row_sum - cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STOREtO_i) + n_epi_corr_tiles = self.pv_mma_tiler[1] // epi_corr_tile_size + for i in range(n_epi_corr_tiles): + tTMEM_LOADtO_epi_i = tTMEM_LOADtO_epi[None, 0, 0, i] + tTMEM_LOADsO_epi_i = tTMEM_LOADsO_epi[None, 0, 0, i] + tTMrO = cute.make_rmem_tensor( + tTMEM_LOADcO_epi[None, 0, 0, i].shape, self.acc_dtype + ) + cute.copy(tiled_tmem_load_epi, tTMEM_LOADtO_epi_i, tTMrO) + for j in range(cute.size(tTMrO)): + tTMrO[j] = tTMrO[j] * inv_row_sum + tSMrO = cute.make_rmem_tensor(tTMrO.shape, self.o_dtype) + o_vec = tTMrO.load() + tSMrO.store(o_vec.to(self.o_dtype)) + cute.copy(tiled_smem_store_epi, tSMrO, tTMEM_LOADsO_epi_i) - cute.arch.fence_view_async_tmem_store() + cute.arch.fence_proxy("async.shared", space="cta") - # 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 + # TMA store: SMEM → GMEM + epi_s_tile = cute.select(self.c_smem_s, mode=[0, 1]) + tma_c_epi, mC_epi = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), c, epi_s_tile, self.epi_tile ) - 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, + tCgC_epi = cute.local_tile(mC_epi, cute.slice_(self.pv_mma_tiler, (None, None, 0)), (None, None, None)) + tCsC_epi = cute.local_tile(sC, cute.slice_(self.epi_tile, (None, None)), (None, None)) + + # Sync before TMA store — all softmax warps must finish SMEM writes + softmax_all_bar = pipeline.NamedBarrier( + barrier_id=5, num_threads=32 * len(self.epilogue_warp_id) ) - c_pipe.producer_tail() + softmax_all_bar.arrive_and_wait() + + # Warp 0 does the TMA store + if sfw_idx < 32: + cute.copy(tma_c_epi, tCsC_epi, tCgC_epi) + cute.arch.cp_async_bulk_commit_group() + cute.arch.cp_async_bulk_wait_group(0, read=True) tmem.relinquish_alloc_permit() tmem.free(tmem_ptr)