From ea66b6ee8d32ad45e7a8ce4e5004b279340bed27 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 02:15:28 +0000 Subject: [PATCH] =?UTF-8?q?diag:=20NO-OP=20TMEM=20round-trip=20test=20?= =?UTF-8?q?=E2=80=94=20load+store=20back=20unchanged?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/unit/test_fmha_v3_stage_c.py | 121 +++++++++++++---------------- 1 file changed, 56 insertions(+), 65 deletions(-) diff --git a/tests/unit/test_fmha_v3_stage_c.py b/tests/unit/test_fmha_v3_stage_c.py index 9ea17b8e..64346644 100644 --- a/tests/unit/test_fmha_v3_stage_c.py +++ b/tests/unit/test_fmha_v3_stage_c.py @@ -399,79 +399,70 @@ class FmhaV3StageCMulti: # Wait for MMA's PV[N-1] to commit before reading O. final_o_bar.arrive_and_wait() - # === Correction epilog: one-way TMEM → reg → SMEM → GMEM with normalize === - # Uses get_tmem_load_op + get_smem_store_op paired atoms (same as CUTLASS correction_epilog). - # NO TMEM round-trip — hand-constructed Ld32x32bOp/St32x32bOp atoms corrupt data. + # DIAG: Test TMEM round-trip with NO-OP (load + store back unchanged) + # If cos drops from 0.999998, the round-trip atoms are the problem. + tTMrO = cute.make_rmem_tensor( + (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype + ) + 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) + # NO-OP: store back without modification + 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 === inv_row_sum = Float32(1.0) / row_sum - # Build the TMEM→reg and reg→SMEM tiled 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) - tiled_copy_t2r, tTR_tO, tTR_rO = utils.gemm.sm100.epilogue_tmem_copy_and_partition( - self, tidx, tCtO, tCgC, epi_tile, 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 - ) - tCgC_epi = cute.flat_divide(tCgC, epi_tile) - bSG_sC, bSG_gC_partitioned = 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( - barrier_id=self.epilog_sync_bar_id, - num_threads=32 * len(self.epilogue_warp_id), + tTMrO = cute.make_rmem_tensor( + (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype ) - # Consume the accumulator pipeline + 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 ) - c_pipe = pipeline.PipelineTmaStore.create( - num_stages=self.num_c_stage, - producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id)), + 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, ) - acc_pipe.consumer_wait(acc_cons_st) - - # Slice to the current tile - tTR_tO_tile = tTR_tO[(None, None, None, None, None, acc_cons_st.index)] - bSG_gC = bSG_gC_partitioned[(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)) - - # Store O to global memory in subtiles, applying 1/row_sum normalize - 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) - - # Apply normalize: multiply by inv_row_sum, then convert to BF16 - acc_vec = tiled_copy_r2s.retile(tTR_rO).load() - # acc_vec is in FP32 — apply scale before conversion - # We can't directly scale the vector, but we can scale the register tensor - for j in cutlass.range(cute.size(tTR_rO), vectorize=True): - tTR_rO[j] = tTR_rO[j] * inv_row_sum - 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_cnt * 0 + subtile_idx # num_prev_subtiles = 0 - c_buffer = c_buffer % 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() tmem.relinquish_alloc_permit()