Slice bSG_gC MMA tile coords for TMA copy

2026-05-27 05:11:26 +00:00
parent 754740d5e5
commit b8b6e8cc0b
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
@@ -628,6 +628,8 @@ class FmhaKernel:
                    cute.group_modes(sC, 0, 2),
                    cute.group_modes(tCgC_epi, 0, 2),
                )
+                # Slice off MMA tile coordinates (same as epilogue_tma_store)
+                bSG_gC = bSG_gC[(None, None, None, Int32(0), Int32(0), Int32(0))]
                c_pipe = pipeline.PipelineTmaStore.create(
                    num_stages=self.num_c_stage,
                    producer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(self.epilogue_warp_id))