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)