From ea5662ab2b368b00d577b047dc33a09dca4d25f3 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 24 May 2026 00:30:38 +0000 Subject: [PATCH] D1.5: Implement correction epilog with get_tmem_load_op + get_smem_store_op paired atoms MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - One-way: TMEM → registers (normalize) → SMEM → GMEM - Eliminates TMEM round-trip error for O normalization - O rescale (kt>0) still uses old atoms (fix later) - Based on CUTLASS FMHA reference's correction_epilog pattern --- dsv4/kernels/attention/fmha.py | 133 +++++++++++++++++---------------- 1 file changed, 68 insertions(+), 65 deletions(-) diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index e39ca3ab..463aba5d 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -9,7 +9,7 @@ from cutlass.cute.nvgpu import cpasync, tcgen05 from cutlass import Float32, BFloat16, Int32, Boolean, const_expr from cutlass.utils import LayoutEnum from cutlass.utils.tmem_allocator import find_tmem_tensor_col_offset -from cutlass.utils.blackwell_helpers import get_smem_store_op +from cutlass.utils.blackwell_helpers import get_smem_store_op, get_tmem_load_op import cuda.bindings.driver as cuda import cutlass.torch as ct import math @@ -413,83 +413,86 @@ class FmhaKernel: # ============================================================ # CORRECTION EPILOG: One-way TMEM → registers → normalize → SMEM → GMEM # ============================================================ - # Replace the old TMEM round-trip (3% error per trip) with the proper - # CUTLASS correction epilog pattern using paired atoms. - # This is a ONE-WAY trip: TMEM → registers (get_tmem_load_op) → - # normalize → SMEM (get_smem_store_op) → GMEM (TMA store). - # No TMEM write-back, no layout mismatch, no data corruption. + # 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. # ============================================================ - # D5a: When normalize=False, we still do the one-way trip but skip the 1/row_sum multiply. + # 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 - # Build the TMEM→reg and reg→SMEM tiled copies using paired atoms + # Correction tile: split O into (128, corr_tile_size) sub-tiles + corr_tile_size = 32 * 8 // self.o_dtype.width # 32*8/16 = 16 for BF16 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 + cO = cute.make_identity_tensor(self.pv_mma_tiler[:2]) + tOcO = pv_thr.partition_C(cO) + tOsO = pv_thr.partition_C(sC) + + tOtO_i = cute.logical_divide(tCtO_base, cute.make_layout((128, corr_tile_size))) + tOcO_i = cute.logical_divide(tOcO, cute.make_layout((128, corr_tile_size))) + tOsO_i = cute.logical_divide(tOsO, cute.make_layout((128, corr_tile_size))) + + # Build TMEM load copy using get_tmem_load_op (paired atom) + epi_subtile = (self.epi_tile[0], corr_tile_size) + tmem_copy_atom = utils.blackwell_helpers.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, ) - 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), + tiled_tmem_load = tcgen05.make_tmem_copy(tmem_copy_atom, tOtO_i[(None, None), 0]) + + # Build SMEM store copy using get_smem_store_op (paired with TMEM load) + smem_copy_atom = utils.blackwell_helpers.get_smem_store_op( + self.c_layout, self.o_dtype, self.acc_dtype, tiled_tmem_load ) + tiled_smem_store = cute.make_tiled_copy_D(smem_copy_atom, tiled_tmem_load) + + # Partition source (TMEM) and destination (SMEM) for each thread + thr_tmem_load = tiled_tmem_load.get_slice(sfw_idx) + thr_smem_store = tiled_smem_store.get_slice(sfw_idx) + + tTMEM_LOADtO = thr_tmem_load.partition_S(tOtO_i[(None, None), None]) + tSMEM_STOREsO = thr_smem_store.partition_D(tOsO_i[(None, None), None]) + tSMEM_STOREcO = thr_smem_store.partition_S(tOcO_i[(None, None), None]) + + n_corr_tiles = self.pv_n_tile // corr_tile_size + + # For each correction tile: TMEM → reg (normalize) → SMEM + for i in range(n_corr_tiles): + tTMEM_LOADtO_i = tTMEM_LOADtO[None, 0, 0, i] + tSMEM_STOREsO_i = tSMEM_STOREsO[None, 0, 0, i] + tSMEM_STOREcO_i = tSMEM_STOREcO[None, 0, 0, i] + tTMrO = cute.make_rmem_tensor(tSMEM_STOREcO_i.shape, self.acc_dtype) + + # Load O from TMEM using paired atom + cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO) + + # Normalize: multiply by inv_row_sum + if const_expr(self.normalize): + for j in cutlass.range(cute.size(tTMrO), vectorize=True): + tTMrO[j] = tTMrO[j] * inv_row_sum + + # 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, tSMrO, tSMEM_STOREsO_i) + + # Fence SMEM writes and sync before TMA store + cute.arch.fence_proxy("async.shared", space="cta") epilog_sync_bar = pipeline.NamedBarrier( barrier_id=self.epilog_sync_bar_id, num_threads=32 * len(self.epilogue_warp_id), ) - - # Consume the accumulator pipeline - 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) - - # 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 - if const_expr(self.normalize): - 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() + + # TMA store: SMEM → GMEM + 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) + c_pipe.producer_acquire() + cute.copy(tma_c, sC, tCgC) + c_pipe.producer_commit() + cute.arch.gpu_bar_sync() c_pipe.producer_tail() # D5a: Write LSE (log-softmax) when normalize=False