diff --git a/tests/fmha_v3_stage_c_example7.py b/tests/fmha_v3_stage_c_example7.py
index afe27334..dac7347d 100644
--- a/tests/fmha_v3_stage_c_example7.py
+++ b/tests/fmha_v3_stage_c_example7.py
@@ -370,56 +370,62 @@ class FmhaV3StageCMulti:
             final_o_bar.arrive_and_wait()
 
             # === O normalization via TMEM load → scale → TMEM store ===
-            # Matches CUTLASS reference's correction_rescale pattern.
-            # Uses Ld32x32bOp / St32x32bOp with the SAME Repetition so the
-            # register tile shapes match (paired atoms).
+            # Matches CUTLASS reference's correction_rescale pattern exactly.
 
             corr_tile_size = 16
 
-            # Sub-tile the O C-fragment
-            tOtO_i = cute.logical_divide(tOtO0, cute.make_layout((128, corr_tile_size)))
             cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1]))
             tOcO = pv_thr.partition_C(cO)
-            tOcO_i = cute.logical_divide(tOcO, cute.make_layout((128, corr_tile_size)))
 
-            # TMEM load + store atoms (paired — same Repetition)
-            tmem_load_o_atom = cute.make_copy_atom(
+            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_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(
+            tmem_store_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[(None, None), 0])
-            tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_o_atom, tOtO_i[(None, None), 0])
-            thr_load_o = tiled_tmem_load_o.get_slice(sfw_idx)
-            thr_store_o = tiled_tmem_store_o.get_slice(sfw_idx)
 
-            tTMEM_LOADtO = thr_load_o.partition_S(tOtO_i[(None, None), None])
-            tTMEM_LOADcO = thr_load_o.partition_D(tOcO_i[(None, None), None])
-            tTMEM_STOREtO = thr_store_o.partition_D(tOtO_i[(None, None), None])
+            tiled_tmem_load_o = tcgen05.make_tmem_copy(tmem_load_atom, tOtO_i)
+            tiled_tmem_store_o = tcgen05.make_tmem_copy(tmem_store_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)
+
+            # 2D register tensor: (frg_shape, n_corr_tiles)
+            tTMrO = cute.make_rmem_tensor(
+                (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype
+            )
 
-            # Scale = 1/row_sum
             inv_row_sum = Float32(1.0) / row_sum
 
-            n_corr = self.pv_mma_tiler[1] // corr_tile_size
-            for i in range(n_corr):
-                tTMEM_LOADtO_i = tTMEM_LOADtO[None, 0, 0, i]
+            for i in range(HEAD_DIM // corr_tile_size):
+                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
                 )
-                tTMrO = cute.make_rmem_tensor(
-                    tTMEM_LOADcO[None, 0, 0, i].shape, self.acc_dtype
-                )
-                cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO)
 
-                # Scale in FP32 registers
-                for j in range(cute.size(tTMrO), vectorize=True):
-                    tTMrO[j] = tTMrO[j] * inv_row_sum
-
-                # Write back to TMEM
-                cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STOREtO_i)
+                cute.copy(tiled_tmem_load_o, tTMEM_LOADtO_i, tTMrO_i)
+                for j in cutlass.range(cute.size(tTMrO_i), vectorize=True):
+                    tTMrO_i[j] = tTMrO_i[j] * inv_row_sum
+                cute.copy(tiled_tmem_store_o, tTMrO_i, tTMEM_STOREtO_i)
 
             cute.arch.fence_view_async_tmem_store()