From 2438826eee0defbec9ab96545ea221511ef92258 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 27 May 2026 05:31:47 +0000 Subject: [PATCH] Use tma_partition with group_modes on both sC_flat and gO --- dsv4/kernels/attention/fmha_smem_acc.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/dsv4/kernels/attention/fmha_smem_acc.py b/dsv4/kernels/attention/fmha_smem_acc.py index 2e19dddd..aa10b5bb 100644 --- a/dsv4/kernels/attention/fmha_smem_acc.py +++ b/dsv4/kernels/attention/fmha_smem_acc.py @@ -624,10 +624,15 @@ class FmhaKernel: cute.arch.fence_proxy("async.shared", space="cta") # Step 2: TMA store sC_flat -> GMEM - # Use cute.copy with tma_c directly (AutoCopy-style) gO = cute.local_tile(mC, cute.slice_(self.pv_mma_tiler, (None, None, 0)), (None, None, None)) + # Group modes to match: sC_flat is 2D, gO needs to be grouped to 2D + tOsC, tOgO = cpasync.tma_partition( + tma_c, 0, cute.make_layout(1), + cute.group_modes(sC_flat, 0, 2), + cute.group_modes(gO, 0, len(cute.shape(gO))), + ) if warp_idx == self.epilogue_warp_id[0]: - cute.copy(tma_c, sC_flat, gO) + cute.copy(tma_c, tOsC, tOgO) cute.arch.cp_async_bulk_commit_group() cute.arch.cp_async_bulk_wait_group(0, read=True)