diff --git a/dsv4/kernels/attention/fmha_smem_acc.py b/dsv4/kernels/attention/fmha_smem_acc.py
index d6cf9004..15ec7b42 100644
--- a/dsv4/kernels/attention/fmha_smem_acc.py
+++ b/dsv4/kernels/attention/fmha_smem_acc.py
@@ -226,7 +226,7 @@ class FmhaKernel:
             sO_acc = smem.allocate_tensor(element_type=Float32, layout=sO_acc_layout, byte_alignment=128)
             # sC_flat: BF16 SMEM buffer with epi_s layout (non-swizzled) for TMA store
             # Used to cast sO_acc (FP32) -> BF16 and TMA store to GMEM
-            sC_flat = smem.allocate_tensor(element_type=self.o_dtype, layout=cute.select(self.c_smem_s, mode=[0, 1]).outer, byte_alignment=128)
+            sC_flat = smem.allocate_tensor(element_type=self.o_dtype, layout=cute.make_layout((128, self.pv_n_tile), stride=(self.pv_n_tile, 1)), byte_alignment=128)
         else:
             sO_acc = smem.allocate_tensor(element_type=Float32, layout=cute.make_layout((1, 1), stride=(1, 1)), byte_alignment=128)
             sC_flat = smem.allocate_tensor(element_type=self.o_dtype, layout=cute.make_layout((1, 1), stride=(1, 1)), byte_alignment=128)