biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: use make_tiled_tma_atom_A for P TMA

Browse Source
This commit is contained in:
biondizzle
2026-05-24 03:32:15 +00:00
parent 0f943ab48c
commit edacdae017
1 changed files with 8 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
dsv4/kernels/attention/fmha.py
View File

@@ -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).