diff --git a/tests/fmha_v3_real_softmax.py b/tests/fmha_v3_real_softmax.py index de4a430d..0a3cc287 100644 --- a/tests/fmha_v3_real_softmax.py +++ b/tests/fmha_v3_real_softmax.py @@ -226,7 +226,23 @@ class FmhaV3RealSoftmax: tScP = cute.make_tensor(tScS.iterator, tScP_layout) tTMEM_STOREcP = thr_store.partition_S(tScP) - # TODO: O normalize setup (add when O rescale is ready) + # O normalize setup: sub-tile O for TMEM read-modify-write + cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) + tOcO = pv_thr.partition_C(cO) + corr_tile_size = 16 + tOtO_i_layout = cute.composition(tOtO.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) + tmem_load_o_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype) + tmem_store_o_atom = cute.make_copy_atom(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) + thr_load_o = tiled_tmem_load_o.get_slice(sfw_idx) + thr_store_o = tiled_tmem_store_o.get_slice(sfw_idx) + tTMEM_LOAD_OtO = thr_load_o.partition_S(tOtO_i) + tTMEM_LOAD_OcO = thr_load_o.partition_D(tOcO_i) + tTMEM_STORE_OtO = thr_store_o.partition_D(tOtO_i) row_max = -Float32.inf row_sum = Float32(0.0) @@ -281,10 +297,30 @@ class FmhaV3RealSoftmax: softmax_done_bar.arrive() # Final O normalization: O = O / row_sum - # TODO: enable after basic softmax works - # if row_sum != Float32(0.0): - # inv_row_sum = Float32(1.0) / row_sum - # ... + if row_sum != Float32(0.0): + inv_row_sum = Float32(1.0) / row_sum + # Load O from TMEM, multiply by 1/row_sum, write back + n_corr = 128 // corr_tile_size + for ci in range(n_corr): + tOtO_ci = cute.make_tensor(tOtO_i.iterator, tOtO_i.layout) + tOcO_ci = cute.make_tensor(tOcO_i.iterator, tOcO_i.layout) + # Sub-tile index offset + tOtO_sub = tOtO_ci[None, ci, None] + tOcO_sub = tOcO_ci[None, ci, None] + # Use the base partitioned tensors with offset + # Actually, just load the full O sub-tile + pass + # Simple approach: load/store via the partitioned tensors + tTMEM_LOAD_OrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype) + cute.copy(tiled_tmem_load_o, tTMEM_LOAD_OtO, tTMEM_LOAD_OrO) + cute.arch.fence_view_async_tmem_load() + # Iterate with proper indexing + n_o_frg = cute.size(tTMEM_LOAD_OrO, mode=[0]) + for fi in range(n_o_frg): + for fj in range(cute.size(tTMEM_LOAD_OrO, mode=[1])): + tTMEM_LOAD_OrO[fi, fj, None] = tTMEM_LOAD_OrO[fi, fj, None] * inv_row_sum + cute.copy(tiled_tmem_store_o, tTMEM_LOAD_OrO, tTMEM_STORE_OtO) + cute.arch.fence_view_async_tmem_store() # Epilogue: TMEM -> SMEM -> GMEM via TMA store tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout)