diff --git a/dsv4/kernels/attention/fmha_smem_acc.py b/dsv4/kernels/attention/fmha_smem_acc.py
index 1e1d326b..dd2c1095 100644
--- a/dsv4/kernels/attention/fmha_smem_acc.py
+++ b/dsv4/kernels/attention/fmha_smem_acc.py
@@ -616,22 +616,24 @@ class FmhaKernel:
                 # Path 2: write sO_acc (FP32) -> sC_flat (BF16) -> TMA store to GMEM
                 # sC_flat has epi_s layout (same as what tma_c was created from)
 
-                # Step 1: Cast sO_acc -> sC_flat (BF16)
+                # Step 1: Cast sO_acc -> sC_flat (BF16) with staged layout
                 for row in cutlass.range(0, 128, unroll=1):
                     for col in cutlass.range(0, self.pv_n_tile, unroll=1):
                         val = sO_acc[row, col]
                         if const_expr(self.normalize):
                             inv_row_sum = Float32(1.0) / row_sum
                             val = val * inv_row_sum
-                        sC_flat[row, col] = val.to(self.o_dtype)
+                        sC_flat_staged[row, col // self.num_c_stage, col % self.num_c_stage] = val.to(self.o_dtype)
                 cute.arch.fence_proxy("async.shared", space="cta")
 
                 # Step 2: TMA store sC_flat -> GMEM
                 # Use tCgC (already partitioned) for the GMEM side of TMA
                 tCgC_xfm = transform_partitioned_tensor_layout(tCgC)
                 tCgC_epi = cute.flat_divide(tCgC_xfm, epi_tile)
-                # sC_flat (128, pv_n_tile) -> split to match TMA stage: (128, pv_n_tile//2, 2)
-                sC_flat_staged = cute.logical_divide(sC_flat, cute.make_layout((128, self.pv_n_tile // 2, 2), stride=(self.pv_n_tile, 2, 1)))
+                # sC_flat (128, pv_n_tile, stride=(pv_n_tile, 1))
+                # Split pv_n_tile into (pv_n_tile//2, 2) to match TMA stage decomposition
+                # Result: (128, pv_n_tile//2, 2, stride=(pv_n_tile, 2, 1))
+                sC_flat_staged = cute.make_tensor(sC_flat.iterator, cute.make_layout((128, self.pv_n_tile // self.num_c_stage, self.num_c_stage), stride=(self.pv_n_tile, self.num_c_stage, 1)))
                 tOsC, tOgO = cpasync.tma_partition(
                     tma_c, 0, cute.make_layout(1),
                     sC_flat_staged,