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

test: add try/except for SMEM-P coord test

Browse Source
This commit is contained in:
biondizzle
2026-05-24 02:15:07 +00:00
parent 16bade9e10
commit 8010e3dda2
1 changed files with 15 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
tests/unit/test_smem_p_coord.py
View File

@@ -44,8 +44,21 @@ def test_smem_p_coords():
mLSE = ct.from_dlpack(lse).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse)) mLSE = ct.from_dlpack(lse).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse))
print("Compiling FmhaKernel (hd=256, SMEM-P, normalize=False)...") print("Compiling FmhaKernel (hd=256, SMEM-P, normalize=False)...")
compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE) try:
compiled(mQ, mK, mV, mC, stream, mLSE) compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE)
except Exception as e:
print(f"COMPILE FAILED: {e}")
import traceback
traceback.print_exc()
return
print("Running...")
try:
compiled(mQ, mK, mV, mC, stream, mLSE)
except Exception as e:
print(f"RUN FAILED: {e}")
import traceback
traceback.print_exc()
return
torch.cuda.synchronize() torch.cuda.synchronize()
# The kernel writes P to sP using the coordinate-indexed approach # The kernel writes P to sP using the coordinate-indexed approach