Slice MMA tile coords from tOgO for TMA copy

2026-05-27 05:39:42 +00:00
parent 1caa737b09
commit 51e456df44
1 changed files with 2 additions and 0 deletions
--- a/dsv4/kernels/attention/fmha_smem_acc.py
+++ b/dsv4/kernels/attention/fmha_smem_acc.py
@@ -640,6 +640,8 @@ class FmhaKernel:
                    sC_flat_staged,
                    tCgC_epi,
                )
+                # Slice off MMA tile coordinates from tOgO
+                tOgO = tOgO[(None, None, None, Int32(0), Int32(0), Int32(0), Int32(0))]
                if warp_idx == self.epilogue_warp_id[0]:
                    cute.copy(tma_c, tOsC, tOgO)
                    cute.arch.cp_async_bulk_commit_group()