From 3a7d87adba7ff58eceac4d96c80482c2ac9a7298 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Wed, 27 May 2026 04:54:23 +0000 Subject: [PATCH] Fix test_smem_acc: use keyword args for lse/row_sums --- tests/unit/test_smem_acc.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_smem_acc.py b/tests/unit/test_smem_acc.py index 547d9d0d..97288564 100644 --- a/tests/unit/test_smem_acc.py +++ b/tests/unit/test_smem_acc.py @@ -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