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

fix: CopyAtomUniversalOp

Browse Source
This commit is contained in:
biondizzle
2026-05-23 03:52:47 +00:00
parent 05173c1992
commit 0dbdc4f865
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -200,8 +200,8 @@ class FmhaKernel:
tScS = qk_thr.partition_C(cS)
tTMEM_LOADcS = thr_load.partition_D(tScS)
# P → SMEM copy setup
p_copy_atom = cute.make_copy_atom(cute.nvgpu.copy.AutoCopyOp(), self.q_dtype)
# P → SMEM copy setup (use SMEM copy atom for register→SMEM)
p_copy_atom = cute.make_copy_atom(cute.CopyAtomUniversalOp(), self.q_dtype)
tiled_p_copy = cute.make_tiled_copy(p_copy_atom, tCrP_smem.layout, tCrP_reg.layout, tidx)
tPS_sP = tiled_p_copy.get_slice(sfw_idx).partition_D(tCrP_smem)
tPS_rP = tiled_p_copy.get_slice(sfw_idx).partition_S(tCrP_reg)