diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py
index 67de010b..89ece768 100644
--- a/dsv4/kernels/attention/fmha.py
+++ b/dsv4/kernels/attention/fmha.py
@@ -362,13 +362,23 @@ class FmhaKernel:
                 # For now, zero sP as a stub — PV will read garbage/zero
                 pass
 
-            # ── O rescale / normalization setup (correction_rescale atoms) ──
+            # ── O rescale / normalization setup (correction_rescale pattern from Stage C) ──
             corr_tile_size = 16
-            o_rescale_atom_ld = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype)
-            o_rescale_atom_st = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(32)), self.qk_acc_dtype)
-            o_rescale_layout = cute.composition(tStS.layout, cute.make_layout((self.pv_mma_tiler[0], corr_tile_size)))
-            tiled_o_ld = tcgen05.make_tmem_copy(o_rescale_atom_ld, cute.make_tensor(tStS.iterator + self.tmem_o0_offset, o_rescale_layout))
-            tiled_o_st = tcgen05.make_tmem_copy(o_rescale_atom_st, cute.make_tensor(tStS.iterator + self.tmem_o0_offset, o_rescale_layout))
+            tOcO = pv_thr.partition_C(cS)
+            tOtO_i_layout = cute.composition(tOtO0.layout, cute.make_layout((128, corr_tile_size)))
+            tOcO_i_layout = cute.composition(tOcO.layout, cute.make_layout((128, corr_tile_size)))
+            tOtO_i = cute.make_tensor(tOtO0.iterator, tOtO_i_layout)
+            tOcO_i = cute.make_tensor(tOcO.iterator, tOcO_i_layout)
+            tmem_load_o_atom = cute.make_copy_atom(tcgen05.copy.Ld32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype)
+            tmem_store_o_atom = cute.make_copy_atom(tcgen05.copy.St32x32bOp(tcgen05.copy.Repetition(corr_tile_size)), self.acc_dtype)
+            tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_o_atom, tOtO_i)
+            tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_i)
+            thr_tmem_load_o = tiled_tmem_load_o.get_slice(sfw_idx)
+            thr_tmem_store_o = tiled_tmem_store_o.get_slice(sfw_idx)
+            tTMEM_LOADtO = thr_tmem_load_o.partition_S(tOtO_i)
+            tTMEM_LOADcO = thr_tmem_load_o.partition_D(tOcO_i)
+            tTMEM_STOREtO = thr_tmem_store_o.partition_D(tOtO_i)
+            n_corr_tiles = self.head_dim // corr_tile_size
 
             # ── Online softmax state ──
             row_max = -Float32.inf
@@ -441,30 +451,35 @@ class FmhaKernel:
 
                 # ── Per-tile O rescale (multiply O by acc_scale when kt > 0) ──
                 if kt > 0:
-                    thr_ld = tiled_o_ld.get_slice(sfw_idx)
-                    thr_st = tiled_o_st.get_slice(sfw_idx)
-                    tOrO_src = thr_ld.partition_S(cute.make_tensor(tStS.iterator + self.tmem_o0_offset, o_rescale_layout))
-                    tOrO_dst = thr_st.partition_D(cute.make_tensor(tStS.iterator + self.tmem_o0_offset, o_rescale_layout))
-                    rO = cute.make_rmem_tensor(tOrO_src.shape, self.qk_acc_dtype)
-                    cute.copy(tiled_o_ld, tOrO_src, rO)
-                    for i in cutlass.range(cute.size(rO), vectorize=True):
-                        rO[i] = rO[i] * acc_scale
-                    cute.copy(tiled_o_st, rO, tOrO_dst)
+                    tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype)
+                    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 k in cutlass.range(cute.size(tTMrO_i), vectorize=True):
+                            tTMrO_i[k] = tTMrO_i[k] * acc_scale
+                        cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STOREtO_i)
+                    cute.arch.fence_view_async_tmem_store()
 
             # ── Wait for MMA's final PV GEMM ──
             final_o_bar.arrive_and_wait()
 
             # ── O normalization: multiply O by 1/row_sum (TMEM round-trip) ──
             inv_row_sum = Float32(1.0) / row_sum
-            thr_ld = tiled_o_ld.get_slice(sfw_idx)
-            thr_st = tiled_o_st.get_slice(sfw_idx)
-            tOrO_src = thr_ld.partition_S(cute.make_tensor(tStS.iterator + self.tmem_o0_offset, o_rescale_layout))
-            tOrO_dst = thr_st.partition_D(cute.make_tensor(tStS.iterator + self.tmem_o0_offset, o_rescale_layout))
-            rO = cute.make_rmem_tensor(tOrO_src.shape, self.qk_acc_dtype)
-            cute.copy(tiled_o_ld, tOrO_src, rO)
-            for i in cutlass.range(cute.size(rO), vectorize=True):
-                rO[i] = rO[i] * inv_row_sum
-            cute.copy(tiled_o_st, rO, tOrO_dst)
+            tTMrO = cute.make_rmem_tensor((tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype)
+            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 k in cutlass.range(cute.size(tTMrO_i), vectorize=True):
+                    tTMrO_i[k] = tTMrO_i[k] * inv_row_sum
+                cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STOREtO_i)
             cute.arch.fence_view_async_tmem_store()
 
             # ── Epilogue: TMA store O → global ──