diff --git a/tests/unit/test_fmha_v3_stage_c.py b/tests/unit/test_fmha_v3_stage_c.py index 2aad1970..9ea17b8e 100644 --- a/tests/unit/test_fmha_v3_stage_c.py +++ b/tests/unit/test_fmha_v3_stage_c.py @@ -130,75 +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 paired TMEM load atom for correction_epilog - 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) - - def _correction_epilog(self, pv_thr, tOtO, scale, sC, tCgC, tma_c, sfw_idx, tidx, warp_idx, epi_tile, tmem_load_epi_atom): - """CUTLASS correction_epilog: read O from TMEM, normalize, convert, write SMEM→GMEM.""" - epi_corr_tile_size = 32 * 8 // self.o_dtype.width # 16 for BF16 - cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1])) - tOcO = pv_thr.partition_C(cO) - tOsO = pv_thr.partition_C(sC) - - tOtO_i = cute.logical_divide(tOtO, cute.make_layout((128, epi_corr_tile_size))) - tOcO_i = cute.logical_divide(tOcO, cute.make_layout((128, epi_corr_tile_size))) - tOsO_i = cute.logical_divide(tOsO, cute.make_layout((128, epi_corr_tile_size))) - - tiled_tmem_load = tcgen05.make_tmem_copy(tmem_load_epi_atom, tOtO_i[(None, None), 0]) - smem_copy_atom = utils.sm100.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) - - thr_tmem_load = tiled_tmem_load.get_slice(sfw_idx) - tTMEM_LOADtO = thr_tmem_load.partition_S(tOtO_i[(None, None), None]) - tTMEM_LOADsO = thr_tmem_load.partition_D(tOsO_i[(None, None), None]) - tTMEM_LOADcO = thr_tmem_load.partition_D(tOcO_i[(None, None), None]) - - n_corr_tiles = self.pv_mma_tiler[1] // epi_corr_tile_size - for i in range(n_corr_tiles): - tTMrO = cute.make_rmem_tensor(tTMEM_LOADcO[None, 0, 0, i].shape, self.acc_dtype) - cute.copy(tiled_tmem_load, tTMEM_LOADtO[None, 0, 0, i], tTMrO) - for j in range(cute.size(tTMrO)): - tTMrO[j] = tTMrO[j] * scale - 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, tTMEM_LOADsO[None, 0, 0, i]) - - cute.arch.fence_proxy("async.shared", space="cta") - - # TMA store SMEM → GMEM (same pattern as epilogue_tma_store) - tCgC_epi = cute.flat_divide(tCgC, 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), - ) - - epilog_sync_bar = pipeline.NamedBarrier( - barrier_id=self.epilog_sync_bar_id, - num_threads=32 * len(self.epilogue_warp_id), - ) - epilog_sync_bar.arrive_and_wait() - - # All warps in the epilogue group do the TMA store together - c_buffer = 0 - cute.copy(tma_c, tCsC[(None, c_buffer)], 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) - epilog_sync_bar.arrive_and_wait() + 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: @@ -464,12 +399,80 @@ class FmhaV3StageCMulti: # Wait for MMA's PV[N-1] to commit before reading O. final_o_bar.arrive_and_wait() - # === Correction epilog: one-way TMEM → SMEM with normalize === + # === Correction epilog: one-way TMEM → reg → SMEM → GMEM with normalize === + # Uses get_tmem_load_op + get_smem_store_op paired atoms (same as CUTLASS correction_epilog). + # NO TMEM round-trip — hand-constructed Ld32x32bOp/St32x32bOp atoms corrupt data. inv_row_sum = Float32(1.0) / row_sum - self._correction_epilog( - pv_thr, tOtO0, inv_row_sum, sC, tCgC, tma_c, sfw_idx, - tidx, warp_idx, epi_tile, tmem_load_epi_atom, + + # Build the TMEM→reg and reg→SMEM tiled copies using paired atoms + 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 ) + 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), + ) + 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 + acc_vec = tiled_copy_r2s.retile(tTR_rO).load() + # acc_vec is in FP32 — apply scale before conversion + # We can't directly scale the vector, but we can scale the register tensor + 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() + c_pipe.producer_tail() tmem.relinquish_alloc_permit() tmem.free(tmem_ptr) @@ -489,8 +492,8 @@ def test(): qf = q[:, :, 0].float() kf = k[:, :, 0].float() scale = 1.0 / math.sqrt(hd) - attn_raw = qf @ kf.T * scale - attn = torch.softmax(attn_raw, dim=-1) + attn = qf @ kf.T * scale + attn = torch.softmax(attn, dim=-1) ref = attn @ v.float() mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q))