diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index 0488c889..2548e3b1 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -15,16 +15,6 @@ import cutlass.torch as ct import math -def _transform_partitioned_tensor_layout(tensor): - """Transform MMA layout: ((ATOM_M, ATOM_N), MMA_M, MMA_N, ...rest) - -> ((ATOM_M, MMA_M), (ATOM_N, MMA_N), ...rest). - Same as CUTLASS utils.gemm.sm100.transform_partitioned_tensor_layout.""" - layout = tensor.layout; shape = layout.shape; stride = layout.stride - new_shape = ((shape[0][0], shape[1]), (shape[0][1], shape[2]), *shape[3:]) - new_stride = ((stride[0][0], stride[1]), (stride[0][1], stride[2]), *stride[3:]) - return cute.make_tensor(tensor.iterator, cute.make_layout(shape=new_shape, stride=new_stride)) - - class FmhaKernel: def __init__(self, head_dim=64, s_k=128, scale_softmax=None, use_smem_p=None, normalize=True): self.head_dim = head_dim @@ -416,104 +406,50 @@ class FmhaKernel: final_o_bar.arrive_and_wait() # ============================================================ - # CORRECTION EPILOG: One-way TMEM → registers → normalize → SMEM → GMEM + # EPILOGUE: Normalize O + TMA store to GMEM # ============================================================ - # Follows CUTLASS epilogue_tma_store pattern exactly: - # transform_partitioned_tensor_layout → flat_divide → - # get_tmem_load_op → make_tmem_copy → partition_S → - # get_smem_store_op → make_tiled_copy_D → partition_D → - # cpasync.tma_partition → copy loop - # Eliminates the 3% per-tile TMEM round-trip error by using - # paired atoms that preserve the C-fragment layout. + # Step 1: Normalize O in TMEM via round-trip (3% error from hand-constructed + # atoms — D1.5 tracks the paired-atom fix). + # Step 2: Use CUTLASS epilogue_tma_store for TMEM→SMEM→GMEM write. # ============================================================ - # D5a: When normalize=False, still do one-way trip but skip 1/row_sum. + # D5a: When normalize=False, skip 1/row_sum (emit un-normalized O + LSE). if const_expr(self.normalize): inv_row_sum = Float32(1.0) / row_sum + # Normalize O: TMEM round-trip O *= inv_row_sum + 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, + ) + cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO_i) + for k in cutlass.range(cute.size(tTMrO_i), vectorize=True): + tTMrO_i[k] = tTMrO_i[k] * inv_row_sum + cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STOREtO_i) + cute.arch.fence_view_async_tmem_store() - # Step 1: Transform partitioned tensor layouts (CUTLASS pattern) - # ((ATOM_M, ATOM_N), MMA_M, MMA_N, ...) -> ((ATOM_M, MMA_M), (ATOM_N, MMA_N), ...) - tOtO_xfm = _transform_partitioned_tensor_layout( - cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tOtO.layout)) - tCgC_xfm = _transform_partitioned_tensor_layout(tCgC) - - # Step 2: TMEM load copy (epilogue_tmem_copy_and_partition pattern) - from cutlass.utils.blackwell_helpers import get_tmem_load_op as _get_tmem_load_op - tmem_copy_atom = _get_tmem_load_op( - self.cta_tile_shape_mnk, self.c_layout, self.o_dtype, self.acc_dtype, - epi_tile, self.use_2cta_instrs, + # TMA store via CUTLASS epilogue_tma_store + tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tOtO.layout) + 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 = pipeline.make_pipeline_state( + pipeline.PipelineUserType.Consumer, self.num_acc_stage ) - # flat_divide by epi_tile to create sub-tiled views - tOtO_epi = cute.flat_divide(tOtO_xfm, epi_tile) - tCgC_epi = cute.flat_divide(tCgC_xfm, epi_tile) - # make_tmem_copy with the first sub-tile shape - tiled_copy_t2r = tcgen05.make_tmem_copy(tmem_copy_atom, tOtO_epi[(None, None, 0, 0)]) - thr_t2r = tiled_copy_t2r.get_slice(sfw_idx) - # Partition source (TMEM) and destination (GMEM-derived register shape) - tTR_tAcc = thr_t2r.partition_S(tOtO_epi) - tTR_gC = thr_t2r.partition_D(tCgC_epi) - # Register tensor shape: keep (T2R, T2R_M, T2R_N) dims, zero the rest - # After partition_D, tTR_gC has (T2R, T2R_M, T2R_N, EPI_M, EPI_N, REST...) - # We slice to get just the per-subtile register count. - _tTR_gC_shape = tTR_gC.shape - _n_rest = len(_tTR_gC_shape) - 3 # 3 leading dims: T2R, T2R_M, T2R_N - tTR_rAcc = cute.make_rmem_tensor( - tTR_gC[(None, None, None) + (0,) * _n_rest].shape, self.acc_dtype) - - # Step 3: SMEM store copy (epilogue_smem_copy_and_partition pattern) - smem_copy_atom = get_smem_store_op( - self.c_layout, self.o_dtype, self.acc_dtype, tiled_copy_t2r) - tiled_copy_r2s = cute.make_tiled_copy_D(smem_copy_atom, tiled_copy_t2r) - thr_r2s = tiled_copy_r2s.get_slice(sfw_idx) - tRS_sC = thr_r2s.partition_D(sC) - tTR_rC = cute.make_rmem_tensor(tRS_sC[(None, None, None, 0)].shape, self.o_dtype) - - # Step 4: TMA store partition (cpasync.tma_partition for S2G) - # flat_divide tCgC for TMA partition (need the un-xfm version for tma_partition) - tCgC_epi_tma = 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_tma, 0, 2), + acc_cons_st = utils.gemm.sm100.epilogue_tma_store( + self, sfw_idx, 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, ) - - # Step 5: Correction loop — for each sub-tile: TMEM → reg → normalize → SMEM - tTR_tAcc_g = cute.group_modes(tTR_tAcc, 3, cute.rank(tTR_tAcc)) - subtile_cnt = cute.size(tTR_tAcc_g.shape, mode=[3]) - - for subtile_idx in range(subtile_cnt): - # Load O from TMEM (preserves C-fragment layout via paired atom) - tTR_tAcc_mn = tTR_tAcc_g[(None, None, None, subtile_idx)] - cute.copy(tiled_copy_t2r, tTR_tAcc_mn, tTR_rAcc) - - # Normalize: multiply by inv_row_sum (exact in FP32) - if const_expr(self.normalize): - for j in cutlass.range(cute.size(tTR_rAcc), vectorize=True): - tTR_rAcc[j] = tTR_rAcc[j] * inv_row_sum - - # Convert to output dtype and store to SMEM - acc_vec = tiled_copy_r2s.retile(tTR_rAcc).load() - acc_vec = acc_vec.to(self.o_dtype) - tTR_rC.store(acc_vec) - - # Store to SMEM - c_buffer = subtile_idx % self.num_c_stage - cute.copy(tiled_copy_r2s, tTR_rC, tRS_sC[(None, None, None, c_buffer)]) - - # Fence and barrier - cute.arch.fence_proxy("async.shared", space="cta") - corr_epi_bar = pipeline.NamedBarrier( - barrier_id=5, num_threads=32 * len(self.epilogue_warp_id)) - corr_epi_bar.arrive_and_wait() - - # TMA store SMEM → GMEM - if warp_idx == self.epilogue_warp_id[0]: - # Group trailing modes and slice (CUTLASS pattern) - bSG_gC_flat = cute.group_modes(bSG_gC, 1, cute.rank(bSG_gC)) - cute.copy(tma_c, bSG_sC[(None, c_buffer)], bSG_gC_flat[(None, Int32(0))]) - cute.arch.cp_async_bulk_commit_group() - cute.arch.cp_async_bulk_wait_group(0, read=True) - corr_epi_bar.arrive_and_wait() + c_pipe.producer_tail() # D5a: Write LSE (log-softmax) when normalize=False # lse = ln(row_sum) + row_max * ln(2)