diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index 48537701..ce44ba31 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -335,67 +335,72 @@ class FmhaKernel: # Wait for MMA's PV[N-1] to commit before reading O. final_o_bar.arrive_and_wait() - # === 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() - - # === Final O normalization: O *= 1/row_sum === + # === Correction epilog: one-way TMEM → reg (normalize) → SMEM → GMEM === + # Uses get_tmem_load_op + get_smem_store_op paired atoms. + # NO TMEM round-trip — hand-constructed atoms corrupt data. inv_row_sum = Float32(1.0) / row_sum - tTMrO = cute.make_rmem_tensor( - (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype + 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), ) - 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 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) - - cute.arch.fence_view_async_tmem_store() - - # 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 ) - 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, + 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) + + 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)) + + 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) + + # Normalize: multiply by inv_row_sum, then convert to BF16 + 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 + 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() diff --git a/tests/unit/test_d1_sweep.py b/tests/unit/test_d1_sweep.py index eb4b587d..934bca8b 100644 --- a/tests/unit/test_d1_sweep.py +++ b/tests/unit/test_d1_sweep.py @@ -1,10 +1,11 @@ -"""D1: Quick test at hd=128 to narrow down the breakage.""" +"""D1 sweep: paired atoms epilogue (no TMEM round-trip).""" import torch, math import cutlass.cute as cute import cutlass.torch as ct import cuda.bindings.driver as cuda from dsv4.kernels.attention.fmha import FmhaKernel +# Single KV tile (n=128) only — O rescale not needed for hd in [64, 128, 256]: torch.manual_seed(42) n = 128; m = 128 @@ -18,7 +19,6 @@ for hd in [64, 128, 256]: attn = qf @ kf.T * scale; attn = torch.softmax(attn, dim=-1) ref = attn @ v.float() - # For hd>256, we'd need N-tiling, but 128 is fine as single tile v_kernel = v.unsqueeze(-1) mQ = ct.from_dlpack(q).mark_layout_dynamic(leading_dim=ct.get_leading_dim(q)) mK = ct.from_dlpack(k).mark_layout_dynamic(leading_dim=ct.get_leading_dim(k)) @@ -34,4 +34,4 @@ for hd in [64, 128, 256]: out = c[:,:,0].float() cos = torch.nn.functional.cosine_similarity(out.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)).item() - print(f'hd={hd}: cos {cos:.6f} {"PASS" if cos >= 0.97 else "FAIL"}') + print(f'hd={hd}: cos {cos:.6f} {"PASS" if cos >= 0.99 else "FAIL (need >=0.99)"}')