From b8b6e8cc0b8c3dfcbf1f791c83d05a5d2c64f5dc Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 27 May 2026 05:11:26 +0000 Subject: [PATCH] Slice bSG_gC MMA tile coords for TMA copy --- dsv4/kernels/attention/fmha_smem_acc.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/dsv4/kernels/attention/fmha_smem_acc.py b/dsv4/kernels/attention/fmha_smem_acc.py index 8e88b4bd..e44a3857 100644 --- 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))