From e7b2cb5a32a060b408d961faa14784ba01db8f93 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 01:25:53 +0000 Subject: [PATCH] fix: O rescale uses 2D register tensor pattern (matching CUTLASS correction_rescale) --- tests/unit/test_fmha_v3_stage_c.py | 117 ++++++++++++----------------- 1 file changed, 48 insertions(+), 69 deletions(-) diff --git a/tests/unit/test_fmha_v3_stage_c.py b/tests/unit/test_fmha_v3_stage_c.py index e11dcac0..62f7f997 100644 --- a/tests/unit/test_fmha_v3_stage_c.py +++ b/tests/unit/test_fmha_v3_stage_c.py @@ -130,19 +130,10 @@ class FmhaV3StageCMulti: tma_v,mV = cute.nvgpu.make_tiled_tma_atom_B(utils.sm100.cluster_shape_to_tma_atom_B(self.cluster_shape_mn,pv_mma.thr_id),v_fmha,v_s,self.pv_mma_tiler,pv_mma,self.cluster_layout_vmnk.shape) epi_s = cute.select(self.c_smem_s,mode=[0,1]) tma_c,mC = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileS2GOp(),c,epi_s,self.epi_tile) - - # Pre-compute epilog atoms for correction_epilog pattern - epi_corr_tile_size = 32 * 8 // self.o_dtype.width # 16 for BF16 - epi_subtile = (self.epi_tile[0], epi_corr_tile_size) - tmem_load_epi_atom = utils.sm100.get_tmem_load_op( - self.pv_mma_tiler, self.c_layout, self.o_dtype, self.acc_dtype, - epi_subtile, use_2cta_instrs=False, - ) - - self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile,tmem_load_epi_atom).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) + self._kernel(qk_mma,pv_mma,tma_q,mQ,tma_k,mK,tma_v,mV,tma_c,mC,self.cluster_layout_vmnk,self.q_smem_s,self.k_smem_s,self.v_smem_s,self.p_tmem_s,self.c_smem_s,self.epi_tile).launch(grid=(1,1,1),block=[self.threads_per_cta,1,1],stream=stream) @cute.kernel - def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile, tmem_load_epi_atom): + def _kernel(self, qk_mma, pv_mma, tma_q, mQ, tma_k, mK, tma_v, mV, tma_c, mC, cl_vmnk, q_smem_s, k_smem_s, v_smem_s, p_tmem_s, c_smem_s, epi_tile): warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) tidx,_,_ = cute.arch.thread_idx() if warp_idx == self.tma_warp_id: @@ -384,8 +375,18 @@ class FmhaV3StageCMulti: cute.arch.fence_view_async_tmem_store() # === Per-tile O rescale: O *= acc_scale for kt > 0 === + # Uses 2D register tensor pattern (same as CUTLASS correction_rescale + # and our final normalize) to preserve data through TMEM round-trip. if kt > 0: + tTMrO = cute.make_rmem_tensor( + (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype + ) 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, @@ -394,12 +395,10 @@ class FmhaV3StageCMulti: tTMEM_STOREtO.iterator + i * corr_tile_size, tTMEM_STOREtO.layout, ) - tTMrO = cute.make_rmem_tensor(tTMEM_LOADcO.shape, self.acc_dtype) - cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO) - cute.arch.fence_view_async_tmem_load() - for k in cutlass.range(cute.size(tTMrO), vectorize=True): - tTMrO[k] = tTMrO[k] * acc_scale - cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STOREtO_i) + 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] * acc_scale + cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STOREtO_i) cute.arch.fence_view_async_tmem_store() si_handle.release() @@ -408,67 +407,47 @@ class FmhaV3StageCMulti: # Wait for MMA's PV[N-1] to commit before reading O. final_o_bar.arrive_and_wait() - # === Final O normalization + epilogue: CUTLASS correction_epilog pattern === - # ONE-WAY trip: TMEM → reg (normalize + FP32→BF16) → SMEM → TMA → GMEM + # === Final O normalization: O *= 1/row_sum === inv_row_sum = Float32(1.0) / row_sum - # Build paired SMEM store atom from the pre-computed TMEM load atom - epi_corr_tile_size = 32 * 8 // self.o_dtype.width # 16 for BF16 - tOtO_epi = cute.logical_divide(tOtO0, cute.make_layout((128, epi_corr_tile_size))) - tiled_tmem_load_epi = tcgen05.make_tmem_copy( - tmem_load_epi_atom, tOtO_epi[(None, None), 0] - ) - smem_store_epi_atom = utils.sm100.get_smem_store_op( - self.c_layout, self.o_dtype, self.acc_dtype, tiled_tmem_load_epi, - ) - tiled_smem_store_epi = cute.make_tiled_copy_D( - smem_store_epi_atom, tiled_tmem_load_epi + tTMrO = cute.make_rmem_tensor( + (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype ) - # Partition SMEM for the epilog output - tOsO = pv_thr.partition_C(sC) - cO_epi = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO_epi = pv_thr.partition_C(cO_epi) - tOsO_epi = cute.logical_divide(tOsO, cute.make_layout((128, epi_corr_tile_size))) - tOcO_epi = cute.logical_divide(tOcO_epi, cute.make_layout((128, epi_corr_tile_size))) - - thr_tmem_load_epi = tiled_tmem_load_epi.get_slice(sfw_idx) - tTMEM_LOADtO_epi = thr_tmem_load_epi.partition_S(tOtO_epi[(None, None), None]) - tTMEM_LOADsO_epi = thr_tmem_load_epi.partition_D(tOsO_epi[(None, None), None]) - tTMEM_LOADcO_epi = thr_tmem_load_epi.partition_D(tOcO_epi[(None, None), None]) - - n_epi_corr_tiles = self.pv_mma_tiler[1] // epi_corr_tile_size - for i in range(n_epi_corr_tiles): - tTMEM_LOADtO_epi_i = tTMEM_LOADtO_epi[None, 0, 0, i] - tTMEM_LOADsO_epi_i = tTMEM_LOADsO_epi[None, 0, 0, i] - tTMrO = cute.make_rmem_tensor( - tTMEM_LOADcO_epi[None, 0, 0, i].shape, self.acc_dtype + 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_epi, tTMEM_LOADtO_epi_i, tTMrO) - for j in range(cute.size(tTMrO)): - tTMrO[j] = tTMrO[j] * inv_row_sum - 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_epi, tSMrO, tTMEM_LOADsO_epi_i) - cute.arch.fence_proxy("async.shared", space="cta") + cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO_i) + 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) - # TMA store: SMEM → GMEM - softmax_all_bar = pipeline.NamedBarrier( - barrier_id=5, num_threads=32 * len(self.epilogue_warp_id) + cute.arch.fence_view_async_tmem_store() + + # Standard epilogue: TMEM → SMEM → GMEM via TMA store. + # O in TMEM is now scaled by 1/row_sum. + 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 ) - softmax_all_bar.arrive_and_wait() - - tCgC_epi = cute.flat_divide(tCgC, self.epi_tile) - tCsC, tCgC_tma = cpasync.tma_partition( - tma_c, 0, cute.make_layout(1), - cute.group_modes(sC, 0, 2), - cute.group_modes(tCgC_epi, 0, 2), + 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, ) - cute.copy(tma_c, tCsC[(None, 0)], tCgC_tma[(None, 0, 0, 0, 0, 0, 0)]) - cute.arch.cp_async_bulk_commit_group() - cute.arch.cp_async_bulk_wait_group(0, read=True) + c_pipe.producer_tail() tmem.relinquish_alloc_permit() tmem.free(tmem_ptr)