diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index fac28721..46ba5611 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, get_tmem_load_op +from cutlass.utils.blackwell_helpers import get_smem_store_op import cuda.bindings.driver as cuda import cutlass.torch as ct import math @@ -27,7 +27,6 @@ class FmhaKernel: self.acc_dtype = Float32; self.qk_acc_dtype = Float32 self.q_dtype = BFloat16; self.o_dtype = BFloat16; self.c_dtype = BFloat16 self.use_2cta_instrs = False; self.epilog_sync_bar_id = 1 - self.iter_acc_early_release_in_epilogue = 0 # No early release self.cluster_shape_mn = (1, 1); self.cta_group = tcgen05.CtaGroup.ONE self.epilogue_warp_id = (0,1,2,3); self.mma_warp_id = 4; self.tma_warp_id = 5 self.threads_per_cta = 192; self.num_c_stage = 2 @@ -410,93 +409,70 @@ class FmhaKernel: # Wait for MMA's PV[N-1] to commit before reading O. final_o_bar.arrive_and_wait() - # ============================================================ - # CORRECTION EPILOG: One-way TMEM → registers → normalize → SMEM → GMEM - # ============================================================ - # 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. - # ============================================================ + # === NO-OP TMEM round-trip: re-map O from MMA layout to epilog layout === + 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() - # D5a: When normalize=False, still do one-way trip but skip 1/row_sum. + # === Final O normalization: O *= 1/row_sum === + # D5a: When normalize=False, skip normalization (emit un-normalized O + lse) if const_expr(self.normalize): inv_row_sum = Float32(1.0) / row_sum - # 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) - 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, + tTMrO = cute.make_rmem_tensor( + (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype ) - 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) + 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 + ) - # Load O from TMEM using paired atom - cute.copy(tiled_tmem_load, tTMEM_LOADtO_i, tTMrO) - - # Normalize: multiply by inv_row_sum + cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO_i) if const_expr(self.normalize): - for j in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[j] = tTMrO[j] * inv_row_sum + 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) - # 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) + cute.arch.fence_view_async_tmem_store() - # 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), + # 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 ) - epilog_sync_bar.arrive_and_wait() - - # TMA store: SMEM → GMEM - # Reuse the existing TMA partition (tCgC) which was set up at kernel start. - # sC was written by the correction epilog. TMA reads from sC → GMEM via tCgC. 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() - # TMA store from sC to GMEM using the pre-partitioned gC - gC = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler,(None,0,None)),(None,None,None)) - cute.copy(tma_c, cute.select(sC, mode=[0, 1]), cute.select(gC, mode=[0, 1])) - c_pipe.producer_commit() - cute.arch.gpu_bar_sync() + 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, + ) c_pipe.producer_tail() # D5a: Write LSE (log-softmax) when normalize=False