From 51e456df44120af20881fc77446e17e74fcae975 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 27 May 2026 05:39:42 +0000 Subject: [PATCH] Slice MMA tile coords from tOgO 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 85d1d215..0767911f 100644 --- 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()