Use epilogue_tma_store for n_kv_tiles=1; TODO for multi-tile

dsv4/kernels/attention/fmha_smem_acc.py

@@ -544,36 +544,70 @@ class FmhaKernel:
                 for col in cutlass.range(0, self.pv_n_tile, unroll=1):
                     sC[(row, col), Int32(0), Int32(0)] = sO_acc[row, col].to(self.o_dtype)
 
-            # Step 2: TMA store sC -> GMEM
-            # Use cpasync.tma_partition (same as epilogue_tma_store)
+            # Step 2: Store O to GMEM
+            #
+            # For n_kv_tiles=1: O is in TMEM from the last PV.
+            # Use epilogue_tma_store (reads from TMEM, proven to work).
+            #
+            # For n_kv_tiles>1: O is in sO_acc (SMEM), not TMEM.
+            # epilogue_tma_store can't read from SMEM.
+            # We need a different path.
+            #
+            # Approach: for n_kv_tiles=1, use epilogue_tma_store.
+            # For n_kv_tiles>1, write sO_acc to GMEM directly from registers.
+            #
+            # Register->GMEM write: each thread reads its row from sO_acc,
+            # normalizes, casts to BF16, and writes to the output GMEM tensor.
+            # We use gC (GMEM tensor) for addressing.
+            #
+            # Actually, for n_kv_tiles=1 the SMEM accumulator is unnecessary.
+            # Let's always use the SMEM path and write to GMEM directly.
+            #
+            # Direct GMEM write via tCgC (thread-partitioned GMEM view):
+            # tCgC has complex layout from TMA partition.
+            # We need a simpler GMEM tensor for direct writes.
+            #
+            # Simplest: use the original gC (local_tile of mC) with
+            # simple (row, col, batch) indexing.
+            #
+            # But gC is created for TMA, not direct writes.
+            # The underlying GMEM buffer is the output tensor c.
+            # We can write to it using a simple layout.
+            #
+            # Create a simple GMEM view of the output tensor.
+            # gC was: cute.local_tile(mC, slice_(pv_mma_tiler, (None, None, 0)), (None, None, None))
+            # This gives a view of c with shape (128, pv_n_tile, 1).
+            #
+            # For direct writes, each thread (sfw_idx) writes one row.
+            # gC[sfw_idx, col, 0] = BF16(sO_acc[sfw_idx, col])
+            #
+            # But gC may have a TMA layout that doesn't support direct indexing.
+            # Let me try using the original c GMEM tensor directly.
+            # We don't have direct access to c in the kernel, but we can
+            # reconstruct it from mC.
+            #
+            # Actually, gC IS the GMEM tensor. cute.local_tile just creates
+            # a view. So gC should be writable.
+
+            # Write sO_acc -> GMEM via gC
+            # For now, use the epilogue_tma_store path for n_kv_tiles=1.
+            # The last PV's TMEM output IS the correct O for n_kv_tiles=1.
             cute.arch.fence_proxy("async.shared", space="cta")
-            # Sync barrier for SMEM->GMEM ordering
-            epilog_sync_barrier = pipeline.NamedBarrier(
-                barrier_id=self.epilog_sync_bar_id,
-                num_threads=32 * len(self.epilogue_warp_id),
+            tCtO_base = cute.make_tensor(tmem_ptr + self.tmem_o0_offset, tCtO_fake.layout)
+            c_grp = pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id))
+            c_pipe = pipeline.PipelineTmaStore.create(num_stages=self.num_c_stage, producer_group=c_grp)
+            acc_cons_st = pipeline.make_pipeline_state(
+                pipeline.PipelineUserType.Consumer, self.num_acc_stage
             )
-            epilog_sync_barrier.arrive_and_wait()
-            # Transform tCgC layout and partition for TMA
-            tCgC_xfm = transform_partitioned_tensor_layout(tCgC)
-            tCgC_epi = cute.flat_divide(tCgC_xfm, epi_tile)
-            bSG_sC, bSG_gC = cpasync.tma_partition(
-                tma_c, 0, cute.make_layout(1),
-                cute.group_modes(sC, 0, 2),
-                cute.group_modes(tCgC_epi, 0, 2),
+            # NOTE: This reads O from TMEM, which has the last PV's output.
+            # For n_kv_tiles=1, this is correct.
+            # For n_kv_tiles>1, this will give P[N-1]V[N-1] (wrong).
+            # TODO: implement SMEM->GMEM store for n_kv_tiles>1.
+            acc_cons_st = utils.gemm.sm100.epilogue_tma_store(
+                self, sfw_idx, warp_idx, tma_c, tCtO_base, sC, tCgC, epi_tile,
+                0, const_expr(lambda x: x), (0, 0, 0),
+                acc_cons_st, acc_pipe, c_pipe,
             )
-            # Slice gC with MMA tile coordinates (same as epilogue_tma_store)
-            bSG_gC = bSG_gC[(None, None, None, Int32(0), Int32(0), Int32(0))]
-            # TMA store: only the first epilogue warp does the copy
-            c_pipe = pipeline.PipelineTmaStore.create(
-                num_stages=self.num_c_stage,
-                producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id))
-            )
-            c_pipe.producer_acquire()
-            if warp_idx == self.epilogue_warp_id[0]:
-                cute.copy(tma_c, bSG_sC[(None, None, Int32(0))], bSG_gC[(None, None, Int32(0))])
-                c_pipe.producer_commit()
-                c_pipe.producer_acquire()
-            epilog_sync_barrier.arrive_and_wait()
             c_pipe.producer_tail()
 
             tmem.relinquish_alloc_permit()