From 22f8870fbdb59ccfbb03089ad277531b967c36d5 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 02:34:33 +0000 Subject: [PATCH] fix: correction_epilog with paired atoms + pre-partitioned TMA store --- tests/unit/test_fmha_v3_stage_c.py | 116 ++++++++++++++--------------- 1 file changed, 58 insertions(+), 58 deletions(-) diff --git a/tests/unit/test_fmha_v3_stage_c.py b/tests/unit/test_fmha_v3_stage_c.py index 534adc3e..80969b1a 100644 --- a/tests/unit/test_fmha_v3_stage_c.py +++ b/tests/unit/test_fmha_v3_stage_c.py @@ -170,6 +170,17 @@ class FmhaV3StageCMulti: gK = cute.local_tile(mK,cute.slice_(self.qk_mma_tiler,(0,None,None)),(None,None,None)) gV = cute.local_tile(mV,cute.slice_(self.pv_mma_tiler,(0,None,None)),(None,None,None)) gC = cute.local_tile(mC,cute.slice_(self.pv_mma_tiler,(None,None,0)),(None,None,None)) + tCgC_epi = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None, None, 0)), (None, None, None)) + + # Pre-partition TMA store tensors (needed for correction_epilog) + epi_s = cute.select(self.c_smem_s, mode=[0, 1]) + tCgC_epi_flat = cute.flat_divide(tCgC_epi, epi_tile) + bSG_sC, bSG_gC = cpasync.tma_partition( + tma_c, 0, cute.make_layout(1), + cute.group_modes(sC, 0, 2), + cute.group_modes(tCgC_epi_flat, 0, 2), + ) + n_kv_tiles = cute.size(gK, mode=[3]) qk_thr = qk_mma.get_slice(0); pv_thr = pv_mma.get_slice(0) @@ -407,69 +418,58 @@ class FmhaV3StageCMulti: # Wait for MMA's PV[N-1] to commit before reading O. final_o_bar.arrive_and_wait() - # DIAG: NO-OP TMEM round-trip before normalize - tTMrO_noop = cute.make_rmem_tensor( - (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype - ) - for i in range(n_corr_tiles): - tTMrO_i_ = tTMrO_noop[None, i] - tTMrO_i_layout = cute.composition( - tTMrO_i_.layout, cute.make_layout(tTMrO_noop.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, - ) - cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO_i) - cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STOREtO_i) - cute.arch.fence_view_async_tmem_store() - - # === Final O normalization: O *= 1/row_sum === + # === Correction epilog: one-way TMEM → reg → SMEM with normalize === + # Uses get_tmem_load_op + get_smem_store_op paired atoms. inv_row_sum = Float32(1.0) / row_sum - tTMrO = cute.make_rmem_tensor( - (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype + epi_corr_tile_size = 32 * 8 // self.o_dtype.width # 16 for BF16 + + tOtO_epi = cute.logical_divide(tOtO0, 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_tile[0], epi_corr_tile_size), self.use_2cta_instrs, ) - - 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 - ) - - 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) - - cute.arch.fence_view_async_tmem_store() - - # 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 + tiled_tmem_load_epi = tcgen05.make_tmem_copy(tmem_load_epi_atom, tOtO_epi[(None, None), 0]) + smem_store_epi_atom = utils.sm100.get_smem_store_op( + self.c_layout, self.o_dtype, self.acc_dtype, tiled_tmem_load_epi, ) - 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, + tiled_smem_store_epi = cute.make_tiled_copy_D(smem_store_epi_atom, tiled_tmem_load_epi) + + 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) + 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))) + + thr_tmem_load_epi = tiled_tmem_load_epi.get_slice(sfw_idx) + 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]) + + n_epi_corr_tiles = self.pv_mma_tiler[1] // epi_corr_tile_size + for i in range(n_epi_corr_tiles): + 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[None, 0, 0, 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[None, 0, 0, i]) + + cute.arch.fence_proxy("async.shared", space="cta") + + # TMA store SMEM → GMEM using pre-partitioned tensors + epi_bar = pipeline.NamedBarrier( + barrier_id=self.epilog_sync_bar_id, + num_threads=32 * len(self.epilogue_warp_id), ) - c_pipe.producer_tail() + epi_bar.arrive_and_wait() + cute.copy(tma_c, bSG_sC[(None, 0)], bSG_gC[(None, 0, 0, 0, 0, 0, 0)]) + 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)