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

Fix test_smem_acc: use keyword args for lse/row_sums

Browse Source
This commit is contained in:
biondizzle
2026-05-27 04:54:23 +00:00
parent 6a621bdf64
commit 3a7d87adba
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tests/unit/test_smem_acc.py
View File

@@ -51,7 +51,7 @@ def test_smem_acc(hd=64, s_k=256, use_smem_p=False, normalize=True):
mRS = ct.from_dlpack(row_sums_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(row_sums_tensor))
print(f' hd={hd}, s_k={s_k} ({n_kv_tiles} KV tiles, pv_n_tile={pv_n_tile}, n_pv_tiles={n_pv_tiles}): Compiling...', flush=True)
compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, mLSE, mRS)
compiled = cute.compile(kernel, mQ, mK, mV, mC, stream, lse=mLSE, row_sums=mRS)
for nt in range(n_pv_tiles):
v_start = nt * pv_n_tile
@@ -69,7 +69,7 @@ def test_smem_acc(hd=64, s_k=256, use_smem_p=False, normalize=True):
mLSE = ct.from_dlpack(lse_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(lse_tensor))
mRS = ct.from_dlpack(row_sums_tensor).mark_layout_dynamic(leading_dim=ct.get_leading_dim(row_sums_tensor))
compiled(mQ, mK, mV, mC, stream, mLSE, mRS)
compiled(mQ, mK, mV, mC, stream, lse=mLSE, row_sums=mRS)
torch.cuda.synchronize()
c[:, v_start:v_end, :] = c_tile