From de7cc61480c27af4523ed4dced9867fbba56c3c8 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 02:53:59 +0000 Subject: [PATCH] fix: use logical_divide (not composition) for O rescale/normalize atoms to match get_tmem_load_op layout --- tests/unit/test_fmha_v3_stage_c.py | 28 +++++++++------------------- 1 file changed, 9 insertions(+), 19 deletions(-) diff --git a/tests/unit/test_fmha_v3_stage_c.py b/tests/unit/test_fmha_v3_stage_c.py index ea84ee35..2ba1c95a 100644 --- a/tests/unit/test_fmha_v3_stage_c.py +++ b/tests/unit/test_fmha_v3_stage_c.py @@ -235,12 +235,11 @@ class FmhaV3StageCMulti: scale_log2 = Float32(self.scale_softmax_log2) # O rescale atoms (hand-constructed, for per-tile O *= acc_scale) + # Use logical_divide (not composition) to match get_tmem_load_op layout corr_tile_size = 16 tOcO = pv_thr.partition_C(cS) - tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size))) - tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size))) - tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout) - tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout) + tOtO_epi_r = cute.logical_divide(tOtO0, cute.make_layout((128, corr_tile_size))) + tOcO_epi_r = cute.logical_divide(tOcO, cute.make_layout((128, corr_tile_size))) tmem_load_o_atom = cute.make_copy_atom( tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype, @@ -249,13 +248,13 @@ class FmhaV3StageCMulti: 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) - tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_i) + tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tOtO_epi_r[(None, None), 0]) + tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_epi_r[(None, None), 0]) thr_tmem_load_o = tiled_tmem_load_o.get_slice(sfw_idx) thr_tmem_store_o = tiled_tmem_store_o.get_slice(sfw_idx) - tTMEM_LOADtO = thr_tmem_load_o.partition_S(tOtO_i) - tTMEM_LOADcO = thr_tmem_load_o.partition_D(tOcO_i) - tTMEM_STOREtO = thr_tmem_store_o.partition_D(tOtO_i) + tTMEM_LOADtO = thr_tmem_load_o.partition_S(tOtO_epi_r[(None, None), None]) + tTMEM_LOADcO = thr_tmem_load_o.partition_D(tOcO_epi_r[(None, None), None]) + tTMEM_STOREtO = thr_tmem_store_o.partition_D(tOtO_epi_r[(None, None), None]) n_corr_tiles = HEAD_DIM // corr_tile_size for kt in range(self.n_kv_tiles): @@ -299,7 +298,7 @@ class FmhaV3StageCMulti: cute.copy(tiled_tmem_store, rP_words, tTMEM_STOREtP) cute.arch.fence_view_async_tmem_store() - # Per-tile O rescale (hand-constructed atoms, only for kt > 0) + # Per-tile O rescale (hand-constructed atoms with logical_divide layout) if kt > 0: tTMrO = cute.make_rmem_tensor( (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype @@ -331,10 +330,6 @@ class FmhaV3StageCMulti: final_o_bar.arrive_and_wait() # === NO-OP TMEM round-trip: re-map O from MMA layout to epilog layout === - # The MMA writes O in the C-fragment TMEM layout, but epilogue_tma_store - # reads using get_tmem_load_op which expects a different layout. A NO-OP - # load-then-store through the hand-constructed atoms re-maps the data. - # TODO: eliminate this by using get_tmem_load_op for the normalize. tTMrO_noop = cute.make_rmem_tensor( (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype ) @@ -357,11 +352,6 @@ class FmhaV3StageCMulti: cute.arch.fence_view_async_tmem_store() # === Final O normalization: O *= 1/row_sum === - # TMEM round-trip using hand-constructed atoms. - # Known issue: hand-constructed Ld32x32bOp/St32x32bOp atoms introduce - # ~3% error due to TMEM column mapping mismatch with get_tmem_load_op. - # TODO: replace with correction_epilog (paired atoms) once TMA store - # region isolation is resolved. inv_row_sum = Float32(1.0) / row_sum tTMrO = cute.make_rmem_tensor(