Different coordinate dims for bSG_sC (2D) and bSG_gC (3D)

2026-05-27 05:13:38 +00:00
parent 75b272c5f2
commit 7d1e0a605d
1 changed files with 1 additions and 1 deletions
--- a/dsv4/kernels/attention/fmha_smem_acc.py
+++ b/dsv4/kernels/attention/fmha_smem_acc.py
@@ -636,7 +636,7 @@ class FmhaKernel:
                )
                c_pipe.producer_acquire()
                if warp_idx == self.epilogue_warp_id[0]:
-                    cute.copy(tma_c, bSG_sC[(None, Int32(0))], bSG_gC[(None, Int32(0))])
+                    cute.copy(tma_c, bSG_sC[(None, Int32(0))], bSG_gC[(None, None, Int32(0))])
                    c_pipe.producer_commit()
                c_pipe.producer_tail()