diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py
index 4393c83a..f42693d4 100644
--- a/dsv4/kernels/attention/fmha.py
+++ b/dsv4/kernels/attention/fmha.py
@@ -362,35 +362,36 @@ class FmhaKernel:
             row_sum = Float32(0.0)
             scale_log2 = Float32(self.scale_softmax_log2)
 
-            # O rescale atoms (CUTLASS correction_rescale pattern)
-            # Always defined — const_expr guards the kt>0 execution, not the atom setup.
-            # CUTLASS correction_rescale defines these unconditionally.
+            # O rescale atoms (hand-constructed, using composition layout like CUTLASS correction_rescale)
+            # Only needed when there are multiple KV tiles (O must be rescaled per-kt).
+            # With n_kv_tiles=1, no rescale is needed (kt is always 0).
             corr_tile_size = 16
             n_corr_tiles = self.pv_n_tile // corr_tile_size
-            cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1]))
-            tOcO = pv_thr.partition_C(cO)
-            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)
-            tTMrO = cute.make_rmem_tensor(
-                (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype
-            )
+            if const_expr(self.n_kv_tiles > 1):
+                cO = cute.make_identity_tensor((self.pv_mma_tiler[0], self.pv_mma_tiler[1]))
+                tOcO = pv_thr.partition_C(cO)
+                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)
+                tTMrO = cute.make_rmem_tensor(
+                    (tTMEM_LOADcO.shape, 128 // corr_tile_size), self.acc_dtype
+                )
 
             for kt in range(self.n_kv_tiles):
                 si_handle = s_cons.wait_and_advance()
@@ -446,9 +447,6 @@ class FmhaKernel:
                             k2 = k_coord // 64
                             _sP_nostage[(m_coord, k0), 0, (k1, k2)] = rP_bf16[(j0, 0), j1, 0, 0]
                     cute.arch.fence_proxy("async.shared", space="cta")
-                # O rescale: multiply existing O accumulator by acc_scale
-                # when processing a new KV tile (kt > 0).
-                # Same pattern as CUTLASS correction_rescale.
                 if const_expr(self.n_kv_tiles > 1):
                     if kt > 0:
                         for i in range(n_corr_tiles):