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

SMEM-P: fix scoping - define tTMEM_LOADcS_frg unconditionally

Browse Source
This commit is contained in:
biondizzle
2026-05-23 19:37:34 +00:00
parent e0aa7ccd19
commit 096a48b5cb
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

@@ -316,9 +316,9 @@ class FmhaKernel:
frg_cnt = 4
frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt
tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile))
# Coordinate fragments for SMEM-P mapping
# Coordinate fragments for SMEM-P mapping (needed unconditionally for scoping)
tTMEM_LOADcS_frg = cute.logical_divide(tTMEM_LOADcS, cute.make_layout(frg_tile))
if self.use_smem_p:
tTMEM_LOADcS_frg = cute.logical_divide(tTMEM_LOADcS, cute.make_layout(frg_tile))
print(f"[SMEM-P CUTLASS] Created tTMEM_LOADcS_frg shape: {cute.shape(tTMEM_LOADcS_frg)}")
for j in range(frg_cnt):