From 74e1c0420a0ea5c90200ea5c5fabe6d50a4b150d Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 24 May 2026 00:41:27 +0000 Subject: [PATCH] D1.5: Implement correction epilog with paired atoms (get_tmem_load_op + get_smem_store_op) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit One-way: TMEM → registers (normalize) → SMEM → GMEM Based on CUTLASS FMHA reference's correction_epilog pattern. Eliminates TMEM round-trip error for O normalization. O rescale (kt>0) still uses old atoms (separate fix). --- dsv4/kernels/attention/fmha.py | 142 ++++++++++++++++++++------------- 1 file changed, 88 insertions(+), 54 deletions(-) diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index d3ed7cbe..a5c9a533 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -409,72 +409,106 @@ class FmhaKernel: # Wait for MMA's PV[N-1] to commit before reading O. final_o_bar.arrive_and_wait() - # === NO-OP TMEM round-trip: re-map O from MMA layout to epilog layout === - # TODO: Replace with correction epilog (D1.5) for zero-error one-way trip - 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() + # ============================================================ + # CORRECTION EPILOG: One-way TMEM → registers → normalize → SMEM + # ============================================================ + # Uses paired atoms from get_tmem_load_op + get_smem_store_op + # to preserve the C-fragment layout. No TMEM write-back. + # Based on CUTLASS FMHA reference's correction_epilog pattern. + # Eliminates the 3% per-tile TMEM round-trip error. + # ============================================================ - # === Final O normalization: O *= 1/row_sum === - # D5a: When normalize=False, skip normalization (emit un-normalized O + lse) + # D5a: When normalize=False, still do one-way trip but skip 1/row_sum. if const_expr(self.normalize): inv_row_sum = Float32(1.0) / row_sum - tTMrO = cute.make_rmem_tensor( - (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype - ) + # Step 1: logical_divide O and sC into correction sub-tiles. + tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tOtO.layout) + tOcO = pv_thr.partition_C(cute.make_identity_tensor(self.pv_mma_tiler[:2])) + tOsO = pv_thr.partition_C(sC) + corr_ts = corr_tile_size # sub-tile N-dim (16 for BF16) + tOtO_i = cute.logical_divide(tCtO_base, cute.make_layout((128, corr_ts))) + tOcO_i = cute.logical_divide(tOcO, cute.make_layout((128, corr_ts))) + tOsO_i = cute.logical_divide(tOsO, cute.make_layout((128, corr_ts))) + # Step 2: Build TMEM load copy using get_tmem_load_op (paired atom). + epi_subtile = (self.epi_tile[0], corr_ts) + from cutlass.utils.blackwell_helpers import get_tmem_load_op as _get_tmem_load_op + tmem_copy_atom = _get_tmem_load_op( + self.pv_mma_tiler, self.c_layout, self.o_dtype, self.acc_dtype, + epi_subtile, use_2cta_instrs=self.use_2cta_instrs, + ) + # tOtO_i has shape ((128, corr_ts), n_corr_tiles) after logical_divide. + # make_tmem_copy needs a tensor with the sub-tile layout. + # Slice to the first sub-tile to get the right layout for the copy atom. + tOtO_sub0 = tOtO_i[(None, None), 0] # first sub-tile + tiled_tmem_load_corr = tcgen05.make_tmem_copy(tmem_copy_atom, tOtO_sub0) + + # Step 3: Build SMEM store copy using get_smem_store_op (paired with TMEM load). + smem_copy_atom = get_smem_store_op( + self.c_layout, self.o_dtype, self.acc_dtype, tiled_tmem_load_corr + ) + tiled_smem_store_corr = cute.make_tiled_copy_D(smem_copy_atom, tiled_tmem_load_corr) + + # Step 4: Partition source (TMEM) and destination (SMEM) for each softmax thread. + thr_tmem_corr = tiled_tmem_load_corr.get_slice(sfw_idx) + thr_smem_corr = tiled_smem_store_corr.get_slice(sfw_idx) + # Partition the sub-tiled O for the correction loop. + tTMEM_CORRtO = thr_tmem_corr.partition_S(tOtO_i[(None, None), None]) + tSMEM_CORRsO = thr_smem_corr.partition_D(tOsO_i[(None, None), None]) + tSMEM_CORRcO = thr_smem_corr.partition_S(tOcO_i[(None, None), None]) + + # Step 5: Correction loop — for each sub-tile: TMEM → reg (normalize) → SMEM 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 - ) + tTMEM_CORRtO_i = tTMEM_CORRtO[None, 0, 0, i] + tSMEM_CORRsO_i = tSMEM_CORRsO[None, 0, 0, i] + # Create register tensor for this sub-tile using the SMEM copy's source layout + tTMrO = cute.make_rmem_tensor(tSMEM_CORRcO[None, 0, 0, i].shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO_i) + # Load O from TMEM using paired atom (preserves C-fragment layout) + cute.copy(tiled_tmem_load_corr, tTMEM_CORRtO_i, tTMrO) + + # Normalize: multiply by inv_row_sum (exact in FP32) if const_expr(self.normalize): - 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) + for j in cutlass.range(cute.size(tTMrO), vectorize=True): + tTMrO[j] = tTMrO[j] * inv_row_sum - cute.arch.fence_view_async_tmem_store() + # Convert to output dtype and store to SMEM via paired atom + 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_corr, tSMrO, tSMEM_CORRsO_i) - # Epilogue: TMEM → SMEM → GMEM via TMA store. - 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 + # Fence SMEM writes and sync before TMA store + cute.arch.fence_proxy("async.shared", space="cta") + # Barrier: ensure all softmax warps have finished writing to SMEM + # before TMA store reads from it. Use a separate barrier ID. + corr_epi_bar = pipeline.NamedBarrier( + barrier_id=5, num_threads=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, + corr_epi_bar.arrive_and_wait() + + # Step 6: TMA store SMEM → GMEM + # The normalized O is now in sC (written by the correction epilog). + # TMA store from sC to the output tensor in GMEM. + # Use cpasync.tma_partition to set up the SMEM/GMEM partition. + gC = cute.local_tile(mC, epi_tile, (Int32(0), Int32(0))) + tCgC_epi = cute.flat_divide(tCgC, 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, 0, 2), ) - c_pipe.producer_tail() + # One TMA store for the full output tile + if warp_idx == self.epilogue_warp_id[0]: + 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_pipe.producer_acquire() + cute.copy(tma_c, bSG_sC[(None, 0)], bSG_gC[(None, 0)]) + c_pipe.producer_commit() + c_pipe.producer_tail() # D5a: Write LSE (log-softmax) when normalize=False # lse = ln(row_sum) + row_max * ln(2)