From edacdae01736266521e0caabb88cb44ad0521e40 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sun, 24 May 2026 03:32:15 +0000 Subject: [PATCH] fix: use make_tiled_tma_atom_A for P TMA --- dsv4/kernels/attention/fmha.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py index 67bd85e4..9c12fed9 100644 --- a/dsv4/kernels/attention/fmha.py +++ b/dsv4/kernels/attention/fmha.py @@ -119,14 +119,19 @@ class FmhaKernel: p_s = cute.slice_(self.p_smem_s,(None,None,None,0)) gP = torch.zeros(128, self.s_k, dtype=torch.bfloat16, device='cuda') mgP = ct.from_dlpack(gP).mark_layout_dynamic(leading_dim=ct.get_leading_dim(gP)) - tma_p = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), mgP, p_s, self.qk_mma_tiler) + tma_p,mgP = cute.nvgpu.make_tiled_tma_atom_A( + utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, pv_mma.thr_id), + mgP, p_s, self.qk_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape + ) else: # Dummy gP and tma_p (not used, dead-code-eliminated) gP = torch.zeros(128, self.s_k, dtype=torch.bfloat16, device='cuda') mgP = ct.from_dlpack(gP).mark_layout_dynamic(leading_dim=ct.get_leading_dim(gP)) - # Create a dummy TMA using the V SMEM layout (same structure, unused) v_s = cute.slice_(self.v_smem_s,(None,None,None,0)) - tma_p = cpasync.make_tiled_tma_atom(cpasync.CopyBulkTensorTileG2SOp(), mgP, v_s, self.qk_mma_tiler) + tma_p,mgP = cute.nvgpu.make_tiled_tma_atom_A( + utils.sm100.cluster_shape_to_tma_atom_A(self.cluster_shape_mn, pv_mma.thr_id), + mgP, v_s, self.qk_mma_tiler, pv_mma, self.cluster_layout_vmnk.shape + ) # Always create a valid mLSE tensor for the kernel. # CuTeDSL doesn't support None parameters in @cute.kernel. # For normalize=True, mLSE is unused (dead-code-eliminated by compiler).