From b457d196afd1cf99aa4e8ac8af7deb03204a2645 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 02:26:57 +0000 Subject: [PATCH] fix: use paired atoms for correction_epilog + cute.copy TMA store --- tests/unit/test_fmha_v3_stage_c.py | 122 +++++++++++++++-------------- 1 file changed, 62 insertions(+), 60 deletions(-) diff --git a/tests/unit/test_fmha_v3_stage_c.py b/tests/unit/test_fmha_v3_stage_c.py index 681b9cdd..b13098c7 100644 --- a/tests/unit/test_fmha_v3_stage_c.py +++ b/tests/unit/test_fmha_v3_stage_c.py @@ -348,75 +348,77 @@ class FmhaV3StageCMulti: final_o_bar.arrive_and_wait() # === Correction epilog: one-way TMEM → reg → SMEM with normalize === - # Replicates epilogue_tma_store logic with inv_row_sum multiply added. # Uses get_tmem_load_op + get_smem_store_op paired atoms. + # NO TMEM round-trip — hand-constructed atoms corrupt data. inv_row_sum = Float32(1.0) / row_sum - # Build TMEM→reg and reg→SMEM copies using paired atoms - tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout) - tCtO = utils.gemm.sm100.transform_partitioned_tensor_layout(tCtO_base) - tCgC_t = utils.gemm.sm100.transform_partitioned_tensor_layout(tCgC) - tiled_copy_t2r, tTR_tO, tTR_rO = utils.gemm.sm100.epilogue_tmem_copy_and_partition( - self, tidx, tCtO, tCgC_t, epi_tile, self.use_2cta_instrs + 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, ) - tTR_rC = cute.make_rmem_tensor(tTR_rO.shape, self.c_dtype) - tiled_copy_r2s, tRS_rC, tRS_sC = utils.gemm.sm100.epilogue_smem_copy_and_partition( - self, tiled_copy_t2r, tTR_rC, tidx, sC + 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, ) - tCgC_epi = cute.flat_divide(tCgC_t, epi_tile) - bSG_sC, bSG_gC_part = cpasync.tma_partition( - tma_c, 0, cute.make_layout(1), - cute.group_modes(sC, 0, 2), - cute.group_modes(tCgC_epi, 0, 2), - ) - epilog_sync_bar = pipeline.NamedBarrier( + 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 epilogue_tma_store helper + # Feed it the unscaled TMEM data — it reads from TMEM, but our scaled + # data is in SMEM. So we need to bypass the TMEM read. + # Instead, use a simple cp_async_bulk from SMEM to GMEM. + epi_bar = pipeline.NamedBarrier( barrier_id=self.epilog_sync_bar_id, num_threads=32 * len(self.epilogue_warp_id), ) + epi_bar.arrive_and_wait() - acc_cons_st = pipeline.make_pipeline_state( - pipeline.PipelineUserType.Consumer, self.num_acc_stage - ) - c_pipe = pipeline.PipelineTmaStore.create( - num_stages=self.num_c_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), - ) - acc_pipe.consumer_wait(acc_cons_st) - - tTR_tO_tile = tTR_tO[(None, None, None, None, None, acc_cons_st.index)] - bSG_gC = bSG_gC_part[(None, None, None, Int32(0), Int32(0), Int32(0))] - tTR_tO_tile = cute.group_modes(tTR_tO_tile, 3, cute.rank(tTR_tO_tile)) - bSG_gC = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) - - subtile_cnt = cute.size(tTR_tO_tile.shape, mode=[3]) - for subtile_idx in range(subtile_cnt): - tTR_tO_mn = tTR_tO_tile[(None, None, None, subtile_idx)] - cute.copy(tiled_copy_t2r, tTR_tO_mn, tTR_rO) - - # *** Normalize: multiply by inv_row_sum *** - for j in cutlass.range(cute.size(tTR_rO), vectorize=True): - tTR_rO[j] = tTR_rO[j] * inv_row_sum - - # Convert FP32 → BF16 - acc_vec = tiled_copy_r2s.retile(tTR_rO).load() - acc_vec = acc_vec.to(self.c_dtype) - tRS_rC.store(acc_vec) - - c_buffer = subtile_idx % self.num_c_stage - cute.copy(tiled_copy_r2s, tRS_rC, tRS_sC[(None, None, None, c_buffer)]) - cute.arch.fence_proxy("async.shared", space="cta") - epilog_sync_bar.arrive_and_wait() - - if warp_idx == self.epilogue_warp_id[0]: - cute.copy(tma_c, bSG_sC[(None, c_buffer)], bSG_gC[(None, subtile_idx)]) - c_pipe.producer_commit() - c_pipe.producer_acquire() - epilog_sync_bar.arrive_and_wait() - - epilog_sync_bar.arrive_and_wait() - acc_pipe.consumer_release(acc_cons_st) - acc_cons_st.advance() - c_pipe.producer_tail() + # Use cute.copy with tma_c atom for TMA store (only from warp 0) + # Need to get the SMEM and GMEM tensors into the right shape + # cute.copy(tma_c, sC_slice, gC) — but need proper partitioning + # Fallback: write O*inv_row_sum back to TMEM (2nd round-trip risk) + # then use epilogue_tma_store. + # OR: just use epilogue_tma_store with a no-op epilogue and + # apply normalize on CPU. But Mike says no shortcuts. + # For now, write back to TMEM (which we know corrupts slightly) + # and see if the corruption is acceptable with the paired atoms. + # Actually — we DON'T have a paired TMEM store. We have paired + # TMEM load + SMEM store. The data is ALREADY in SMEM (scaled). + # We need TMA store from SMEM → GMEM. + # The cleanest: use the TMA atom directly. + # tma_c was created with make_tiled_tma_atom(S2G, c, epi_s, epi_tile) + # cute.copy(tma_c, sC, gC) should work if shapes match. + cute.copy(tma_c, sC, gC) + 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)