diff --git a/tests/unit/test_fmha_v3_stage_c_full.py b/tests/unit/test_fmha_v3_stage_c_full.py
index 588cfebd..209053f6 100644
--- a/tests/unit/test_fmha_v3_stage_c_full.py
+++ b/tests/unit/test_fmha_v3_stage_c_full.py
@@ -284,9 +284,6 @@ class FmhaV3StageC:
                 softmax_done_bar.arrive()
 
             # --- Normalize O in TMEM by row_sum ---
-            # O is at tmem_o0_offset in TMEM. Load each element, divide by row_sum, store back.
-            # Use the O TMEM layout (pv_thr C-fragment) for the load/store.
-            # We need a tiled TMEM copy for O.
             cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1]))
             tOcO = pv_thr.partition_C(cO)
             corr_tile_size = 16
@@ -308,13 +305,11 @@ class FmhaV3StageC:
             for i in range(o_col_tiles):
                 tTMEM_LOAD_O_i = cute.make_tensor(tTMEM_LOAD_OtO.iterator + i * corr_tile_size, tTMEM_LOAD_OtO.layout)
                 tTMEM_STORE_O_i = cute.make_tensor(tTMEM_STORE_OtO.iterator + i * corr_tile_size, tTMEM_STORE_OtO.layout)
-                tTMrO_i_ = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype)
-                tTMrO_i_layout = cute.composition(tTMrO_i_.layout, cute.make_layout(tTMEM_LOAD_OcO.shape[0]))
-                tTMrO_i = cute.make_tensor(tTMrO_i_.iterator, tTMrO_i_layout)
-                cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_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_STORE_O_i)
+                tTMrO = cute.make_rmem_tensor(tTMEM_LOAD_OcO.shape, self.acc_dtype)
+                cute.copy(tiled_tmem_load_o, tTMEM_LOAD_O_i, tTMrO)
+                for k in cutlass.range(cute.size(tTMrO), vectorize=True):
+                    tTMrO[k] = tTMrO[k] * inv_row_sum
+                cute.copy(tiled_tmem_store_o, tTMrO, tTMEM_STORE_O_i)
             cute.arch.fence_view_async_tmem_store()
 
             # --- Epilogue: write O from TMEM to GMEM ---